Repository: bojand/ghz
Pull request #218 (2020-08-24 00:56:49)

Add template functions
Adds template functions:

`newUUID` - Generates a new UUID for each invocation.

`randomString` - Generates a new random string for each incovation. Accepts a length parameter. If the argument is `<= 0` then a ranom string is generated with a random length between length of `2` and `16`.

Example usage with data:

```sh
-d '{"order_id":"{{newUUID}}", "item_id":"{{newUUID}}", "sku":"{{randomString 8 }}", "product_name":"{{ranomdString 0}}"}'
```

Would result in data with JSON representation:

```json
{
  "order_id": "3974e7b3-5946-4df5-bed3-8c3dc9a0be19",
  "item_id": "cd9c2604-cd9b-43a8-9cbb-d1ad26ca93a4",
  "sku": "HlFTAxcm",
  "product_name": "xg3NEC"
}
```

Also adds `WithTemplateFuncs()` API option.

Addresses #213.


Patch:
diff --git a/runner/call_template_data.go b/runner/call_template_data.go
index d3fcd08..be95f6f 100644
--- a/runner/call_template_data.go
+++ b/runner/call_template_data.go
@@ -3,6 +3,7 @@ package runner
 import (
 	"bytes"
 	"encoding/json"
+	"math/rand"
 	"text/template"
 	"time"
 
@@ -10,6 +11,12 @@ import (
 	"github.com/jhump/protoreflect/desc"
 )
 
+const charset = "abcdefghijklmnopqrstuvwxyz" +
+	"ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
+
+var seededRand *rand.Rand = rand.New(
+	rand.NewSource(time.Now().UnixNano()))
+
 // call template data
 type callTemplateData struct {
 	WorkerID           string // unique worker ID
@@ -26,13 +33,34 @@ type callTemplateData struct {
 	TimestampUnixMilli int64  // timestamp of the call as unix time in milliseconds
 	TimestampUnixNano  int64  // timestamp of the call as unix time in nanoseconds
 	UUID               string // generated UUIDv4 for each call
+
+	templateFuncs template.FuncMap
+}
+
+var tmplFuncMap = template.FuncMap{
+	"newUUID":      newUUID,
+	"randomString": randomString,
 }
 
 // newCallTemplateData returns new call template data
-func newCallTemplateData(mtd *desc.MethodDescriptor, workerID string, reqNum int64) *callTemplateData {
+func newCallTemplateData(
+	mtd *desc.MethodDescriptor,
+	funcs template.FuncMap,
+	workerID string, reqNum int64) *callTemplateData {
 	now := time.Now()
 	newUUID, _ := uuid.NewRandom()
 
+	fns := make(template.FuncMap, len(funcs)+2)
+	for k, v := range tmplFuncMap {
+		fns[k] = v
+	}
+
+	if len(funcs) > 0 {
+		for k, v := range funcs {
+			fns[k] = v
+		}
+	}
+
 	return &callTemplateData{
 		WorkerID:           workerID,
 		RequestNumber:      reqNum,
@@ -48,11 +76,12 @@ func newCallTemplateData(mtd *desc.MethodDescriptor, workerID string, reqNum int
 		TimestampUnixMilli: now.UnixNano() / 1000000,
 		TimestampUnixNano:  now.UnixNano(),
 		UUID:               newUUID.String(),
+		templateFuncs:      fns,
 	}
 }
 
 func (td *callTemplateData) execute(data string) (*bytes.Buffer, error) {
-	t := template.Must(template.New("call_template_data").Parse(data))
+	t := template.Must(template.New("call_template_data").Funcs(td.templateFuncs).Parse(data))
 	var tpl bytes.Buffer
 	err := t.Execute(&tpl, td)
 	return &tpl, err
@@ -72,7 +101,7 @@ func (td *callTemplateData) executeData(data string) ([]byte, error) {
 	return []byte{}, nil
 }
 
-func (td *callTemplateData) executeMetadata(metadata string) (*map[string]string, error) {
+func (td *callTemplateData) executeMetadata(metadata string) (map[string]string, error) {
 	var mdMap map[string]string
 
 	if len(metadata) > 0 {
@@ -88,5 +117,29 @@ func (td *callTemplateData) executeMetadata(metadata string) (*map[string]string
 		}
 	}
 
-	return &mdMap, nil
+	return mdMap, nil
+}
+
+func newUUID() string {
+	newUUID, _ := uuid.NewRandom()
+	return newUUID.String()
+}
+
+const maxLen = 16
+const minLen = 2
+
+func stringWithCharset(length int, charset string) string {
+	b := make([]byte, length)
+	for i := range b {
+		b[i] = charset[seededRand.Intn(len(charset))]
+	}
+	return string(b)
+}
+
+func randomString(length int) string {
+	if length <= 0 {
+		length = seededRand.Intn(maxLen-minLen+1) + minLen
+	}
+
+	return stringWithCharset(length, charset)
 }
diff --git a/runner/options.go b/runner/options.go
index f42fd57..91c294b 100644
--- a/runner/options.go
+++ b/runner/options.go
@@ -11,6 +11,7 @@ import (
 	"path/filepath"
 	"runtime"
 	"strings"
+	"text/template"
 	"time"
 
 	"github.com/pkg/errors"
@@ -70,6 +71,8 @@ type RunConfig struct {
 	cpus      int
 	tags      []byte
 	skipFirst int
+
+	funcs template.FuncMap
 }
 
 // Option controls some aspect of run
@@ -546,6 +549,15 @@ func WithLogger(log Logger) Option {
 	}
 }
 
+// WithTemplateFuncs adds additional tempalte functions
+func WithTemplateFuncs(funcMap template.FuncMap) Option {
+	return func(o *RunConfig) error {
+		o.funcs = funcMap
+
+		return nil
+	}
+}
+
 func newConfig(call, host string, options ...Option) (*RunConfig, error) {
 	call = strings.TrimSpace(call)
 	host = strings.TrimSpace(host)
diff --git a/runner/worker.go b/runner/worker.go
index d371ef3..f5a8fb6 100644
--- a/runner/worker.go
+++ b/runner/worker.go
@@ -65,7 +65,7 @@ func (w *Worker) makeRequest() error {
 
 	reqNum := atomic.AddInt64(w.reqCounter, 1)
 
-	ctd := newCallTemplateData(w.mtd, w.workerID, reqNum)
+	ctd := newCallTemplateData(w.mtd, w.config.funcs, w.workerID, reqNum)
 
 	var inputs []*dynamic.Message
 	var err error
@@ -88,8 +88,8 @@ func (w *Worker) makeRequest() error {
 	}
 
 	var reqMD *metadata.MD
-	if mdMap != nil && len(*mdMap) > 0 {
-		md := metadata.New(*mdMap)
+	if len(mdMap) > 0 {
+		md := metadata.New(mdMap)
 		reqMD = &md
 	} else {
 		reqMD = &metadata.MD{}
diff --git a/web/ui/package-lock.json b/web/ui/package-lock.json
index ffaab7c..724a3f6 100644
--- a/web/ui/package-lock.json
+++ b/web/ui/package-lock.json
@@ -3953,9 +3953,9 @@
       "dev": true
     },
     "elliptic": {
-      "version": "6.5.3",
-      "resolved": "https://registry.npmjs.org/elliptic/-/elliptic-6.5.3.tgz",
-      "integrity": "sha512-IMqzv5wNQf+E6aHeIqATs0tOLeOTwj1QKbRcS3jBbYkl5oLAserA8yJTT7/VyHUYG91PRmPyeQDObKLPpeS4dw==",
+      "version": "6.5.1",
+      "resolved": "https://registry.npmjs.org/elliptic/-/elliptic-6.5.1.tgz",
+      "integrity": "sha512-xvJINNLbTeWQjrl6X+7eQCrIy/YPv5XCpKW6kB5mKvtnGILoLDcySuwomfdzt0BMdLNVnuRNTuzKNHj0bva1Cg==",
       "dev": true,
       "requires": {
         "bn.js": "^4.4.0",
diff --git a/www/docs/calldata.md b/www/docs/calldata.md
index c9b98c9..6ee7cc2 100644
--- a/www/docs/calldata.md
+++ b/www/docs/calldata.md
@@ -53,6 +53,16 @@ type callTemplateData struct {
 }
 ```
 
+**Functions**
+
+There are also two template functions available:
+
+`newUUID` - Generates a new UUID for each invocation.
+
+`randomString` - Generates a new random string for each incovation. Accepts a length parameter. If the argument is `<= 0` then a ranom string is generated with a random length between length of `2` and `16`.
+
+**Examples**
+
 This can be useful to inject variable information into the message data JSON or metadata JSON payloads for each request, such as timestamp or unique request number. For example:
 
 ```sh
@@ -69,4 +79,19 @@ Would result in server getting the following metadata map represented here in JS
 }
 ```
 
+```sh
+-d '{"order_id":"{{newUUID}}", "item_id":"{{newUUID}}", "sku":"{{randomString 8 }}", "product_name":"{{ranomdString 0}}"}'
+```
+
+Would result in data with JSON representation:
+
+```json
+{
+  "order_id": "3974e7b3-5946-4df5-bed3-8c3dc9a0be19",
+  "item_id": "cd9c2604-cd9b-43a8-9cbb-d1ad26ca93a4",
+  "sku": "HlFTAxcm",
+  "product_name": "xg3NEC"
+}
+```
+
 See [example calls](examples.md) for some more usage examples.
diff --git a/www/website/package-lock.json b/www/website/package-lock.json
index 5f86bac..355d039 100644
--- a/www/website/package-lock.json
+++ b/www/website/package-lock.json
@@ -7039,9 +7039,9 @@
       "dev": true
     },
     "prismjs": {
-      "version": "1.21.0",
-      "resolved": "https://registry.npmjs.org/prismjs/-/prismjs-1.21.0.tgz",
-      "integrity": "sha512-uGdSIu1nk3kej2iZsLyDoJ7e9bnPzIgY0naW/HdknGj61zScaprVEVGHrPoXqI+M9sP0NDnTK2jpkvmldpuqDw==",
+      "version": "1.20.0",
+      "resolved": "https://registry.npmjs.org/prismjs/-/prismjs-1.20.0.tgz",
+      "integrity": "sha512-AEDjSrVNkynnw6A+B1DsFkd6AVdTnp+/WoUixFRULlCLZVRZlVQMVWio/16jv7G1FscUxQxOQhWwApgbnxr6kQ==",
       "dev": true,
       "requires": {
         "clipboard": "^2.0.0"


Test patch:
diff --git a/runner/call_template_data_test.go b/runner/call_template_data_test.go
index 9523210..2807468 100644
--- a/runner/call_template_data_test.go
+++ b/runner/call_template_data_test.go
@@ -1,9 +1,12 @@
 package runner
 
 import (
+	"strings"
 	"testing"
+	"text/template"
 
 	"github.com/bojand/ghz/protodesc"
+	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -12,7 +15,7 @@ func TestCallTemplateData_New(t *testing.T) {
 	assert.NoError(t, err)
 	assert.NotNil(t, md)
 
-	ctd := newCallTemplateData(md, "worker_id_123", 100)
+	ctd := newCallTemplateData(md, nil, "worker_id_123", 100)
 
 	assert.NotNil(t, ctd)
 	assert.Equal(t, "worker_id_123", ctd.WorkerID)
@@ -38,7 +41,7 @@ func TestCallTemplateData_ExecuteData(t *testing.T) {
 	assert.NoError(t, err)
 	assert.NotNil(t, md)
 
-	ctd := newCallTemplateData(md, "worker_id_123", 200)
+	ctd := newCallTemplateData(md, nil, "worker_id_123", 200)
 
 	assert.NotNil(t, ctd)
 
@@ -89,7 +92,7 @@ func TestCallTemplateData_ExecuteMetadata(t *testing.T) {
 	assert.NoError(t, err)
 	assert.NotNil(t, md)
 
-	ctd := newCallTemplateData(md, "worker_id_123", 200)
+	ctd := newCallTemplateData(md, nil, "worker_id_123", 200)
 
 	assert.NotNil(t, ctd)
 
@@ -101,17 +104,17 @@ func TestCallTemplateData_ExecuteMetadata(t *testing.T) {
 	}{
 		{"no template",
 			`{"trace_id":"asdf"}`,
-			&map[string]string{"trace_id": "asdf"},
+			map[string]string{"trace_id": "asdf"},
 			false,
 		},
 		{"with template",
 			`{"trace_id":"{{.RequestNumber}} asdf {{.FullyQualifiedName}} {{.MethodName}} {{.ServiceName}} {{.InputName}} {{.OutputName}} {{.IsClientStreaming}} {{.IsServerStreaming}}"}`,
-			&map[string]string{"trace_id": "200 asdf helloworld.Greeter.SayHello SayHello Greeter HelloRequest HelloReply false false"},
+			map[string]string{"trace_id": "200 asdf helloworld.Greeter.SayHello SayHello Greeter HelloRequest HelloReply false false"},
 			false,
 		},
 		{"with unknown action",
 			`{"trace_id":"asdf {{.Something}} {{.MethodName}} bob"}`,
-			&map[string]string{"trace_id": "asdf {{.Something}} {{.MethodName}} bob"},
+			map[string]string{"trace_id": "asdf {{.Something}} {{.MethodName}} bob"},
 			false,
 		},
 	}
@@ -129,3 +132,140 @@ func TestCallTemplateData_ExecuteMetadata(t *testing.T) {
 		})
 	}
 }
+
+func TestCallTemplateData_ExecuteFuncs(t *testing.T) {
+	md, err := protodesc.GetMethodDescFromProto("helloworld.Greeter/SayHello", "../testdata/greeter.proto", []string{})
+	assert.NoError(t, err)
+	assert.NotNil(t, md)
+
+	ctd := newCallTemplateData(md, nil, "worker_id_123", 200)
+
+	assert.NotNil(t, ctd)
+
+	t.Run("newUUID", func(t *testing.T) {
+
+		// no template
+		r, err := ctd.executeData(`{"trace_id":"asdf"}`)
+		assert.NoError(t, err)
+		assert.Equal(t, `{"trace_id":"asdf"}`, string(r))
+
+		rm, err := ctd.executeMetadata(`{"trace_id":"asdf"}`)
+		assert.NoError(t, err)
+		assert.Equal(t, map[string]string{"trace_id": "asdf"}, rm)
+
+		// new uuid
+		r, err = ctd.executeData(`{"trace_id":"{{newUUID}}"}`)
+		assert.NoError(t, err)
+		rs := strings.Replace(string(r), `{"trace_id":"`, "", -1)
+		rs = strings.Replace(rs, `"}`, "", -1)
+		assert.NotEmpty(t, rs)
+		parsed, err := uuid.Parse(rs)
+		assert.NoError(t, err)
+		assert.NotEmpty(t, parsed)
+
+		rm2, err := ctd.executeMetadata(`{"trace_id":"{{newUUID}}"}`)
+		assert.NoError(t, err)
+		rs2 := rm2["trace_id"]
+		assert.NotEmpty(t, rs)
+		assert.NotEqual(t, rs, rs2)
+		parsed, err = uuid.Parse(rs2)
+		assert.NoError(t, err)
+		assert.NotEmpty(t, parsed)
+
+		rm3, err := ctd.executeMetadata(`{"span_id":"{{newUUID}}","trace_id":"{{newUUID}}"}`)
+		assert.NoError(t, err)
+		assert.NotEmpty(t, rs)
+		assert.NotEqual(t, rm3["span_id"], rs2)
+		assert.NotEqual(t, rm3["span_id"], rs)
+		assert.NotEqual(t, rm3["trace_id"], rs2)
+		assert.NotEqual(t, rm3["trace_id"], rs)
+		assert.NotEqual(t, rm3["trace_id"], rm3["span_id"])
+		parsed, err = uuid.Parse(rm3["span_id"])
+		assert.NoError(t, err)
+		assert.NotEmpty(t, parsed)
+		parsed, err = uuid.Parse(rm3["trace_id"])
+		assert.NoError(t, err)
+		assert.NotEmpty(t, parsed)
+	})
+
+	t.Run("randomString", func(t *testing.T) {
+
+		// no template
+		r, err := ctd.executeData(`{"trace_id":"asdf"}`)
+		assert.NoError(t, err)
+		assert.Equal(t, `{"trace_id":"asdf"}`, string(r))
+
+		rm, err := ctd.executeMetadata(`{"trace_id":"asdf"}`)
+		assert.NoError(t, err)
+		assert.Equal(t, map[string]string{"trace_id": "asdf"}, rm)
+
+		// default length when 0
+		r, err = ctd.executeData(`{"trace_id":"{{randomString 0}}"}`)
+		assert.NoError(t, err)
+		rs := strings.Replace(string(r), `{"trace_id":"`, "", -1)
+		rs = strings.Replace(rs, `"}`, "", -1)
+		assert.NotEmpty(t, rs)
+		assert.True(t, len(rs) >= 2)
+		assert.True(t, len(rs) <= 16)
+
+		// default length when -1
+		r, err = ctd.executeData(`{"trace_id":"{{randomString -1}}"}`)
+		assert.NoError(t, err)
+		rs2 := strings.Replace(string(r), `{"trace_id":"`, "", -1)
+		rs2 = strings.Replace(rs2, `"}`, "", -1)
+		assert.NotEmpty(t, rs2)
+		assert.NotEqual(t, rs, rs2)
+		assert.True(t, len(rs2) >= 2)
+		assert.True(t, len(rs2) <= 16)
+
+		// specific length
+		r, err = ctd.executeData(`{"trace_id":"{{randomString 10}}"}`)
+		assert.NoError(t, err)
+		rs = strings.Replace(string(r), `{"trace_id":"`, "", -1)
+		rs = strings.Replace(rs, `"}`, "", -1)
+		assert.NotEmpty(t, rs)
+		assert.Len(t, rs, 10)
+
+		rm, err = ctd.executeMetadata(`{"trace_id":"{{randomString 0}}"}`)
+		assert.NoError(t, err)
+		assert.True(t, len(rm["trace_id"]) >= 2)
+		assert.True(t, len(rm["trace_id"]) <= 16)
+
+		rm, err = ctd.executeMetadata(`{"span_id":"{{randomString -1}}","trace_id":"{{randomString 0}}"}`)
+		assert.NoError(t, err)
+		assert.True(t, len(rm["trace_id"]) >= 2)
+		assert.True(t, len(rm["trace_id"]) <= 16)
+		assert.True(t, len(rm["span_id"]) >= 2)
+		assert.True(t, len(rm["span_id"]) <= 16)
+		assert.NotEqual(t, rm["trace_id"], rm["span_id"])
+
+		rm, err = ctd.executeMetadata(`{"span_id":"{{randomString 12}}","trace_id":"{{randomString 12}}"}`)
+		assert.NoError(t, err)
+		assert.Len(t, rm["trace_id"], 12)
+		assert.Len(t, rm["span_id"], 12)
+		assert.NotEqual(t, rm["trace_id"], rs)
+		assert.NotEqual(t, rm["trace_id"], rs2)
+		assert.NotEqual(t, rm["trace_id"], rm["span_id"])
+	})
+
+	t.Run("custom functions", func(t *testing.T) {
+		ctd = newCallTemplateData(md, template.FuncMap{
+			"getSKU": func() string {
+				return "custom-sku"
+			},
+			"newUUID": func() string {
+				return "custom-uuid"
+			},
+		}, "worker_id_123", 200)
+
+		r, err := ctd.executeData(`{"trace_id":"{{newUUID}}", "span_id":"{{getSKU}}"}`)
+		assert.NoError(t, err)
+		assert.Equal(t, `{"trace_id":"custom-uuid", "span_id":"custom-sku"}`, string(r))
+
+		rm, err := ctd.executeMetadata(`{"span_id":"{{randomString 12}}","trace_id":"{{newUUID}}", "sku":"{{getSKU}}"}`)
+		assert.NoError(t, err)
+		assert.Len(t, rm["span_id"], 12)
+		assert.Equal(t, "custom-uuid", rm["trace_id"])
+		assert.Equal(t, "custom-sku", rm["sku"])
+	})
+}