From 908b5552e4696a5a775e5b35bb9627f700df7efc Mon Sep 17 00:00:00 2001
From: Matheus Nogueira <matheus.nogueira2008@gmail.com>
Date: Fri, 3 Jun 2022 16:22:21 -0300
Subject: [PATCH] feat: support env variables in definition file (#658)

* feat: support env variables in definition file

* remove partial parallel support
---
 cli/file/definition.go                        | 26 +++++++++++
 cli/file/definition_test.go                   | 46 +++++++++++++++++++
 ...ttp_test_definition_with_env_variables.yml | 19 ++++++++
 3 files changed, 91 insertions(+)
 create mode 100644 cli/testdata/definitions/valid_http_test_definition_with_env_variables.yml

diff --git a/cli/file/definition.go b/cli/file/definition.go
index 86d6fc5276..ed5ef7c6a8 100644
--- a/cli/file/definition.go
+++ b/cli/file/definition.go
@@ -3,6 +3,8 @@ package file
 import (
 	"fmt"
 	"os"
+	"regexp"
+	"strings"
 
 	"github.com/kubeshop/tracetest/cli/definition"
 	"gopkg.in/yaml.v2"
@@ -14,6 +16,11 @@ func LoadDefinition(file string) (definition.Test, error) {
 		return definition.Test{}, fmt.Errorf("could not read test definition file %s: %w", file, err)
 	}
 
+	fileBytes, err = injectEnvVariables(fileBytes)
+	if err != nil {
+		return definition.Test{}, fmt.Errorf("could not inject env variables into definition file: %w", err)
+	}
+
 	test := definition.Test{}
 	err = yaml.Unmarshal(fileBytes, &test)
 	if err != nil {
@@ -36,3 +43,22 @@ func SaveDefinition(file string, definition definition.Test) error {
 
 	return nil
 }
+
+func injectEnvVariables(fileBytes []byte) ([]byte, error) {
+	envVarRegex, err := regexp.Compile(`\$\{\w+\}`)
+	if err != nil {
+		return []byte{}, fmt.Errorf("could not compile env variable regex: %w", err)
+	}
+
+	fileString := string(fileBytes)
+	allEnvVariables := envVarRegex.FindAllString(fileString, -1)
+
+	for _, envVariableExpression := range allEnvVariables {
+		envVarName := envVariableExpression[2 : len(envVariableExpression)-1] // removes '${' and '}'
+		envVarValue := os.Getenv(envVarName)
+
+		fileString = strings.Replace(fileString, envVariableExpression, envVarValue, -1)
+	}
+
+	return []byte(fileString), nil
+}
diff --git a/cli/file/definition_test.go b/cli/file/definition_test.go
index 384138581c..909efa419b 100644
--- a/cli/file/definition_test.go
+++ b/cli/file/definition_test.go
@@ -1,6 +1,7 @@
 package file_test
 
 import (
+	"os"
 	"testing"
 
 	"github.com/kubeshop/tracetest/cli/definition"
@@ -15,6 +16,7 @@ func TestLoadDefinition(t *testing.T) {
 		File               string
 		ExpectedDefinition definition.Test
 		ShouldSucceed      bool
+		EnvVariables       map[string]string
 	}{
 		{
 			Name:          "Should_parse_valid_definition_file",
@@ -131,10 +133,54 @@ func TestLoadDefinition(t *testing.T) {
 				},
 			},
 		},
+		{
+			Name:          "Should_parse_env_variables",
+			File:          "../testdata/definitions/valid_http_test_definition_with_env_variables.yml",
+			ShouldSucceed: true,
+			EnvVariables: map[string]string{
+				"POKEMON_APP_API_KEY": "my secret key",
+			},
+			ExpectedDefinition: definition.Test{
+				Name:        "POST import pokemon",
+				Description: "Import a pokemon using its ID",
+				Trigger: definition.TestTrigger{
+					Type: "http",
+					HTTPRequest: definition.HttpRequest{
+						URL:    "http://pokemon-demo.tracetest.io/pokemon/import",
+						Method: "POST",
+						Headers: []definition.HTTPHeader{
+							{Key: "Content-Type", Value: "application/json"},
+						},
+						Body: definition.HTTPBody{
+							Type: "raw",
+							Raw:  `{ "id": 52 }`,
+						},
+						Authentication: definition.HTTPAuthentication{
+							Type: "apiKey",
+							ApiKey: definition.HTTPAPIKeyAuth{
+								Key:   "X-Key",
+								Value: "my secret key",
+								In:    "header",
+							},
+						},
+					},
+				},
+			},
+		},
 	}
 
 	for _, testCase := range testCases {
 		t.Run(testCase.Name, func(t *testing.T) {
+			for key, value := range testCase.EnvVariables {
+				os.Setenv(key, value)
+			}
+
+			t.Cleanup(func() {
+				for key := range testCase.EnvVariables {
+					os.Unsetenv(key)
+				}
+			})
+
 			definition, err := file.LoadDefinition(testCase.File)
 			if testCase.ShouldSucceed {
 				require.NoError(t, err, "LoadDefinition should not fail")
diff --git a/cli/testdata/definitions/valid_http_test_definition_with_env_variables.yml b/cli/testdata/definitions/valid_http_test_definition_with_env_variables.yml
new file mode 100644
index 0000000000..bb66c9f3e5
--- /dev/null
+++ b/cli/testdata/definitions/valid_http_test_definition_with_env_variables.yml
@@ -0,0 +1,19 @@
+name: POST import pokemon
+description: Import a pokemon using its ID
+trigger:
+  type: http
+  http_request:
+    url: http://pokemon-demo.tracetest.io/pokemon/import
+    method: POST
+    headers:
+    - key: Content-Type
+      value: application/json
+    authentication:
+      type: apiKey
+      apiKey:
+        key: X-Key
+        value: ${POKEMON_APP_API_KEY}
+        in: header
+    body:
+      type: raw
+      raw: '{ "id": 52 }'
\ No newline at end of file