fix: use custom template funcs for metadata

bojand · May 4, 2021 · cfcfaac · cfcfaac
1 parent 79e4749
commit cfcfaac
Show file tree

Hide file tree

Showing 4 changed files with 160 additions and 7 deletions.
diff --git a/runner/data.go b/runner/data.go
@@ -218,9 +218,9 @@ func (dp *dataProvider) getMessages(ctd *CallData, i int, inputData []byte) ([]*
 	return inputs, nil
 }
 
-func newMetadataProvider(mtd *desc.MethodDescriptor, mdData []byte) (*mdProvider, error) {
+func newMetadataProvider(mtd *desc.MethodDescriptor, mdData []byte, funcs template.FuncMap) (*mdProvider, error) {
 	// Test if we can preseed data
-	ctd := newCallData(mtd, nil, "", 0)
+	ctd := newCallData(mtd, funcs, "", 0)
 	ha, err := ctd.hasAction(string(mdData))
 	if err != nil {
 		return nil, err

diff --git a/runner/data_test.go b/runner/data_test.go
@@ -3,6 +3,7 @@ package runner
 import (
 	"encoding/json"
 	"testing"
+	"text/template"
 
 	"github.com/golang/protobuf/proto"
 
@@ -274,7 +275,7 @@ func TestMetadata_newMetadataProvider(t *testing.T) {
 			nil)
 		assert.NoError(t, err)
 
-		mdp, err := newMetadataProvider(mtdUnary, []byte(`{"token":"asdf"}`))
+		mdp, err := newMetadataProvider(mtdUnary, []byte(`{"token":"asdf"}`), nil)
 		assert.NoError(t, err)
 		assert.NotNil(t, mdp)
 		assert.NotNil(t, mdp.preseed)
@@ -288,7 +289,7 @@ func TestMetadata_newMetadataProvider(t *testing.T) {
 			nil)
 		assert.NoError(t, err)
 
-		mdp, err := newMetadataProvider(mtdUnary, []byte(`{"token":"{{ .RequestNumber }}"}`))
+		mdp, err := newMetadataProvider(mtdUnary, []byte(`{"token":"{{ .RequestNumber }}"}`), nil)
 		assert.NoError(t, err)
 		assert.NotNil(t, mdp)
 		assert.Nil(t, mdp.preseed)
@@ -303,7 +304,7 @@ func TestMetadata_getMetadataForCall(t *testing.T) {
 			nil)
 		assert.NoError(t, err)
 
-		mdp, err := newMetadataProvider(mtdUnary, []byte(`{"token":"asdf"}`))
+		mdp, err := newMetadataProvider(mtdUnary, []byte(`{"token":"asdf"}`), nil)
 		assert.NoError(t, err)
 		assert.NotNil(t, mdp.preseed)
 
@@ -323,7 +324,7 @@ func TestMetadata_getMetadataForCall(t *testing.T) {
 			nil)
 		assert.NoError(t, err)
 
-		mdp, err := newMetadataProvider(mtdUnary, []byte(`{"token":"{{ .RequestNumber }}"}`))
+		mdp, err := newMetadataProvider(mtdUnary, []byte(`{"token":"{{ .RequestNumber }}"}`), nil)
 		assert.NoError(t, err)
 		assert.Nil(t, mdp.preseed)
 
@@ -344,4 +345,30 @@ func TestMetadata_getMetadataForCall(t *testing.T) {
 		assert.NotSame(t, md1, md2)
 		assert.NotEqual(t, md1, md2)
 	})
+
+	t.Run("with cunstom function", func(t *testing.T) {
+		mtdUnary, err := protodesc.GetMethodDescFromProto(
+			"helloworld.Greeter.SayHello",
+			"../testdata/greeter.proto",
+			nil)
+		assert.NoError(t, err)
+
+		funcs := template.FuncMap{
+			"customFunc": func() string {
+				return "custom-value"
+			},
+		}
+
+		mdp, err := newMetadataProvider(mtdUnary, []byte(`{"token":"{{ customFunc }}"}`), funcs)
+		assert.NoError(t, err)
+		assert.Nil(t, mdp.preseed)
+
+		cd := newCallData(mtdUnary, funcs, "123", 1)
+
+		md1, err := mdp.getMetadataForCall(cd)
+		assert.NoError(t, err)
+		assert.NotNil(t, md1)
+		assert.Equal(t, []string{"custom-value"}, md1.Get("token"))
+		assert.NotSame(t, mdp.preseed, md1)
+	})
 }
diff --git a/runner/requester.go b/runner/requester.go
@@ -135,7 +135,7 @@ func NewRequester(c *RunConfig) (*Requester, error) {
 	if c.mdProviderFunc != nil {
 		reqr.metadataProvider = c.mdProviderFunc
 	} else {
-		defaultMDProvider, err := newMetadataProvider(reqr.mtd, c.metadata)
+		defaultMDProvider, err := newMetadataProvider(reqr.mtd, c.metadata, c.funcs)
 		if err != nil {
 			return nil, err
 		}

diff --git a/runner/run_test.go b/runner/run_test.go
@@ -740,6 +740,63 @@ func TestRunUnary(t *testing.T) {
 		assert.Equal(t, []string{"0", "1", "2", "3", "4"}, names)
 	})
 
+	t.Run("with data custom function", func(t *testing.T) {
+		gs.ResetCounters()
+
+		data := make(map[string]interface{})
+		data["name"] = "{{customFunc}}"
+
+		report, err := Run(
+			"helloworld.Greeter.SayHello",
+			internal.TestLocalhost,
+			WithProtoFile("../testdata/greeter.proto", []string{}),
+			WithTotalRequests(3),
+			WithConcurrency(1),
+			WithTimeout(time.Duration(20*time.Second)),
+			WithDialTimeout(time.Duration(20*time.Second)),
+			WithData(data),
+			WithTemplateFuncs(template.FuncMap{
+				"customFunc": func() string {
+					return "custom-value"
+				},
+			}),
+			WithInsecure(true),
+		)
+
+		assert.NoError(t, err)
+
+		assert.NotNil(t, report)
+
+		assert.Equal(t, 3, int(report.Count))
+		assert.NotZero(t, report.Average)
+		assert.NotZero(t, report.Fastest)
+		assert.NotZero(t, report.Slowest)
+		assert.NotZero(t, report.Rps)
+		assert.Empty(t, report.Name)
+		assert.NotEmpty(t, report.Date)
+		assert.NotEmpty(t, report.Options)
+		assert.NotEmpty(t, report.Details)
+		assert.Equal(t, true, report.Options.Insecure)
+		assert.NotEmpty(t, report.LatencyDistribution)
+		assert.Equal(t, ReasonNormalEnd, report.EndReason)
+		assert.Empty(t, report.ErrorDist)
+
+		assert.NotEqual(t, report.Average, report.Slowest)
+		assert.NotEqual(t, report.Average, report.Fastest)
+		assert.NotEqual(t, report.Slowest, report.Fastest)
+
+		count := gs.GetCount(callType)
+		assert.Equal(t, 3, count)
+
+		calls := gs.GetCalls(callType)
+		assert.NotNil(t, calls)
+		assert.Len(t, calls, 3)
+		fmt.Printf("%+v\n", calls)
+		assert.Equal(t, "custom-value", calls[0][0].GetName())
+		assert.Equal(t, "custom-value", calls[1][0].GetName())
+		assert.Equal(t, "custom-value", calls[2][0].GetName())
+	})
+
 	t.Run("with metadata provider", func(t *testing.T) {
 		gs.ResetCounters()
 
@@ -812,6 +869,75 @@ func TestRunUnary(t *testing.T) {
 
 		assert.Equal(t, []string{"0", "__record_metadata__||token:secret1", "2", "__record_metadata__||token:secret3", "4"}, names)
 	})
+
+	t.Run("with metadata custom function", func(t *testing.T) {
+		gs.ResetCounters()
+
+		data := make(map[string]interface{})
+		data["name"] = "__record_metadata__"
+
+		metadata := make(map[string]string)
+		metadata["token"] = "{{customFunc}}"
+
+		report, err := Run(
+			"helloworld.Greeter.SayHello",
+			internal.TestLocalhost,
+			WithProtoFile("../testdata/greeter.proto", []string{}),
+			WithTotalRequests(3),
+			WithConcurrency(1),
+			WithTimeout(time.Duration(20*time.Second)),
+			WithDialTimeout(time.Duration(20*time.Second)),
+			WithData(data),
+			WithMetadata(metadata),
+			WithTemplateFuncs(template.FuncMap{
+				"customFunc": func() string {
+					return "custom-value"
+				},
+			}),
+			WithInsecure(true),
+		)
+
+		assert.NoError(t, err)
+
+		assert.NotNil(t, report)
+
+		assert.Equal(t, 3, int(report.Count))
+		assert.NotZero(t, report.Average)
+		assert.NotZero(t, report.Fastest)
+		assert.NotZero(t, report.Slowest)
+		assert.NotZero(t, report.Rps)
+		assert.Empty(t, report.Name)
+		assert.NotEmpty(t, report.Date)
+		assert.NotEmpty(t, report.Options)
+		assert.NotEmpty(t, report.Details)
+		assert.Equal(t, true, report.Options.Insecure)
+		assert.NotEmpty(t, report.LatencyDistribution)
+		assert.Equal(t, ReasonNormalEnd, report.EndReason)
+		assert.Empty(t, report.ErrorDist)
+
+		assert.NotEqual(t, report.Average, report.Slowest)
+		assert.NotEqual(t, report.Average, report.Fastest)
+		assert.NotEqual(t, report.Slowest, report.Fastest)
+
+		count := gs.GetCount(callType)
+		assert.Equal(t, 3, count)
+
+		calls := gs.GetCalls(callType)
+		assert.NotNil(t, calls)
+		assert.Len(t, calls, 3)
+		names := make([]string, 0)
+		for _, msgs := range calls {
+			for _, msg := range msgs {
+				names = append(names, msg.GetName())
+			}
+		}
+
+		assert.Equal(t, []string{
+			"__record_metadata__||token:custom-value",
+			"__record_metadata__||token:custom-value",
+			"__record_metadata__||token:custom-value",
+		}, names)
+	})
 }
 
 func TestRunServerStreaming(t *testing.T) {