diff --git a/api/types/load_traffic.go b/api/types/load_traffic.go index d648224..72584e5 100644 --- a/api/types/load_traffic.go +++ b/api/types/load_traffic.go @@ -210,6 +210,12 @@ type RequestPostDel struct { } type PostDelPodOptions struct { + // Labels adds metadata labels to created Pods. + Labels map[string]string `json:"labels,omitempty" yaml:"labels,omitempty"` + // ServiceAccountName selects a pre-created ServiceAccount. + ServiceAccountName string `json:"serviceAccountName,omitempty" yaml:"serviceAccountName,omitempty"` + // SchedulerName selects a scheduler; an unhandled name keeps Pods unscheduled. + SchedulerName string `json:"schedulerName,omitempty" yaml:"schedulerName,omitempty"` // NodeSelector schedules workload Pods separately from runner Pods, for // example on a dedicated node pool: // @@ -440,14 +446,18 @@ func (r *RequestPostDel) Validate() error { return fmt.Errorf("payloadSize must >= 0: %v", r.PayloadSize) } - if r.Pod != nil && len(r.Pod.NodeSelector) > 0 { + if r.Pod != nil { if r.Group != "" || r.Version != "v1" || r.Resource != "pods" { - return fmt.Errorf("pod nodeSelector only supports core/v1/pods") + return fmt.Errorf("pod options only support core/v1/pods") } - - errs := metav1validation.ValidateLabels(r.Pod.NodeSelector, field.NewPath("pod", "nodeSelector")) - if len(errs) > 0 { - return fmt.Errorf("invalid pod nodeSelector: %v", errs.ToAggregate()) + if errs := metav1validation.ValidateLabels(r.Pod.Labels, field.NewPath("pod", "labels")); len(errs) > 0 { + return fmt.Errorf("invalid pod labels: %v", errs.ToAggregate()) + } + if len(r.Pod.NodeSelector) > 0 { + errs := metav1validation.ValidateLabels(r.Pod.NodeSelector, field.NewPath("pod", "nodeSelector")) + if len(errs) > 0 { + return fmt.Errorf("invalid pod nodeSelector: %v", errs.ToAggregate()) + } } } diff --git a/api/types/load_traffic_test.go b/api/types/load_traffic_test.go index fface1f..490f1ff 100644 --- a/api/types/load_traffic_test.go +++ b/api/types/load_traffic_test.go @@ -315,6 +315,32 @@ func TestRequestPostDelNodeSelectorValidation(t *testing.T) { } } +func TestRequestPostDelLabelsValidation(t *testing.T) { + for _, tt := range []struct { + name string + labels map[string]string + wantErr bool + }{ + {name: "omitted"}, + {name: "valid", labels: map[string]string{"example.com/workload": "benchmark", "app": ""}}, + {name: "invalid key", labels: map[string]string{"invalid key": "benchmark"}, wantErr: true}, + {name: "invalid value", labels: map[string]string{"app": "invalid value"}, wantErr: true}, + } { + t.Run(tt.name, func(t *testing.T) { + request := RequestPostDel{ + KubeGroupVersionResource: KubeGroupVersionResource{Version: "v1", Resource: "pods"}, + Pod: &PostDelPodOptions{Labels: tt.labels}, + } + err := request.Validate() + if tt.wantErr { + require.ErrorContains(t, err, "invalid pod labels") + } else { + require.NoError(t, err) + } + }) + } +} + func TestRequestPostDelNodeSelectorUnmarshalFromYAML(t *testing.T) { input := ` version: v1 diff --git a/contrib/internal/manifests/workload/pods/templates/pod.tpl b/contrib/internal/manifests/workload/pods/templates/pod.tpl index 2703889..08cb770 100644 --- a/contrib/internal/manifests/workload/pods/templates/pod.tpl +++ b/contrib/internal/manifests/workload/pods/templates/pod.tpl @@ -8,8 +8,19 @@ metadata: name: {{ $name }} namespace: {{ $namespace }} labels: + {{- range $key, $value := .Values.labels }} + "{{ $key }}": "{{ $value }}" + {{- end }} + {{- if eq .Values.labels.app nil }} app: fake-pod + {{- end }} spec: + {{- if .Values.serviceAccountName }} + serviceAccountName: "{{ .Values.serviceAccountName }}" + {{- end }} + {{- if .Values.schedulerName }} + schedulerName: "{{ .Values.schedulerName }}" + {{- end }} {{- if $nodeSelector }} nodeSelector: {{- range $key, $value := $nodeSelector }} diff --git a/request/random.go b/request/random.go index f5d1967..d2ce6c8 100644 --- a/request/random.go +++ b/request/random.go @@ -443,14 +443,17 @@ func (b *requestPatchBuilder) Build(cli rest.Interface) Requester { } type requestPostDelBuilder struct { - version schema.GroupVersion - resource string - resourceVersion string - namespace string - deleteRatio float64 - payloadSize int - nodeSelector map[string]string - maxRetries int + version schema.GroupVersion + resource string + resourceVersion string + namespace string + deleteRatio float64 + payloadSize int + nodeSelector map[string]string + labels map[string]string + serviceAccountName string + schedulerName string + maxRetries int // Per-builder cache for created resources cache *Cache @@ -460,22 +463,23 @@ type requestPostDelBuilder struct { } func newRequestPostDelBuilder(src *types.RequestPostDel, resourceVersion string, maxRetries int) *requestPostDelBuilder { - var nodeSelector map[string]string - if src.Pod != nil { - nodeSelector = src.Pod.NodeSelector - } - - return &requestPostDelBuilder{ + b := &requestPostDelBuilder{ version: schema.GroupVersion{Group: src.Group, Version: src.Version}, resource: src.Resource, resourceVersion: resourceVersion, namespace: src.Namespace, deleteRatio: src.DeleteRatio, payloadSize: src.PayloadSize, - nodeSelector: nodeSelector, maxRetries: maxRetries, cache: InitCache(), // Initialize the cache } + if p := src.Pod; p != nil { + b.nodeSelector = p.NodeSelector + b.labels = p.Labels + b.serviceAccountName = p.ServiceAccountName + b.schedulerName = p.SchedulerName + } + return b } // Build implements RequestBuilder.Build. @@ -524,10 +528,13 @@ func (b *requestPostDelBuilder) Build(cli rest.Interface) Requester { name := fmt.Sprintf("%d-%d", timestamp, counter) body, err := utils.RenderTemplate(b.resource, map[string]interface{}{ - "namePattern": name, - "namespace": b.namespace, - "payload": RandomPayload(b.payloadSize), - "nodeSelector": b.nodeSelector, + "namePattern": name, + "namespace": b.namespace, + "payload": RandomPayload(b.payloadSize), + "nodeSelector": b.nodeSelector, + "labels": b.labels, + "serviceAccountName": b.serviceAccountName, + "schedulerName": b.schedulerName, }) if err != nil { panic(fmt.Errorf("failed to render %s template: %w", b.resource, err)) diff --git a/request/random_test.go b/request/random_test.go index 81e6996..f5fbe7b 100644 --- a/request/random_test.go +++ b/request/random_test.go @@ -4,13 +4,23 @@ package request import ( + "context" + "encoding/json" + "io" + "net/http" "strconv" + "strings" "sync" "testing" "github.com/Azure/kperf/api/types" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + corev1 "k8s.io/api/core/v1" + "k8s.io/apimachinery/pkg/runtime/schema" + "k8s.io/client-go/kubernetes/scheme" + "k8s.io/client-go/rest/fake" ) func TestRenderPatchBody(t *testing.T) { @@ -126,3 +136,58 @@ func TestRequestPatchBuilderRendersConcurrencySafe(t *testing.T) { // Random bodies are effectively unique across all requests. assert.Equal(t, goroutines*perGoroutine, len(seen)) } + +func TestPodOptions(t *testing.T) { + for _, tt := range []struct { + name string + options *types.PostDelPodOptions + labels map[string]string + }{ + {name: "defaults", labels: map[string]string{"app": "fake-pod"}}, + {name: "empty app", options: &types.PostDelPodOptions{ + Labels: map[string]string{"app": ""}, + }, labels: map[string]string{"app": ""}}, + {name: "other labels", options: &types.PostDelPodOptions{ + Labels: map[string]string{"team": "benchmark"}, + }, labels: map[string]string{"app": "fake-pod", "team": "benchmark"}}, + {name: "string labels", options: &types.PostDelPodOptions{ + Labels: map[string]string{"true": "false", "123": "456", "null": "null"}, + }, labels: map[string]string{"app": "fake-pod", "true": "false", "123": "456", "null": "null"}}, + {name: "custom", options: &types.PostDelPodOptions{ + Labels: map[string]string{"app": "custom-app", "team": "benchmark"}, NodeSelector: map[string]string{"agentpool": "benchmark"}, + ServiceAccountName: "custom-account", SchedulerName: "custom-scheduler", + }, labels: map[string]string{"app": "custom-app", "team": "benchmark"}}, + } { + t.Run(tt.name, func(t *testing.T) { + builder := newRequestPostDelBuilder(&types.RequestPostDel{ + KubeGroupVersionResource: types.KubeGroupVersionResource{Version: "v1", Resource: "pods"}, + Namespace: "default", + Pod: tt.options, + }, "", 0) + var pod corev1.Pod + cli := &fake.RESTClient{ + GroupVersion: schema.GroupVersion{Version: "v1"}, + NegotiatedSerializer: scheme.Codecs.WithoutConversion(), + Client: fake.CreateHTTPClient(func(r *http.Request) (*http.Response, error) { + require.Equal(t, "POST", r.Method) + require.NoError(t, json.NewDecoder(r.Body).Decode(&pod)) + return &http.Response{ + StatusCode: http.StatusCreated, Header: http.Header{"Content-Type": {"application/json"}}, + Body: io.NopCloser(strings.NewReader("{}")), + }, nil + }), + } + _, err := builder.Build(cli).Do(context.Background()) + require.NoError(t, err) + require.Equal(t, "default", pod.Namespace) + options := tt.options + if options == nil { + options = &types.PostDelPodOptions{} + } + require.Equal(t, tt.labels, pod.Labels) + require.Equal(t, options.NodeSelector, pod.Spec.NodeSelector) + require.Equal(t, options.ServiceAccountName, pod.Spec.ServiceAccountName) + require.Equal(t, options.SchedulerName, pod.Spec.SchedulerName) + }) + } +}