From 74a905b3ac811444a2fed0575c8327330ec475b3 Mon Sep 17 00:00:00 2001
From: Fedor Shiriaev <fedor.shiriaev@careem.com>
Date: Thu, 19 Nov 2020 16:03:24 +0100
Subject: [PATCH] feat: add ability to pass the experience id within the
 context (#248)

Add ability to pass the experience id with the context on SDK get/post request. Useful when one SDK client used per different customers.
---
 client.go      | 40 ++++++++++++++++++++++++++++++++++++----
 client_test.go | 38 ++++++++++++++++++++++++++++++++------
 2 files changed, 68 insertions(+), 10 deletions(-)

diff --git a/client.go b/client.go
index b19091e..9164f60 100644
--- a/client.go
+++ b/client.go
@@ -50,8 +50,14 @@ type ClientOption func(*Client) error
 
 var defaultRequestsPerSecond = 50
 
+type contextKey string
+func (c contextKey) String() string {
+	return "maps " + string(c)
+}
+
 const (
 	ExperienceIdHeaderName = "X-GOOG-MAPS-EXPERIENCE-ID"
+	contextExperienceId    = contextKey("EXP-IDS")
 )
 
 // NewClient constructs a new Client which can make requests to the Google Maps
@@ -118,6 +124,23 @@ func WithAPIKeyAndSignature(apiKey, signature string) ClientOption {
 	}
 }
 
+// ExperienceIdContext injects the experienceIds in the context, from where they will be pull out
+// in the post/get handlers. Useful if a customer uses one client instance per different experiences calls
+func ExperienceIdContext(ctx context.Context, experienceIds ...string) context.Context {
+	if ctx != nil {
+		return context.WithValue(ctx, contextExperienceId, experienceIds)
+	}
+	return ctx
+}
+
+// ExperienceIdFromContext returns experienceIds from context if presented
+func ExperienceIdFromContext(ctx context.Context) []string {
+	if experiencesId := ctx.Value(contextExperienceId); experiencesId != nil {
+		return experiencesId.([]string)
+	}
+	return nil
+}
+
 // WithBaseURL configures a Maps API client with a custom base url
 func WithBaseURL(baseURL string) ClientOption {
 	return func(c *Client) error {
@@ -205,7 +228,7 @@ func (c *Client) get(ctx context.Context, config *apiConfig, apiReq apiRequest)
 		return nil, err
 	}
 
-	c.setExperienceIdHeader(req)
+	c.setExperienceIdHeader(ctx, req)
 
 	q, err := c.generateAuthQuery(config.path, apiReq.params(), config.acceptsClientID, config.acceptsSignature)
 	if err != nil {
@@ -235,7 +258,7 @@ func (c *Client) post(ctx context.Context, config *apiConfig, apiReq interface{}
 	}
 	req.Header.Set("Content-Type", "application/json")
 
-	c.setExperienceIdHeader(req)
+	c.setExperienceIdHeader(ctx, req)
 
 	q, err := c.generateAuthQuery(config.path, url.Values{}, config.acceptsClientID, config.acceptsSignature)
 	if err != nil {
@@ -294,9 +317,18 @@ func (c *Client) clearExperienceId() {
 	c.experienceId = nil
 }
 
-func (c *Client) setExperienceIdHeader(req *http.Request) {
+func (c *Client) setExperienceIdHeader(ctx context.Context, req *http.Request) {
+	var ids []string
 	if len(c.experienceId) > 0 {
-		req.Header.Set(ExperienceIdHeaderName, strings.Join(c.experienceId, ","))
+		ids = append(ids, c.experienceId...)
+	}
+	if experiencesId := ExperienceIdFromContext(ctx); experiencesId != nil {
+		for _, v := range experiencesId {
+			ids = append(ids, v)
+		}
+	}
+	if len(ids) != 0 {
+		req.Header.Set(ExperienceIdHeaderName, strings.Join(ids, ","))
 	}
 }
 
diff --git a/client_test.go b/client_test.go
index 30c9a8d..98f99a0 100644
--- a/client_test.go
+++ b/client_test.go
@@ -15,6 +15,7 @@
 package maps
 
 import (
+	"context"
 	"net/http"
 	"strings"
 	"testing"
@@ -69,20 +70,45 @@ func TestClientSetExperienceIdHeader(t *testing.T) {
 	// slice has two elements
 	c.experienceId = ids
 	req, _ := http.NewRequest("GET", "/", nil)
-	c.setExperienceIdHeader(req)
-	assert.Equal(t, req.Header.Get(ExperienceIdHeaderName), strings.Join(ids, ","))
+	c.setExperienceIdHeader(context.Background(), req)
+	assert.Equal(t, strings.Join(ids, ","), req.Header.Get(ExperienceIdHeaderName))
 
 	// slice is nil
 	c.experienceId = nil
 	req, _ = http.NewRequest("GET", "/", nil)
-	c.setExperienceIdHeader(req)
-	assert.Equal(t, req.Header.Get(ExperienceIdHeaderName), "")
+	c.setExperienceIdHeader(context.Background(), req)
+	assert.Equal(t, "", req.Header.Get(ExperienceIdHeaderName))
 
 	// slice is empty
 	c.experienceId = []string{}
 	req, _ = http.NewRequest("GET", "/", nil)
-	c.setExperienceIdHeader(req)
-	assert.Equal(t, req.Header.Get(ExperienceIdHeaderName), "")
+	c.setExperienceIdHeader(context.Background(), req)
+	assert.Equal(t, "", req.Header.Get(ExperienceIdHeaderName))
+
+	var ctx context.Context
+	// context has one element
+	c.experienceId = []string{}
+	ctx = context.Background()
+	ctx = ExperienceIdContext(ctx, "foo")
+	req, _ = http.NewRequest("GET", "/", nil)
+	c.setExperienceIdHeader(ctx, req)
+	assert.Equal(t, "foo", req.Header.Get(ExperienceIdHeaderName))
+
+	// context has two elements
+	c.experienceId = []string{}
+	ctx = context.Background()
+	ctx = ExperienceIdContext(ctx, ids...)
+	req, _ = http.NewRequest("GET", "/", nil)
+	c.setExperienceIdHeader(ctx, req)
+	assert.Equal(t, strings.Join(ids, ","), req.Header.Get(ExperienceIdHeaderName))
+
+	// context has two elements and client has two elements
+	c.experienceId = ids
+	ctx = context.Background()
+	ctx = ExperienceIdContext(ctx, ids...)
+	req, _ = http.NewRequest("GET", "/", nil)
+	c.setExperienceIdHeader(ctx, req)
+	assert.Equal(t, strings.Join(ids, ",") + "," + strings.Join(ids, ","), req.Header.Get(ExperienceIdHeaderName))
 }
 
 func TestClientExperienceIdSample(t *testing.T) {