PLZ: configure custom image and pass env vars

api/v1alpha1/privateloadzone_types.go

@@ -36,6 +36,7 @@ type PrivateLoadZoneSpec struct {
 	Resources          corev1.ResourceRequirements `json:"resources"`
 	ServiceAccountName string                      `json:"serviceAccountName,omitempty"`
 	NodeSelector       map[string]string           `json:"nodeSelector,omitempty"`
+	Image              string                      `json:"image,omitempty"`
 }
 
 // PrivateLoadZoneStatus defines the observed state of PrivateLoadZone
@@ -81,6 +82,9 @@ func (plz *PrivateLoadZone) Register(ctx context.Context, logger logr.Logger, cl
 			CPU:    plz.Spec.Resources.Limits.Cpu().String(),
 			Memory: plz.Spec.Resources.Limits.Memory().String(),
 		},
+		LZConfig: cloud.LZConfig{
+			RunnerImage: plz.Spec.Image,
+		},
 		UID: uid,
 	}
 

config/crd/bases/k6.io_privateloadzones.yaml

@@ -26,6 +26,8 @@ spec:
             type: object
           spec:
             properties:
+              image:
+                type: string
               nodeSelector:
                 additionalProperties:
                   type: string

pkg/cloud/plz.go

@@ -15,6 +15,10 @@ const (
 func RegisterPLZ(client *cloudapi.Client, data PLZRegistrationData) error {
 	url := fmt.Sprintf("%s/cloud-resources/v1/load-zones", strings.TrimSuffix(client.BaseURL(), "/v1"))
 
+	data.LZConfig = LZConfig{
+		RunnerImage: data.RunnerImage,
+	}
+
 	req, err := client.NewRequest("POST", url, data)
 	if err != nil {
 		return err

pkg/cloud/types.go

@@ -6,6 +6,7 @@ import (
 	"go.k6.io/k6/cloudapi"
 	"go.k6.io/k6/lib/types"
 	"go.k6.io/k6/metrics"
+	corev1 "k8s.io/api/core/v1"
 )
 
 // InspectOutput is the parsed output from `k6 inspect --execution-requirements`.
@@ -65,15 +66,29 @@ type TestRunData struct {
 }
 
 type LZConfig struct {
-	RunnerImage   string `json:"load_runner_image"`
-	InstanceCount int    `json:"instance_count"`
-	ArchiveURL    string `json:"k6_archive_temp_public_url"`
+	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"`
 }
 
 func (trd *TestRunData) TestRunID() string {
 	return fmt.Sprintf("%d", trd.TestRunId)
 }
 
+func (lz *LZConfig) EnvVars() []corev1.EnvVar {
+	ev := make([]corev1.EnvVar, len(lz.Environment))
+	i := 0
+	for k, v := range lz.Environment {
+		ev[i] = corev1.EnvVar{
+			Name:  k,
+			Value: v,
+		}
+		i++
+	}
+	return ev
+}
+
 type TestRunStatus cloudapi.RunStatus
 
 func (trs TestRunStatus) Aborted() bool {
@@ -89,6 +104,9 @@ type PLZRegistrationData struct {
 	// defined by user as `name`
 	LoadZoneID string       `json:"k6_load_zone_id"`
 	Resources  PLZResources `json:"pod_tiers"`
+
+	LZConfig `json:"config"`
+
 	// Unique identifier of PLZ, generated by k6-operator
 	// during PLZ registration. It's purpose is to distinguish
 	// between PLZs with accidentally duplicate names.

pkg/testrun/plz.go

@@ -34,6 +34,11 @@ func NewPLZTestRun(plz *v1alpha1.PrivateLoadZone, trData *cloud.TestRunData, ing
 		volumeMount,
 	)
 
+	envVars := append(trData.EnvVars(), corev1.EnvVar{
+		Name:  "K6_CLOUD_HOST",
+		Value: ingestUrl,
+	})
+
 	return &v1alpha1.TestRun{
 		ObjectMeta: metav1.ObjectMeta{
 			Name:      TestName(trData.TestRunID()),
@@ -54,10 +59,7 @@ func NewPLZTestRun(plz *v1alpha1.PrivateLoadZone, trData *cloud.TestRunData, ing
 				InitContainers: []v1alpha1.InitContainer{
 					initContainer,
 				},
-				Env: []corev1.EnvVar{{
-					Name:  "K6_CLOUD_HOST",
-					Value: ingestUrl,
-				}},
+				Env: envVars,
 			},
 			Starter: v1alpha1.Pod{
 				ServiceAccountName: plz.Spec.ServiceAccountName,

pkg/testrun/plz_test.go

@@ -0,0 +1,232 @@
+package testrun
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/go-test/deep"
+	"github.com/grafana/k6-operator/api/v1alpha1"
+	"github.com/grafana/k6-operator/pkg/cloud"
+	"github.com/grafana/k6-operator/pkg/resources/containers"
+	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/api/resource"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+)
+
+func Test_NewPLZTestRun(t *testing.T) {
+	var (
+		mainIngest = "https://ingest.k6.io"
+
+		volumeMount = corev1.VolumeMount{
+			Name:      "archive-volume",
+			MountPath: "/test",
+		}
+		// zero-values test run definition
+		defaultTestRun = v1alpha1.TestRun{
+			ObjectMeta: metav1.ObjectMeta{
+				Name: TestName("0"),
+			},
+			Spec: v1alpha1.TestRunSpec{
+				Runner: v1alpha1.Pod{
+					Volumes: []corev1.Volume{{
+						Name: "archive-volume",
+						VolumeSource: corev1.VolumeSource{
+							EmptyDir: &corev1.EmptyDirVolumeSource{},
+						},
+					},
+					},
+					VolumeMounts: []corev1.VolumeMount{volumeMount},
+					InitContainers: []v1alpha1.InitContainer{
+						containers.NewS3InitContainer(
+							"",
+							"ghcr.io/grafana/k6-operator:latest-starter",
+							volumeMount,
+						),
+					},
+					Env: []corev1.EnvVar{{
+						Name:  "K6_CLOUD_HOST",
+						Value: mainIngest,
+					}},
+				},
+				Script: v1alpha1.K6Script{
+					LocalFile: "/test/archive.tar",
+				},
+				Parallelism: int32(0),
+				Separate:    false,
+				Arguments:   "--out cloud --no-thresholds",
+				Cleanup:     v1alpha1.Cleanup("post"),
+
+				TestRunID: "0",
+			},
+		}
+
+		// non-empty values to use int test cases
+		someToken        = "some-token"
+		someSA           = "some-service-account"
+		someNodeSelector = map[string]string{"foo": "bar"}
+		someNS           = "some-ns"
+		resourceLimits   = corev1.ResourceList{
+			corev1.ResourceCPU:    resource.MustParse("200m"),
+			corev1.ResourceMemory: resource.MustParse("1G"),
+		}
+		someTestRunID   = 6543
+		someRunnerImage = "grafana/k6:0.52.0"
+		someInstances   = 10
+		someArchiveURL  = "https://foo.s3.amazonaws.com"
+		someEnvVars     = map[string]string{
+			"ENV": "VALUE",
+			"foo": "bar",
+		}
+
+		// TestRuns expected in different cases;
+		// see how they are populated below
+		requiredFieldsTestRun = defaultTestRun
+		optionalFieldsTestRun = defaultTestRun //nolint:ineffassign
+		cloudFieldsTestRun    = defaultTestRun //nolint:ineffassign
+		cloudEnvVarsTestRun   = defaultTestRun //nolint:ineffassign
+	)
+
+	// populate TestRuns for different test cases
+
+	requiredFieldsTestRun.Spec.Token = someToken
+	requiredFieldsTestRun.Spec.Runner.Resources.Limits = resourceLimits
+
+	optionalFieldsTestRun = requiredFieldsTestRun // build up on top of required field case
+	optionalFieldsTestRun.Namespace = someNS
+	optionalFieldsTestRun.Spec.Runner.ServiceAccountName = someSA
+	optionalFieldsTestRun.Spec.Runner.NodeSelector = someNodeSelector
+	optionalFieldsTestRun.Spec.Starter.ServiceAccountName = someSA
+	optionalFieldsTestRun.Spec.Starter.NodeSelector = someNodeSelector
+
+	cloudFieldsTestRun = requiredFieldsTestRun // build up on top of required field case
+	cloudFieldsTestRun.ObjectMeta.Name = TestName(fmt.Sprintf("%d", someTestRunID))
+	cloudFieldsTestRun.Spec.TestRunID = fmt.Sprintf("%d", someTestRunID)
+	cloudFieldsTestRun.Spec.Runner.InitContainers = []v1alpha1.InitContainer{
+		containers.NewS3InitContainer(
+			someArchiveURL,
+			"ghcr.io/grafana/k6-operator:latest-starter",
+			volumeMount,
+		),
+	}
+	cloudFieldsTestRun.Spec.Runner.Image = someRunnerImage
+	cloudFieldsTestRun.Spec.Parallelism = int32(someInstances)
+
+	cloudEnvVarsTestRun = cloudFieldsTestRun // build up on top of cloud fields case
+	cloudEnvVarsTestRun.Spec.Runner.Env = []corev1.EnvVar{
+		{
+			Name:  "ENV",
+			Value: "VALUE",
+		},
+		{
+			Name:  "foo",
+			Value: "bar",
+		},
+		{
+			Name:  "K6_CLOUD_HOST",
+			Value: mainIngest,
+		},
+	}
+
+	testCases := []struct {
+		name      string
+		plz       *v1alpha1.PrivateLoadZone
+		cloudData *cloud.TestRunData
+		ingestUrl string
+		expected  *v1alpha1.TestRun
+	}{
+		{
+			name:      "empty input gets a zero-values TestRun",
+			plz:       &v1alpha1.PrivateLoadZone{},
+			cloudData: &cloud.TestRunData{},
+			ingestUrl: mainIngest,
+			expected:  &defaultTestRun,
+		},
+		{
+			name: "required fields in PLZ",
+			plz: &v1alpha1.PrivateLoadZone{
+				Spec: v1alpha1.PrivateLoadZoneSpec{
+					Token: someToken,
+					Resources: corev1.ResourceRequirements{
+						Limits: resourceLimits,
+					},
+				},
+			},
+			cloudData: &cloud.TestRunData{},
+			ingestUrl: mainIngest,
+			expected:  &requiredFieldsTestRun,
+		},
+		{
+			name: "optional fields in PLZ",
+			plz: &v1alpha1.PrivateLoadZone{
+				ObjectMeta: metav1.ObjectMeta{
+					Namespace: someNS,
+				},
+				Spec: v1alpha1.PrivateLoadZoneSpec{
+					Token: someToken,
+					Resources: corev1.ResourceRequirements{
+						Limits: resourceLimits,
+					},
+					ServiceAccountName: someSA,
+					NodeSelector:       someNodeSelector,
+				},
+			},
+			cloudData: &cloud.TestRunData{},
+			ingestUrl: mainIngest,
+			expected:  &optionalFieldsTestRun,
+		},
+		{
+			name: "basic cloud fields",
+			plz: &v1alpha1.PrivateLoadZone{
+				Spec: v1alpha1.PrivateLoadZoneSpec{
+					Token: someToken,
+					Resources: corev1.ResourceRequirements{
+						Limits: resourceLimits,
+					},
+				},
+			},
+			cloudData: &cloud.TestRunData{
+				TestRunId: someTestRunID,
+				LZConfig: cloud.LZConfig{
+					RunnerImage:   someRunnerImage,
+					InstanceCount: someInstances,
+					ArchiveURL:    someArchiveURL,
+				},
+			},
+			ingestUrl: mainIngest,
+			expected:  &cloudFieldsTestRun,
+		},
+		{
+			name: "cloud fields with env vars",
+			plz: &v1alpha1.PrivateLoadZone{
+				Spec: v1alpha1.PrivateLoadZoneSpec{
+					Token: someToken,
+					Resources: corev1.ResourceRequirements{
+						Limits: resourceLimits,
+					},
+				},
+			},
+			cloudData: &cloud.TestRunData{
+				TestRunId: someTestRunID,
+				LZConfig: cloud.LZConfig{
+					RunnerImage:   someRunnerImage,
+					InstanceCount: someInstances,
+					ArchiveURL:    someArchiveURL,
+					Environment:   someEnvVars,
+				},
+			},
+			ingestUrl: mainIngest,
+			expected:  &cloudEnvVarsTestRun,
+		},
+	}
+
+	for _, testCase := range testCases {
+		testCase := testCase
+		t.Run(testCase.name, func(t *testing.T) {
+			t.Parallel()
+			got := NewPLZTestRun(testCase.plz, testCase.cloudData, testCase.ingestUrl)
+			if diff := deep.Equal(got, testCase.expected); diff != nil {
+				t.Errorf("NewPLZTestRun returned unexpected data, diff: %s", diff)
+			}
+		})
+	}
+}