From a8877260019d526c53f9d4de6e14ab8e08ce80da Mon Sep 17 00:00:00 2001
From: Gav <vermagav@users.noreply.github.com>
Date: Thu, 13 Aug 2020 16:53:19 -0700
Subject: [PATCH] Extend TestGateway request creation with header multi-values
 (#737)

In order to use the changes in #733 in services that consume Zanzibar, we extend the TestGateway interface to mirror its existing MakeRequest method with one that creates a test request with all, disparate values for each header's key. We choose extending the interface instead of refactoring MakeRequest to preserve backwards compatibility with all of its many consumers.
---
 CHANGELOG.md                            |  3 +-
 test/lib/bench_gateway/bench_gateway.go | 27 +++++++++++++++++
 test/lib/test_gateway/test_gateway.go   | 39 +++++++++++++++++++++++++
 3 files changed, 68 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 39ea6cf24..17ff73e78 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,10 +4,11 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
-## 0.6.4 - 2020-07-27
+## 0.6.5 - 2020-08-10
 ### Added
 - Added support for fetching multiple header values. https://github.com/uber/zanzibar/pull/733.
 - Set explicit import alias for github.com/uber/zanzibar/runtime in templates. https://github.com/uber/zanzibar/pull/734.
+- Added support for multiple header values in TestGateway utility. https://github.com/uber/zanzibar/pull/737.
 
 ## 0.6.2 - 2020-07-17
 ### Added
diff --git a/test/lib/bench_gateway/bench_gateway.go b/test/lib/bench_gateway/bench_gateway.go
index ad58e6993..78bbfe9a9 100644
--- a/test/lib/bench_gateway/bench_gateway.go
+++ b/test/lib/bench_gateway/bench_gateway.go
@@ -271,14 +271,41 @@ func (gateway *BenchGateway) MakeRequest(
 	fullURL := "http://" + gateway.ActualGateway.RealHTTPAddr + url
 
 	req, err := http.NewRequest(method, fullURL, body)
+	if err != nil {
+		return nil, err
+	}
+
 	for headerName, headerValue := range headers {
 		req.Header.Set(headerName, headerValue)
 	}
 
+	return client.Do(req)
+}
+
+// MakeRequestWithHeaderValues helper
+func (gateway *BenchGateway) MakeRequestWithHeaderValues(
+	method string, url string, headers zanzibar.Header, body io.Reader,
+) (*http.Response, error) {
+	client := gateway.httpClient
+
+	fullURL := "http://" + gateway.ActualGateway.RealHTTPAddr + url
+
+	req, err := http.NewRequest(method, fullURL, body)
 	if err != nil {
 		return nil, err
 	}
 
+	// For each key, fetch every disparate header value and add
+	// it to the bench gateway request.
+	keys := headers.Keys()
+	for _, key := range keys {
+		if values, found := headers.Values(key); found {
+			for _, value := range values {
+				req.Header.Add(key, value)
+			}
+		}
+	}
+
 	return client.Do(req)
 }
 
diff --git a/test/lib/test_gateway/test_gateway.go b/test/lib/test_gateway/test_gateway.go
index 89792d895..6476e56a3 100644
--- a/test/lib/test_gateway/test_gateway.go
+++ b/test/lib/test_gateway/test_gateway.go
@@ -52,6 +52,15 @@ type TestGateway interface {
 		headers map[string]string,
 		body io.Reader,
 	) (*http.Response, error)
+	// MakeRequestWithHeaderValues is an alternate version of `MakeRequest` that uses `zanzibar.Header`
+	// instead of a `map[string]string` to represent headers. This allows us to fetch multiple values
+	// for a given header key.
+	MakeRequestWithHeaderValues(
+		method string,
+		url string,
+		headers zanzibar.Header,
+		body io.Reader,
+	) (*http.Response, error)
 	MakeTChannelRequest(
 		ctx context.Context,
 		thriftService string,
@@ -333,14 +342,44 @@ func (gateway *ChildProcessGateway) MakeRequest(
 	fullURL := "http://" + gateway.RealHTTPAddr + url
 
 	req, err := http.NewRequest(method, fullURL, body)
+	if err != nil {
+		return nil, err
+	}
+
 	for headerName, headerValue := range headers {
 		req.Header.Set(headerName, headerValue)
 	}
 
+	return client.Do(req)
+}
+
+// MakeRequestWithHeaderValues helper
+func (gateway *ChildProcessGateway) MakeRequestWithHeaderValues(
+	method string,
+	url string,
+	headers zanzibar.Header,
+	body io.Reader,
+) (*http.Response, error) {
+	client := gateway.HTTPClient
+
+	fullURL := "http://" + gateway.RealHTTPAddr + url
+
+	req, err := http.NewRequest(method, fullURL, body)
 	if err != nil {
 		return nil, err
 	}
 
+	// For each key, fetch every disparate header value and add
+	// it to the test gateway request.
+	keys := headers.Keys()
+	for _, key := range keys {
+		if values, found := headers.Values(key); found {
+			for _, value := range values {
+				req.Header.Add(key, value)
+			}
+		}
+	}
+
 	return client.Do(req)
 }