PLZ: configuring CLI args and cleaner env vars

pkg/cloud/test_runs.go

@@ -103,13 +103,12 @@ func getTestRun(client *cloudapi.Client, url string) (*TestRunData, error) {
 	if err = client.Do(req, &trData); err != nil {
 		return nil, err
 	}
-
 	return &trData, nil
 }
 
 // called by PLZworker
 func GetTestRunData(client *cloudapi.Client, refID string) (*TestRunData, error) {
-	url := fmt.Sprintf("%s/loadtests/v4/test_runs(%s)?$select=id,run_status,k8s_load_zones_config,k6_runtime_config,test_run_token,secrets_config", strings.TrimSuffix(client.BaseURL(), "/v1"), refID)
+	url := fmt.Sprintf("%s/loadtests/v4/test_runs(%s)?$select=id,run_status,k8s_load_zones_config,k6_runtime_config,test_run_token,secrets_config,load_zone_distribution", strings.TrimSuffix(client.BaseURL(), "/v1"), refID)
 	return getTestRun(client, url)
 }
 

pkg/cloud/types.go

@@ -2,14 +2,38 @@ package cloud
 
 import (
 	"fmt"
+	"maps"
 	"sort"
+	"strings"
 
 	"go.k6.io/k6/cloudapi"
 	"go.k6.io/k6/lib/types"
 	"go.k6.io/k6/metrics"
 	corev1 "k8s.io/api/core/v1"
 )
 
+// GCk6 can set only a limited number of env vars to k6 process:
+// these are known and whitelisted with this "const" map.
+var reservedGCk6EnvVars = map[string]struct{}{
+	"K6_CLOUD_TOKEN": struct{}{},
+}
+
+const (
+	// Reserved vars set for PLZ tests, as described here:
+	// https://grafana.com/docs/grafana-cloud/testing/k6/author-run/cloud-scripting-extras/cloud-execution-context-variables/
+	// These are not passed from GCk6, but set by k6-operator directly.
+	lzCloudExecVar    = "K6_CLOUDRUN_LOAD_ZONE"
+	distrCloudExecVar = "K6_CLOUDRUN_DISTRIBUTION"
+	trIDCloudExecVar  = "K6_CLOUDRUN_TEST_RUN_ID"
+	// IIDCloudExecVar is exported as it must be set in external package, as part of TestRun CRD flow
+	IIDCloudExecVar = "K6_CLOUDRUN_INSTANCE_ID"
+
+	secretSourceEnvVar      = "K6_SECRET_SOURCE"
+	secretSourceURLTemplate = "K6_SECRET_SOURCE_URL_URL_TEMPLATE"
+	secretSourceURLRespPath = "K6_SECRET_SOURCE_URL_RESPONSE_PATH"
+	secretSourceURLAuthKey  = "K6_SECRET_SOURCE_URL_HEADER_AUTHORIZATION"
+)
+
 // InspectOutput is the parsed output from `k6 inspect --execution-requirements`.
 type InspectOutput struct {
 	External struct { // legacy way of defining the options.cloud
@@ -63,13 +87,6 @@ type SecretsConfig struct {
 	ResponsePath string `json:"response_path"`
 }
 
-const (
-	secretSourceEnvVar      = "K6_SECRET_SOURCE"
-	secretSourceURLTemplate = "K6_SECRET_SOURCE_URL_URL_TEMPLATE"
-	secretSourceURLRespPath = "K6_SECRET_SOURCE_URL_RESPONSE_PATH"
-	secretSourceURLAuthKey  = "K6_SECRET_SOURCE_URL_HEADER_AUTHORIZATION"
-)
-
 // TestRunData holds the output from /loadtests/v4/test_runs(%s)
 type TestRunData struct {
 	TestRunId     int `json:"id"`
@@ -78,41 +95,137 @@ type TestRunData struct {
 	RunStatus     cloudapi.RunStatus `json:"run_status"`
 	RuntimeConfig cloudapi.Config    `json:"k6_runtime_config"`
 	// SecretsToken is a short-lived, test-run-scoped token for read-only access to secrets.
-	SecretsToken  string             `json:"test_run_token,omitempty"`
-	SecretsConfig *SecretsConfig     `json:"secrets_config,omitempty"`
+	SecretsToken  string         `json:"test_run_token,omitempty"`
+	SecretsConfig *SecretsConfig `json:"secrets_config,omitempty"`
+	// LZDistribution holds label -> distribution mapping relevant
+	// for the given script and PLZ
+	LZDistribution `json:"load_zone_distribution,omitempty"`
+
+	// Pre-processed k6 arguments, populated by Preprocess().
+	TagArgs      string `json:"-"`
+	EnvArgs      string `json:"-"`
+	UserAgentArg string `json:"-"`
+}
+
+func (trd *TestRunData) TestRunID() string {
+	return fmt.Sprintf("%d", trd.TestRunId)
+}
+
+// Preprocess adds specific for GCk6 tags and env vars to data,
+// and produces sorted CLI argument strings for tags and environment.
+// Returns error if distribution is empty.
+func (trd *TestRunData) Preprocess() error {
+	if len(trd.LZDistribution) != 1 {
+		return fmt.Errorf("only tests with one load zone are supported, provided: %+v", trd.LZDistribution)
+	}
+
+	if len(trd.UserAgent) > 0 {
+		trd.UserAgentArg = fmt.Sprintf(`--user-agent="%s"`, trd.UserAgent)
+	}
+
+	if trd.Tags == nil {
+		trd.Tags = make(map[string]string)
+	}
+
+	if trd.Environment == nil {
+		trd.Environment = make(map[string]string)
+	}
+
+	// The potential overwrite here is deliberate: these keys are reserved by
+	// GCk6, described in docs and considered higher priority in PLZ tests
+	// for the sake of consistency between public & private cloud tests.
+
+	trd.Tags["load_zone"] = trd.LZLabel()
+
+	trd.Environment[lzCloudExecVar] = trd.LZName()
+	trd.Environment[distrCloudExecVar] = trd.LZLabel()
+	trd.Environment[trIDCloudExecVar] = trd.TestRunID()
+
+	trd.TagArgs = sortedArgs("--tag", trd.Tags)
+	trd.EnvArgs = sortedArgs("-e", trd.Environment)
+
+	return nil
+}
+
+// sortedArgs builds a CLI argument string from a map, sorted by key.
+func sortedArgs(flag string, m map[string]string) string {
+	keys := make([]string, 0, len(m))
+	for k := range m {
+		keys = append(keys, k)
+	}
+	sort.Strings(keys)
+
+	parts := make([]string, 0, len(keys))
+	for _, k := range keys {
+		parts = append(parts, flag+" "+k+"="+m[k])
+	}
+	return strings.Join(parts, " ")
 }
 
 type LZConfig struct {
-	RunnerImage   string            `json:"load_runner_image,omitempty"`
-	InstanceCount int               `json:"instance_count,omitempty"`
-	ArchiveURL    string            `json:"k6_archive_temp_public_url,omitempty"`
-	Environment   map[string]string `json:"environment,omitempty"`
+	RunnerImage   string `json:"load_runner_image,omitempty"`
+	InstanceCount int    `json:"instance_count,omitempty"`
+	ArchiveURL    string `json:"k6_archive_temp_public_url,omitempty"`
+	CLIArgs       `json:"cli_flags,omitempty"`
+	// Environment holds values passed by user via:
+	// 1. `-e` CLI option of k6
+	// 2. cloud environment variables of GCk6 -> Settings
+	// They are passed to k6 runners via `-e`
+	Environment map[string]string `json:"environment,omitempty"`
+	// GCk6EnvVars holds key-value pairs generated by GCk6 and
+	// meant to configure k6 process with reserved env vars.
+	GCk6EnvVars map[string]string `json:"gck6_env_vars,omitempty"`
 }
 
-func (trd *TestRunData) TestRunID() string {
-	return fmt.Sprintf("%d", trd.TestRunId)
+type CLIArgs struct {
+	BlacklistIPs         []string          `json:"blacklist_ips,omitempty"`
+	BlockedHostnames     []string          `json:"blocked_hostnames,omitempty"`
+	IncludeSystemEnvVars bool              `json:"include_system_env_vars,omitempty"`
+	Tags                 map[string]string `json:"tags,omitempty"`
+	UserAgent            string            `json:"user_agent,omitempty"`
 }
 
+type LZDistribution map[string]Distribution
+
+type Distribution struct {
+	LoadZone string `json:"loadZone"`
+	Percent  int    `json:"percent"`
+}
+
+// EnvVars makes up the corev1 struct from Go map.
 func (lz *LZConfig) EnvVars() []corev1.EnvVar {
-	ev := make([]corev1.EnvVar, len(lz.Environment))
+	whitelisted := maps.Collect(
+		func(yield func(_, _ string) bool) {
+			for k, v := range lz.GCk6EnvVars {
+				if _, ok := reservedGCk6EnvVars[k]; ok {
+					if !yield(k, v) {
+						return
+					}
+				}
+			}
+		},
+	)
+
+	ev := make([]corev1.EnvVar, len(whitelisted))
 	i := 0
-	for k, v := range lz.Environment {
+	for k, v := range whitelisted {
 		ev[i] = corev1.EnvVar{
 			Name:  k,
 			Value: v,
 		}
 		i++
 	}
+
 	// to have deterministic order in the array
 	sort.Slice(ev, func(i, j int) bool {
 		return ev[i].Name < ev[j].Name
 	})
-
 	return ev
 }
 
 // SecretsEnvVars returns the env vars required by the k6 URL secret source.
 // Returns nil when no secrets configuration is present.
+// TODO: make this private and move it to EnvVars() / Preprocess()
 func (trd *TestRunData) SecretsEnvVars() []corev1.EnvVar {
 	if trd.SecretsConfig == nil {
 		return nil
@@ -131,6 +244,22 @@ func (trd *TestRunData) SecretsEnvVars() []corev1.EnvVar {
 	return ev
 }
 
+// LZLabel assumes there is only one LZ.
+func (lzd *LZDistribution) LZLabel() string {
+	for k := range *lzd {
+		return k
+	}
+	return "unknown_lz_label"
+}
+
+// LZName assumes there is only one LZ.
+func (lzd *LZDistribution) LZName() string {
+	for _, v := range *lzd {
+		return v.LoadZone
+	}
+	return "unknown_lz_name"
+}
+
 type TestRunStatus cloudapi.RunStatus
 
 func (trs TestRunStatus) Aborted() bool {

pkg/cloud/types_test.go

@@ -111,6 +111,72 @@ func TestTestRunData_SecretsEnvVars(t *testing.T) {
 	}
 }
 
+func TestLZConfig_EnvVars(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name     string
+		lz       LZConfig
+		reserved map[string]struct{}
+		expected []corev1.EnvVar
+	}{
+		{
+			name:     "empty GCk6EnvVar",
+			lz:       LZConfig{},
+			expected: []corev1.EnvVar{},
+		},
+		{
+			name: "no keys match whitelist",
+			lz: LZConfig{
+				GCk6EnvVars: map[string]string{"NOT_RESERVED": "val"},
+			},
+			reserved: map[string]struct{}{"RESERVED_A": {}},
+			expected: []corev1.EnvVar{},
+		},
+		{
+			name: "all keys match whitelist",
+			lz: LZConfig{
+				GCk6EnvVars: map[string]string{"VAR_B": "b", "VAR_A": "a"},
+			},
+			reserved: map[string]struct{}{"VAR_A": {}, "VAR_B": {}},
+			expected: []corev1.EnvVar{
+				{Name: "VAR_A", Value: "a"},
+				{Name: "VAR_B", Value: "b"},
+			},
+		},
+		{
+			name: "some keys match whitelist",
+			lz: LZConfig{
+				GCk6EnvVars: map[string]string{"VAR_B": "b", "SKIP": "no", "VAR_A": "a"},
+			},
+			reserved: map[string]struct{}{"VAR_B": {}, "VAR_A": {}},
+			expected: []corev1.EnvVar{
+				{Name: "VAR_A", Value: "a"},
+				{Name: "VAR_B", Value: "b"},
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// Deliberately not parallel: subtests mutate package-level reservedGCk6EnvVars.
+			orig := reservedGCk6EnvVars
+			reservedGCk6EnvVars = tt.reserved
+			defer func() { reservedGCk6EnvVars = orig }()
+
+			got := tt.lz.EnvVars()
+			if len(got) != len(tt.expected) {
+				t.Fatalf("len = %d, want %d", len(got), len(tt.expected))
+			}
+			for i := range got {
+				if got[i] != tt.expected[i] {
+					t.Errorf("EnvVars()[%d] = %v, want %v", i, got[i], tt.expected[i])
+				}
+			}
+		})
+	}
+}
+
 func TestInspectOutput_SetTestName(t *testing.T) {
 	t.Parallel()
 

pkg/plz/worker.go

@@ -3,6 +3,8 @@ package plz
 import (
 	"context"
 	"fmt"
+	"slices"
+	"strings"
 
 	"github.com/go-logr/logr"
 	"github.com/grafana/k6-operator/api/v1alpha1"
@@ -185,10 +187,30 @@ func (w *PLZWorker) complete(tr *v1alpha1.TestRun, trData *cloud.TestRunData) {
 	}
 	tr.Spec.Runner.Env = envVars
 	tr.Spec.Parallelism = int32(trData.InstanceCount)
-	tr.Spec.Arguments = fmt.Sprintf(`--out cloud --no-thresholds --log-output=loki=https://cloudlogs.k6.io/api/v1/push,label.lz=%s,label.test_run_id=%s,header.Authorization="Token $(K6_CLOUD_TOKEN)"`,
-		w.plz.Name,
-		trData.TestRunID())
+
 	tr.Spec.TestRunID = trData.TestRunID()
+
+	// building argument list to k6
+	bips := `--blacklist-ip="` + strings.Join(trData.BlacklistIPs, ",") + `"`
+	bhns := `--block-hostnames="` + strings.Join(trData.BlockedHostnames, ",") + `"`
+
+	args := []string{
+		"--out cloud",
+		bips,
+		bhns,
+		trData.TagArgs,
+		"--no-thresholds",
+		trData.UserAgentArg,
+		fmt.Sprintf(`--log-output=loki=https://cloudlogs.k6.io/api/v1/push,label.lz=%s,label.test_run_id=%s,header.Authorization="Token $(K6_CLOUD_TOKEN)"`, w.plz.Name, trData.TestRunID()),
+		trData.EnvArgs,
+	}
+	if trData.IncludeSystemEnvVars {
+		args = append(args, "--include-system-env-vars", "--verbose")
+	}
+
+	args = slices.DeleteFunc(args, func(s string) bool { return s == "" })
+	tr.Spec.Arguments = strings.Join(args, " ")
+
 }
 
 // handle creates a new PLZ TestRun from the given test run id
@@ -214,6 +236,11 @@ func (w *PLZWorker) handle(testRunId string) {
 		w.logger.Error(err, fmt.Sprintf("Failed to retrieve test run data for `%s`", testRunId))
 		return
 	}
+	if err = trData.Preprocess(); err != nil {
+		w.logger.Error(err, fmt.Sprintf("Failed to sort out test run data for `%s`", testRunId))
+		return
+	}
+
 	w.complete(tr, trData)
 
 	w.logger.Info(fmt.Sprintf("PLZ test run has been prepared with image `%s` and `%d` instances",