Add `bricks configure` command to bricks CLI (#18)

* bricks configure

* remove t.setenv

* Read token and host from stdin

* Update .vscode/testing.code-snippets

Co-authored-by: Serge Smertin <259697+nfx@users.noreply.github.com>

Co-authored-by: Serge Smertin <259697+nfx@users.noreply.github.com>

.vscode/testing.code-snippets

@@ -19,12 +19,12 @@
 	"Method with single error result": {
 		"prefix": "me",
 		"body": [
-		  "func (a $1) $2() error {",
-		  "    return nil",
-		  "}"
+			"func (a $1) $2() error {",
+			"    return nil",
+			"}"
 		],
 		"description": "Method with single error result"
-	  },
+	},
 	"if err != nil return err": {
 		"prefix": "ife",
 		"body": [
@@ -55,29 +55,38 @@
 	"assert.EqualError": {
 		"prefix": "aee",
 		"body": [
-		  "assert.EqualError(t, err, \"..\")"
+			"assert.EqualError(t, err, \"..\")"
 		],
 		"description": "assert.EqualError"
 	},
 	"assert.Equal": {
 		"prefix": "ae",
 		"body": [
-		  "assert.Equal(t, \"..\", $1)"
+			"assert.Equal(t, \"..\", $1)"
 		],
 		"description": "assert.Equal"
 	},
 	"assert.NoError": {
 		"prefix": "anoe",
 		"body": [
-		  "assert.NoError(t, err)"
+			"assert.NoError(t, err)"
 		],
 		"description": "assert.NoError"
 	},
 	"err :=": {
 		"prefix": "e",
 		"body": [
-		  "err := "
+			"err := "
 		],
 		"description": "err :="
-	  }
-}
+	},
+	"Golang Error Wrapper": {
+		"prefix": "ifew",
+		"body": [
+			"if err != nil {",
+			"  return fmt.Errorf(\"$1: %w\", err)",
+			"}$2"
+		],
+		"description": "Golang Error Wrapper"
+	}
+}

cmd/configure/configure.go

@@ -0,0 +1,141 @@
+package configure
+
+import (
+	"bytes"
+	"errors"
+	"fmt"
+	"os"
+	"path/filepath"
+
+	"github.com/databricks/bricks/cmd/prompt"
+	"github.com/databricks/bricks/cmd/root"
+	"github.com/databricks/bricks/project"
+	"github.com/spf13/cobra"
+	"gopkg.in/ini.v1"
+)
+
+type Configs struct {
+	Host  string `ini:"host"`
+	Token string `ini:"token"`
+}
+
+var noInteractive bool
+
+func (cfg *Configs) readFromStdin() error {
+	n, err := fmt.Scanf("%s %s\n", &cfg.Host, &cfg.Token)
+	if err != nil {
+		return err
+	}
+	if n != 2 {
+		return fmt.Errorf("exactly 2 arguments are required")
+	}
+	return nil
+}
+
+func (cfg *Configs) prompt() error {
+	res := prompt.Results{}
+	err := prompt.Questions{prompt.Text{
+		Key:   "host",
+		Label: "Databricks Host",
+		Default: func(res prompt.Results) string {
+			return cfg.Host
+		},
+		Callback: func(ans prompt.Answer, prj *project.Project, res prompt.Results) {
+			cfg.Host = ans.Value
+		},
+	}, prompt.Text{
+		Key:   "token",
+		Label: "Databricks Token",
+		Default: func(res prompt.Results) string {
+			return cfg.Token
+		},
+		Callback: func(ans prompt.Answer, prj *project.Project, res prompt.Results) {
+			cfg.Token = ans.Value
+		},
+	}}.Ask(res)
+	if err != nil {
+		return err
+	}
+
+	for _, answer := range res {
+		if answer.Callback != nil {
+			answer.Callback(answer, nil, res)
+		}
+	}
+
+	return nil
+}
+
+var configureCmd = &cobra.Command{
+	Use:   "configure",
+	Short: "Configure authentication",
+	RunE: func(cmd *cobra.Command, args []string) error {
+		var err error
+		path := os.Getenv("DATABRICKS_CONFIG_FILE")
+		if path == "" {
+			path, err = os.UserHomeDir()
+			if err != nil {
+				return fmt.Errorf("homedir: %w", err)
+			}
+		}
+		if filepath.Base(path) == ".databrickscfg" {
+			path = filepath.Dir(path)
+		}
+		err = os.MkdirAll(path, os.ModeDir|os.ModePerm)
+		if err != nil {
+			return fmt.Errorf("create config dir: %w", err)
+		}
+		cfgPath := filepath.Join(path, ".databrickscfg")
+		_, err = os.Stat(cfgPath)
+		if errors.Is(err, os.ErrNotExist) {
+			file, err := os.Create(cfgPath)
+			if err != nil {
+				return fmt.Errorf("create config file: %w", err)
+			}
+			file.Close()
+		} else if err != nil {
+			return fmt.Errorf("open config file: %w", err)
+		}
+
+		ini_cfg, err := ini.Load(cfgPath)
+		if err != nil {
+			return fmt.Errorf("load config file: %w", err)
+		}
+		cfg := &Configs{"", ""}
+		err = ini_cfg.MapTo(cfg)
+		if err != nil {
+			return fmt.Errorf("unmarshal loaded config: %w", err)
+		}
+
+		if noInteractive {
+			err = cfg.readFromStdin()
+		} else {
+			err = cfg.prompt()
+		}
+		if err != nil {
+			return fmt.Errorf("reading configs: %w", err)
+		}
+
+		var buffer bytes.Buffer
+		buffer.WriteString("[DEFAULT]\n")
+		err = ini_cfg.ReflectFrom(cfg)
+		if err != nil {
+			return fmt.Errorf("marshall config: %w", err)
+		}
+		_, err = ini_cfg.WriteTo(&buffer)
+		if err != nil {
+			return fmt.Errorf("write config to buffer: %w", err)
+		}
+		err = os.WriteFile(cfgPath, buffer.Bytes(), os.ModePerm)
+		if err != nil {
+			return fmt.Errorf("write congfig to file: %w", err)
+		}
+
+		return nil
+	},
+}
+
+func init() {
+	root.RootCmd.AddCommand(configureCmd)
+	configureCmd.Flags().BoolVar(&noInteractive, "no-interactive", false, "Don't show interactive prompts for inputs. Read directly from stdin")
+}

cmd/configure/configure_test.go

@@ -0,0 +1,101 @@
+package configure
+
+import (
+	"context"
+	"os"
+	"path/filepath"
+	"runtime"
+	"testing"
+
+	"github.com/databricks/bricks/cmd/root"
+	"github.com/databricks/bricks/tests"
+	"github.com/stretchr/testify/assert"
+	"gopkg.in/ini.v1"
+)
+
+func assertKeyValueInSection(t *testing.T, section *ini.Section, keyName, expectedValue string) {
+	key, err := section.GetKey(keyName)
+	assert.NoError(t, err)
+	assert.Equal(t, key.Value(), expectedValue)
+}
+
+func setup(t *testing.T) string {
+	tempHomeDir := t.TempDir()
+	homeEnvVar := "HOME"
+	if runtime.GOOS == "windows" {
+		homeEnvVar = "USERPROFILE"
+	}
+	tests.SetTestEnv(t, homeEnvVar, tempHomeDir)
+	tests.SetTestEnv(t, "DATABRICKS_CONFIG_FILE", "")
+	return tempHomeDir
+}
+
+func getTempFileWithContent(t *testing.T, tempHomeDir string, content string) *os.File {
+	inp, err := os.CreateTemp(tempHomeDir, "input")
+	assert.NoError(t, err)
+	_, err = inp.WriteString(content)
+	assert.NoError(t, err)
+	err = inp.Sync()
+	assert.NoError(t, err)
+	_, err = inp.Seek(0, 0)
+	assert.NoError(t, err)
+	return inp
+}
+
+func TestDefaultConfigureNoInteractive(t *testing.T) {
+	ctx := context.Background()
+	tempHomeDir := setup(t)
+	inp := getTempFileWithContent(t, tempHomeDir, "host token\n")
+	oldStdin := os.Stdin
+	defer func() { os.Stdin = oldStdin }()
+	os.Stdin = inp
+
+	root.RootCmd.SetArgs([]string{"configure", "--no-interactive"})
+
+	err := root.RootCmd.ExecuteContext(ctx)
+	assert.NoError(t, err)
+
+	cfgPath := filepath.Join(tempHomeDir, ".databrickscfg")
+	_, err = os.Stat(cfgPath)
+	assert.NoError(t, err)
+
+	cfg, err := ini.Load(cfgPath)
+	assert.NoError(t, err)
+
+	defaultSection, err := cfg.GetSection("DEFAULT")
+	assert.NoError(t, err)
+
+	assertKeyValueInSection(t, defaultSection, "host", "host")
+	assertKeyValueInSection(t, defaultSection, "token", "token")
+}
+
+func TestConfigFileFromEnvNoInteractive(t *testing.T) {
+	//TODO: Replace with similar test code from go SDK, once we start using it directly
+	ctx := context.Background()
+	tempHomeDir := setup(t)
+	cfgFileDir := filepath.Join(tempHomeDir, "test")
+	tests.SetTestEnv(t, "DATABRICKS_CONFIG_FILE", cfgFileDir)
+
+	inp := getTempFileWithContent(t, tempHomeDir, "host token\n")
+	oldStdin := os.Stdin
+	defer func() { os.Stdin = oldStdin }()
+	os.Stdin = inp
+
+	root.RootCmd.SetArgs([]string{"configure", "--no-interactive"})
+
+	err := root.RootCmd.ExecuteContext(ctx)
+	assert.NoError(t, err)
+
+	cfgPath := filepath.Join(cfgFileDir, ".databrickscfg")
+	_, err = os.Stat(cfgPath)
+	assert.NoError(t, err)
+
+	cfg, err := ini.Load(cfgPath)
+	assert.NoError(t, err)
+
+	defaultSection, err := cfg.GetSection("DEFAULT")
+	assert.NoError(t, err)
+
+	assertKeyValueInSection(t, defaultSection, "host", "host")
+	assertKeyValueInSection(t, defaultSection, "token", "token")
+}

cmd/init/init.go

@@ -6,6 +6,7 @@ import (
 	"os"
 	"path"
 
+	"github.com/databricks/bricks/cmd/prompt"
 	"github.com/databricks/bricks/cmd/root"
 	"github.com/databricks/bricks/project"
 	"github.com/ghodss/yaml"
@@ -30,23 +31,44 @@ var initCmd = &cobra.Command{
 			return err
 		}
 		wd, _ := os.Getwd()
-		q := Questions{
-			Text{"name", "Project name", func(res Results) string {
-				return path.Base(wd)
-			}, func(ans Answer, prj *project.Project, res Results) {
-				prj.Name = ans.Value
-			}},
+		q := prompt.Questions{
+			prompt.Text{
+				Key:   "name",
+				Label: "Project name",
+				Default: func(res prompt.Results) string {
+					return path.Base(wd)
+				},
+				Callback: func(ans prompt.Answer, prj *project.Project, res prompt.Results) {
+					prj.Name = ans.Value
+				},
+			},
 			*profileChoice,
-			Choice{"language", "Project language", Answers{
-				{"Python", "Machine learning and data engineering focused projects", nil},
-				{"Scala", "Data engineering focused projects with strong typing", nil},
+			prompt.Choice{Key: "language", Label: "Project language", Answers: prompt.Answers{
+				{
+					Value:    "Python",
+					Details:  "Machine learning and data engineering focused projects",
+					Callback: nil,
+				},
+				{
+					Value:    "Scala",
+					Details:  "Data engineering focused projects with strong typing",
+					Callback: nil,
+				},
 			}},
-			Choice{"isolation", "Deployment isolation", Answers{
-				{"None", "Use shared Databricks workspace resources for all project team members", nil},
-				{"Soft", "Prepend prefixes to each team member's deployment", func(
-					ans Answer, prj *project.Project, res Results) {
-					prj.Isolation = project.Soft
-				}},
+			prompt.Choice{Key: "isolation", Label: "Deployment isolation", Answers: prompt.Answers{
+				{
+					Value:    "None",
+					Details:  "Use shared Databricks workspace resources for all project team members",
+					Callback: nil,
+				},
+				{
+					Value:   "Soft",
+					Details: "Prepend prefixes to each team member's deployment",
+					Callback: func(
+						ans prompt.Answer, prj *project.Project, res prompt.Results) {
+						prj.Isolation = project.Soft
+					},
+				},
 			}},
 			// DBR selection
 			// Choice{"cloud", "Cloud", Answers{
@@ -65,7 +87,7 @@ var initCmd = &cobra.Command{
 			// 	{"PyCharm", "Create project conf and other things", nil},
 			// }},
 		}
-		res := Results{}
+		res := prompt.Results{}
 		err = q.Ask(res)
 		if err != nil {
 			return err

cmd/init/legacy-cli.go

@@ -3,12 +3,13 @@ package init
 import (
 	"fmt"
 
+	"github.com/databricks/bricks/cmd/prompt"
 	"github.com/databricks/bricks/project"
 	"github.com/mitchellh/go-homedir"
 	"gopkg.in/ini.v1"
 )
 
-func loadCliProfiles() (profiles []Answer, err error) {
+func loadCliProfiles() (profiles []prompt.Answer, err error) {
 	file, err := homedir.Expand("~/.databrickscfg")
 	if err != nil {
 		return
@@ -24,10 +25,10 @@ func loadCliProfiles() (profiles []Answer, err error) {
 			continue
 		}
 		// TODO: verify these tokens to work, becaus they may be expired
-		profiles = append(profiles, Answer{
+		profiles = append(profiles, prompt.Answer{
 			Value:   v.Name(),
 			Details: fmt.Sprintf(`Connecting to "%s" workspace`, host),
-			Callback: func(ans Answer, prj *project.Project, _ Results) {
+			Callback: func(ans prompt.Answer, prj *project.Project, _ prompt.Results) {
 				prj.Profile = ans.Value
 			},
 		})
@@ -35,14 +36,14 @@ func loadCliProfiles() (profiles []Answer, err error) {
 	return
 }
 
-func getConnectionProfile() (*Choice, error) {
+func getConnectionProfile() (*prompt.Choice, error) {
 	profiles, err := loadCliProfiles()
 	if err != nil {
 		return nil, err
 	}
 	// TODO: propmt for password and create ~/.databrickscfg
-	return &Choice{
-		key:     "profile",
+	return &prompt.Choice{
+		Key:     "profile",
 		Label:   "Databricks CLI profile",
 		Answers: profiles,
 	}, err

cmd/prompt/prompt.go

@@ -1,7 +1,8 @@
-package init
+package prompt
 
 import (
 	"fmt"
+	"io"
 
 	"github.com/databricks/bricks/project"
 	"github.com/manifoldco/promptui"
@@ -27,10 +28,11 @@ func (qq Questions) Ask(res Results) error {
 }
 
 type Text struct {
-	key      string
+	Key      string
 	Label    string
 	Default  func(res Results) string
 	Callback AnswerCallback
+	Stdin    io.ReadCloser
 }
 
 func (t Text) Ask(res Results) (string, Answer, error) {
@@ -41,17 +43,19 @@ func (t Text) Ask(res Results) (string, Answer, error) {
 	v, err := (&promptui.Prompt{
 		Label:   t.Label,
 		Default: def,
+		Stdin:   t.Stdin,
 	}).Run()
-	return t.key, Answer{
+	return t.Key, Answer{
 		Value:    v,
 		Callback: t.Callback,
 	}, err
 }
 
 type Choice struct {
-	key     string
+	Key     string
 	Label   string
 	Answers []Answer
+	Stdin   io.ReadCloser
 }
 
 func (q Choice) Ask(res Results) (string, Answer, error) {
@@ -64,9 +68,10 @@ func (q Choice) Ask(res Results) (string, Answer, error) {
 			Details:  `{{ .Details | green }}`,
 			Selected: fmt.Sprintf(`{{ "%s" | faint }}: {{ .Value | bold }}`, q.Label),
 		},
+		Stdin: q.Stdin,
 	}
 	i, _, err := prompt.Run()
-	return q.key, q.Answers[i], err
+	return q.Key, q.Answers[i], err
 }
 
 type Answers []Answer

go.mod

@@ -1,6 +1,6 @@
 module github.com/databricks/bricks
 
-go 1.16
+go 1.18
 
 require (
 	github.com/atotto/clipboard v0.1.4
@@ -20,3 +20,72 @@ require (
 	golang.org/x/mod v0.5.1 // BSD-3-Clause
 	gopkg.in/ini.v1 v1.67.0 // Apache 2.0
 )
+
+require (
+	cloud.google.com/go/compute v1.6.1 // indirect
+	github.com/Azure/go-autorest v14.2.0+incompatible // indirect
+	github.com/Azure/go-autorest/autorest v0.11.27 // indirect
+	github.com/Azure/go-autorest/autorest/adal v0.9.19 // indirect
+	github.com/Azure/go-autorest/autorest/azure/auth v0.5.11 // indirect
+	github.com/Azure/go-autorest/autorest/azure/cli v0.4.5 // indirect
+	github.com/Azure/go-autorest/autorest/date v0.3.0 // indirect
+	github.com/Azure/go-autorest/logger v0.2.1 // indirect
+	github.com/Azure/go-autorest/tracing v0.6.0 // indirect
+	github.com/agext/levenshtein v1.2.2 // indirect
+	github.com/apparentlymart/go-textseg/v13 v13.0.0 // indirect
+	github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e // indirect
+	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/dimchansky/utfbom v1.1.1 // indirect
+	github.com/fatih/color v1.7.0 // indirect
+	github.com/golang-jwt/jwt/v4 v4.4.1 // indirect
+	github.com/golang/groupcache v0.0.0-20200121045136-8c9f03a8e57e // indirect
+	github.com/golang/protobuf v1.5.2 // indirect
+	github.com/google/go-cmp v0.5.8 // indirect
+	github.com/google/go-querystring v1.1.0 // indirect
+	github.com/hashicorp/errwrap v1.0.0 // indirect
+	github.com/hashicorp/go-checkpoint v0.5.0 // indirect
+	github.com/hashicorp/go-cleanhttp v0.5.2 // indirect
+	github.com/hashicorp/go-cty v1.4.1-0.20200414143053-d3edf31b6320 // indirect
+	github.com/hashicorp/go-hclog v1.2.0 // indirect
+	github.com/hashicorp/go-multierror v1.1.1 // indirect
+	github.com/hashicorp/go-plugin v1.4.3 // indirect
+	github.com/hashicorp/go-retryablehttp v0.7.1 // indirect
+	github.com/hashicorp/go-uuid v1.0.3 // indirect
+	github.com/hashicorp/hcl/v2 v2.12.0 // indirect
+	github.com/hashicorp/logutils v1.0.0 // indirect
+	github.com/hashicorp/terraform-plugin-go v0.9.0 // indirect
+	github.com/hashicorp/terraform-plugin-log v0.4.0 // indirect
+	github.com/hashicorp/terraform-plugin-sdk/v2 v2.16.0 // indirect
+	github.com/hashicorp/terraform-registry-address v0.0.0-20210412075316-9b2996cce896 // indirect
+	github.com/hashicorp/terraform-svchost v0.0.0-20200729002733-f050f53b9734 // indirect
+	github.com/hashicorp/yamux v0.0.0-20181012175058-2f1d1f20f75d // indirect
+	github.com/inconshreveable/mousetrap v1.0.0 // indirect
+	github.com/mattn/go-colorable v0.1.4 // indirect
+	github.com/mattn/go-isatty v0.0.10 // indirect
+	github.com/mitchellh/copystructure v1.2.0 // indirect
+	github.com/mitchellh/go-testing-interface v1.14.1 // indirect
+	github.com/mitchellh/go-wordwrap v1.0.0 // indirect
+	github.com/mitchellh/mapstructure v1.5.0 // indirect
+	github.com/mitchellh/reflectwalk v1.0.2 // indirect
+	github.com/oklog/run v1.0.0 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	github.com/spf13/pflag v1.0.5 // indirect
+	github.com/vmihailenco/msgpack v4.0.4+incompatible // indirect
+	github.com/vmihailenco/msgpack/v4 v4.3.12 // indirect
+	github.com/vmihailenco/tagparser v0.1.1 // indirect
+	github.com/zclconf/go-cty v1.10.0 // indirect
+	go.opencensus.io v0.23.0 // indirect
+	golang.org/x/crypto v0.0.0-20211215153901-e495a2d5b3d3 // indirect
+	golang.org/x/net v0.0.0-20220425223048-2871e0cb64e4 // indirect
+	golang.org/x/oauth2 v0.0.0-20220411215720-9780585627b5 // indirect
+	golang.org/x/sys v0.0.0-20220503163025-988cb79eb6c6 // indirect
+	golang.org/x/text v0.3.7 // indirect
+	golang.org/x/time v0.0.0-20210723032227-1f47c861a9ac // indirect
+	google.golang.org/api v0.79.0 // indirect
+	google.golang.org/appengine v1.6.7 // indirect
+	google.golang.org/genproto v0.0.0-20220505152158-f39f71e6c8f3 // indirect
+	google.golang.org/grpc v1.46.0 // indirect
+	google.golang.org/protobuf v1.28.0 // indirect
+	gopkg.in/yaml.v2 v2.4.0 // indirect
+	gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c // indirect
+)

go.sum

@@ -26,7 +26,6 @@ cloud.google.com/go v0.93.3/go.mod h1:8utlLll2EF5XMAV15woO4lSbWQlk8rer9aLOfLh7+Y
 cloud.google.com/go v0.94.1/go.mod h1:qAlAugsXlC+JWO+Bke5vCtc9ONxjQT3drlTTnAplMW4=
 cloud.google.com/go v0.97.0/go.mod h1:GF7l59pYBVlXQIBLx3a761cZ41F9bBH3JUlihCt2Udc=
 cloud.google.com/go v0.99.0/go.mod h1:w0Xx2nLzqWJPuozYQX+hFfCSI8WioryfRDzkoI/Y2ZA=
-cloud.google.com/go v0.100.2 h1:t9Iw5QH5v4XtlEQaCtUY7x6sCABps8sW0acw7e2WQ6Y=
 cloud.google.com/go v0.100.2/go.mod h1:4Xra9TjzAeYHrl5+oeLlzbM2k3mjVhZh4UqTZ//w99A=
 cloud.google.com/go/bigquery v1.0.1/go.mod h1:i/xbL2UlR5RvWAURpBYZTtm/cXjCha9lbfbpx4poX+o=
 cloud.google.com/go/bigquery v1.3.0/go.mod h1:PjpwJnslEMmckchkHFfq+HTD2DmtT67aNFKH1/VBDHE=
@@ -93,7 +92,6 @@ github.com/apparentlymart/go-cidr v1.1.0/go.mod h1:EBcsNrHc3zQeuaeCeCtQruQm+n9/Y
 github.com/apparentlymart/go-dump v0.0.0-20180507223929-23540a00eaa3/go.mod h1:oL81AME2rN47vu18xqj1S1jPIPuN7afo62yKTNn3XMM=
 github.com/apparentlymart/go-dump v0.0.0-20190214190832-042adf3cf4a0 h1:MzVXffFUye+ZcSR6opIgz9Co7WcDx6ZcY+RjfFHoA0I=
 github.com/apparentlymart/go-dump v0.0.0-20190214190832-042adf3cf4a0/go.mod h1:oL81AME2rN47vu18xqj1S1jPIPuN7afo62yKTNn3XMM=
-github.com/apparentlymart/go-textseg v1.0.0 h1:rRmlIsPEEhUTIKQb7T++Nz/A5Q6C9IuX2wFoYVvnCs0=
 github.com/apparentlymart/go-textseg v1.0.0/go.mod h1:z96Txxhf3xSFMPmb5X/1W05FF/Nj9VFpLOpjS5yuumk=
 github.com/apparentlymart/go-textseg/v12 v12.0.0/go.mod h1:S/4uRK2UtaQttw1GenVJEynmyUenKwP++x/+DdGV/Ec=
 github.com/apparentlymart/go-textseg/v13 v13.0.0 h1:Y+KvPE1NYz0xl601PVImeQfFyEy6iT90AvPUL1NNfNw=

main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	_ "github.com/databricks/bricks/cmd/configure"
 	_ "github.com/databricks/bricks/cmd/fs"
 	_ "github.com/databricks/bricks/cmd/init"
 	_ "github.com/databricks/bricks/cmd/launch"

tests/utils.go

@@ -0,0 +1,23 @@
+package tests
+
+import (
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func SetTestEnv(t *testing.T, key, value string) {
+	originalValue, isSet := os.LookupEnv(key)
+
+	err := os.Setenv(key, value)
+	assert.NoError(t, err)
+
+	t.Cleanup(func() {
+		if isSet {
+			os.Setenv(key, originalValue)
+		} else {
+			os.Unsetenv(key)
+		}
+	})
+}