From 4e3bc9e63e36755802b1b8912308f943d5069712 Mon Sep 17 00:00:00 2001
From: Balazs Andorko <balazs.andorko@gmail.com>
Date: Fri, 1 Dec 2023 14:12:04 +0100
Subject: [PATCH] Support xk6 environment variables, explicitly go get xk6-g0

---
 README.md                   | 10 ++++++++++
 internal/builder/builder.go | 21 ++++++++++++++++++---
 2 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index 2a095a1..30a1c45 100644
--- a/README.md
+++ b/README.md
@@ -30,6 +30,16 @@ If you want to build only the custom k6 binary, then you can use the build subco
 k6-go build --output custom-k6 script.go
 ```
 
+## environment variables
+
+k6-go supports the following environment variables of xk6:
+
+- K6_VERSION
+- XK6_K6_REPO
+- XK6_RACE_DETECTOR
+- XK6_SKIP_CLEANUP
+
+More information : https://github.com/grafana/xk6#environment-variables
 ## example
 
 [petstore-load-k6](https://github.com/bandorko/petstore-load-k6) is an example load test project written in golang
diff --git a/internal/builder/builder.go b/internal/builder/builder.go
index 7aa0989..e11b582 100644
--- a/internal/builder/builder.go
+++ b/internal/builder/builder.go
@@ -36,7 +36,7 @@ func init() {
 }	
 `
 
-//Build builds the custom k6 binary including all the needed libraries for running the given go script.
+// Build builds the custom k6 binary including all the needed libraries for running the given go script.
 func Build(filename string, out string, silent bool) error {
 
 	origStrdOut := os.Stdout
@@ -71,6 +71,7 @@ func Build(filename string, out string, silent bool) error {
 		if err != nil {
 			return err
 		}
+
 		moduleName, err = getModuleName(moduleDir)
 		if err != nil {
 			return err
@@ -153,7 +154,11 @@ func createExportsFile(dir string, externalImports []string) error {
 }
 
 func generate(dir string) error {
-	_, err := exec.RunCommand(context.Background(), 0, dir, "go", "mod", "tidy")
+	_, err := exec.RunCommand(context.Background(), 0, dir, "go", "get", "github.com/szkiba/xk6-g0")
+	if err != nil {
+		return err
+	}
+	_, err = exec.RunCommand(context.Background(), 0, dir, "go", "mod", "tidy")
 	if err != nil {
 		return err
 	}
@@ -177,11 +182,17 @@ func getModuleName(dir string) (string, error) {
 }
 
 func build(modulename string, moduledir string, out string) error {
+	k6Version := os.Getenv("K6_VERSION")
+	k6Repo := os.Getenv("XK6_K6_REPO")
+	raceDetector := os.Getenv("XK6_RACE_DETECTOR") == "1"
+	skipCleanup := os.Getenv("XK6_SKIP_CLEANUP") == "1"
+
 	extensions := []xk6.Dependency{
 		{
 			PackagePath: "github.com/szkiba/xk6-g0",
 			Version:     "latest",
-		}}
+		},
+	}
 	replacements := []xk6.Replace{}
 	if modulename != "" {
 		extensions = append(extensions, xk6.Dependency{
@@ -191,6 +202,10 @@ func build(modulename string, moduledir string, out string) error {
 		replacements = append(replacements, xk6.NewReplace(modulename, moduledir))
 	}
 	builder := xk6.Builder{
+		K6Version:    k6Version,
+		K6Repo:       k6Repo,
+		RaceDetector: raceDetector,
+		SkipCleanup:  skipCleanup,
 		Extensions:   extensions,
 		Replacements: replacements,
 	}