From aae9a3ef9818f726190ba44883ee8d140716f45f Mon Sep 17 00:00:00 2001
From: Oleg Bespalov <oleg.bespalov@grafana.com>
Date: Thu, 14 Sep 2023 11:02:09 +0200
Subject: [PATCH] Extracting metadata construction logic

Extracting logic + few more tests
---
 grpc/params.go      | 64 +++++++++++++++++++++++++++++----------------
 grpc/params_test.go | 46 ++++++++++++++++++++++++++++++--
 2 files changed, 85 insertions(+), 25 deletions(-)

diff --git a/grpc/params.go b/grpc/params.go
index 6c08c7d..04983ad 100644
--- a/grpc/params.go
+++ b/grpc/params.go
@@ -41,31 +41,12 @@ func newCallParams(vu modules.VU, input goja.Value) (*callParams, error) {
 	for _, k := range params.Keys() {
 		switch k {
 		case "metadata":
-			v := params.Get(k).Export()
-			rawHeaders, ok := v.(map[string]interface{})
-			if !ok {
-				return result, errors.New("metadata must be an object with key-value pairs")
+			md, err := newMetadata(params.Get(k))
+			if err != nil {
+				return result, fmt.Errorf("invalid metadata param: %w", err)
 			}
 
-			for hk, kv := range rawHeaders {
-				var val string
-
-				// The gRPC spec defines that Binary-valued keys end in -bin
-				// https://grpc.io/docs/what-is-grpc/core-concepts/#metadata
-				if strings.HasSuffix(hk, "-bin") {
-					var binVal []byte
-					if binVal, ok = kv.([]byte); !ok {
-						return result, fmt.Errorf("metadata %q value must be binary", hk)
-					}
-
-					// https://github.com/grpc/grpc-go/blob/v1.57.0/Documentation/grpc-metadata.md#storing-binary-data-in-metadata
-					val = string(binVal)
-				} else if val, ok = kv.(string); !ok {
-					return result, fmt.Errorf("metadata %q value must be a string", hk)
-				}
-
-				result.Metadata.Append(hk, val)
-			}
+			result.Metadata = md
 		case "tags":
 			if err := common.ApplyCustomUserTags(rt, &result.TagsAndMeta, params.Get(k)); err != nil {
 				return result, fmt.Errorf("metric tags: %w", err)
@@ -85,6 +66,43 @@ func newCallParams(vu modules.VU, input goja.Value) (*callParams, error) {
 	return result, nil
 }
 
+// newMetadata constructs a metadata.MD from the input value.
+func newMetadata(input goja.Value) (metadata.MD, error) {
+	md := metadata.New(nil)
+
+	if common.IsNullish(input) {
+		return md, nil
+	}
+
+	v := input.Export()
+
+	rawHeaders, ok := v.(map[string]interface{})
+	if !ok {
+		return md, errors.New("must be an object with key-value pairs")
+	}
+
+	for hk, kv := range rawHeaders {
+		var val string
+		// The gRPC spec defines that Binary-valued keys end in -bin
+		// https://grpc.io/docs/what-is-grpc/core-concepts/#metadata
+		if strings.HasSuffix(hk, "-bin") {
+			var binVal []byte
+			if binVal, ok = kv.([]byte); !ok {
+				return md, fmt.Errorf("%q value must be binary", hk)
+			}
+
+			// https://github.com/grpc/grpc-go/blob/v1.57.0/Documentation/grpc-metadata.md#storing-binary-data-in-metadata
+			val = string(binVal)
+		} else if val, ok = kv.(string); !ok {
+			return md, fmt.Errorf("%q value must be a string", hk)
+		}
+
+		md.Append(hk, val)
+	}
+
+	return md, nil
+}
+
 // SetSystemTags sets the system tags for the call.
 func (p *callParams) SetSystemTags(state *lib.State, addr string, methodName string) {
 	if state.Options.SystemTags.Has(metrics.TagURL) {
diff --git a/grpc/params_test.go b/grpc/params_test.go
index 641cb4f..c43c6d9 100644
--- a/grpc/params_test.go
+++ b/grpc/params_test.go
@@ -13,10 +13,11 @@ import (
 	"go.k6.io/k6/js/modulestest"
 	"go.k6.io/k6/lib"
 	"go.k6.io/k6/metrics"
+	"google.golang.org/grpc/metadata"
 	"gopkg.in/guregu/null.v3"
 )
 
-func TestParamsInvalidInput(t *testing.T) {
+func TestCallParamsInvalidInput(t *testing.T) {
 	t.Parallel()
 
 	testCases := []struct {
@@ -39,6 +40,11 @@ func TestParamsInvalidInput(t *testing.T) {
 			JSON:        `{ timeout: "please" }`,
 			ErrContains: `invalid duration`,
 		},
+		{
+			Name:        "InvalidMetadata",
+			JSON:        `{ metadata: "lorem" }`,
+			ErrContains: `invalid metadata param: must be an object with key-value pairs`,
+		},
 	}
 
 	for _, tc := range testCases {
@@ -56,7 +62,43 @@ func TestParamsInvalidInput(t *testing.T) {
 	}
 }
 
-func TestParamsTimeOutParse(t *testing.T) {
+func TestCallParamsMetadata(t *testing.T) {
+	t.Parallel()
+
+	testCases := []struct {
+		Name             string
+		JSON             string
+		ExpectedMetadata metadata.MD
+	}{
+		{
+			Name:             "EmptyMetadata",
+			JSON:             `{}`,
+			ExpectedMetadata: metadata.New(nil),
+		},
+		{
+			Name:             "Metadata",
+			JSON:             `{metadata: {foo: "bar", baz: "qux"}}`,
+			ExpectedMetadata: metadata.New(map[string]string{"foo": "bar", "baz": "qux"}),
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+
+		t.Run(tc.Name, func(t *testing.T) {
+			t.Parallel()
+
+			testRuntime, params := newParamsTestRuntime(t, tc.JSON)
+
+			p, err := newCallParams(testRuntime.VU, params)
+
+			require.NoError(t, err)
+			assert.Equal(t, tc.ExpectedMetadata, p.Metadata)
+		})
+	}
+}
+
+func TestCallParamsTimeOutParse(t *testing.T) {
 	t.Parallel()
 
 	testCases := []struct {