fix(experiments): propagate rolouts labels to experiments and replica…

… sets

Enhancement ticket: argoproj#4113
This can be useful for complying to Tagging Policy imposed at cluster level
since there is no mechanism to add custom labels to Experiments and their
associated ReplicaSets.
Indeed, this could enable automatic reporting but also give required information
for investigation.

Signed-off-by: Ahmed YUREIDINI NOGALES <[email protected]>

experiments/replicaset.go

@@ -166,17 +166,21 @@ func newReplicaSetFromTemplate(experiment *v1alpha1.Experiment, template v1alpha
 	newRSTemplate.Labels = labelsutil.CloneAndAddLabel(newRSTemplate.Labels, v1alpha1.DefaultRolloutUniqueLabelKey, podHash)
 	// Add podTemplateHash label to selector.
 	newRSSelector := labelsutil.CloneSelectorAndAddLabel(template.Selector, v1alpha1.DefaultRolloutUniqueLabelKey, podHash)
+	newRSLabels := map[string]string{}
+	// enrich with template labels
+	for k, v := range newRSTemplate.Labels {
+		newRSLabels[k] = v
+	}
+	newRSLabels[v1alpha1.DefaultRolloutUniqueLabelKey] = podHash
 
 	// The annotations must be different for each template because annotations are used to match
 	// replicasets to templates. We inject the experiment and template name in the replicaset
 	// annotations to ensure uniqueness.
 	rs := appsv1.ReplicaSet{
 		ObjectMeta: metav1.ObjectMeta{
-			Name:      fmt.Sprintf("%s-%s", experiment.Name, template.Name),
-			Namespace: experiment.Namespace,
-			Labels: map[string]string{
-				v1alpha1.DefaultRolloutUniqueLabelKey: podHash,
-			},
+			Name:            fmt.Sprintf("%s-%s", experiment.Name, template.Name),
+			Namespace:       experiment.Namespace,
+			Labels:          newRSLabels,
 			OwnerReferences: []metav1.OwnerReference{*metav1.NewControllerRef(experiment, controllerKind)},
 			Annotations:     replicaSetAnnotations,
 		},

experiments/replicaset_test.go

@@ -1,6 +1,7 @@
 package experiments
 
 import (
+	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -164,3 +165,35 @@ func TestNameCollisionWithEquivalentPodTemplateAndControllerUID(t *testing.T) {
 		validatePatch(t, patch, "", NoChange, templateStatuses, cond)
 	}
 }
+
+// TestNewReplicaSetFromTemplate tests the creation of a new ReplicaSet from a given template.
+// It verifies that the ReplicaSet is correctly initialized with the expected name, namespace,
+// annotations, labels, and container specifications based on the provided experiment and template.
+// The test ensures that:
+// - The ReplicaSet name is a combination of the experiment name and template name.
+// - The ReplicaSet namespace matches the experiment namespace.
+// - The ReplicaSet annotations include the experiment name and template name.
+// - The ReplicaSet labels include the default rollout unique label key and a specific key from the template.
+// - The ReplicaSet selector and template labels include the default rollout unique label key.
+// - The ReplicaSet container specifications match those defined in the template.
+func TestNewReplicaSetFromTemplate(t *testing.T) {
+
+	templates := generateTemplates("bar")
+	template := templates[0]
+	experiment := newExperiment("foo", templates, "")
+	collisionCount := int32(0)
+	rs := newReplicaSetFromTemplate(experiment, template, &collisionCount)
+
+	assert.Equal(t, fmt.Sprintf("%s-%s", experiment.Name, template.Name), rs.Name)
+	assert.Equal(t, experiment.Namespace, rs.Namespace)
+	assert.Equal(t, experiment.Name, rs.Annotations[v1alpha1.ExperimentNameAnnotationKey])
+	assert.NotNil(t, rs.ObjectMeta.Labels[v1alpha1.DefaultRolloutUniqueLabelKey])
+	assert.NotNil(t, rs.ObjectMeta.Labels["key"])
+	assert.Equal(t, template.Template.ObjectMeta.Labels["key"], rs.ObjectMeta.Labels["key"])
+	assert.Equal(t, template.Name, rs.Annotations[v1alpha1.ExperimentTemplateNameAnnotationKey])
+	assert.NotNil(t, rs.Spec.Selector.MatchLabels[v1alpha1.DefaultRolloutUniqueLabelKey])
+	assert.NotNil(t, rs.Spec.Template.ObjectMeta.Labels[v1alpha1.DefaultRolloutUniqueLabelKey])
+	assert.Equal(t, template.Template.Labels["key"], rs.Spec.Template.Labels["key"])
+	assert.Equal(t, template.Template.Spec.Containers[0].Name, rs.Spec.Template.Spec.Containers[0].Name)
+	assert.Equal(t, template.Template.Spec.Containers[0].Image, rs.Spec.Template.Spec.Containers[0].Image)
+}

rollout/canary_test.go

@@ -33,7 +33,9 @@ import (
 
 func newCanaryRollout(name string, replicas int, revisionHistoryLimit *int32, steps []v1alpha1.CanaryStep, stepIndex *int32, maxSurge, maxUnavailable intstr.IntOrString) *v1alpha1.Rollout {
 	selector := map[string]string{"foo": "bar"}
+	labels := map[string]string{"custom": "label"}
 	rollout := newRollout(name, replicas, revisionHistoryLimit, selector)
+	rollout.ObjectMeta.Labels = labels
 	rollout.Spec.Strategy.Canary = &v1alpha1.CanaryStrategy{
 		MaxUnavailable: &maxUnavailable,
 		MaxSurge:       &maxSurge,

rollout/experiment.go

@@ -35,14 +35,19 @@ func GetExperimentFromTemplate(r *v1alpha1.Rollout, stableRS, newRS *appsv1.Repl
 	if r.Annotations != nil {
 		revision = r.Annotations[annotations.RevisionAnnotation]
 	}
+	newExperimentLabels := map[string]string{}
+	// enrich with template labels
+	for k, v := range r.Labels {
+		newExperimentLabels[k] = v
+	}
+	newExperimentLabels[v1alpha1.DefaultRolloutUniqueLabelKey] = podHash
+
 	experiment := &v1alpha1.Experiment{
 		ObjectMeta: metav1.ObjectMeta{
 			Name:            fmt.Sprintf("%s-%s-%s-%d", r.Name, podHash, revision, currentStep),
 			Namespace:       r.Namespace,
 			OwnerReferences: []metav1.OwnerReference{*metav1.NewControllerRef(r, controllerKind)},
-			Labels: map[string]string{
-				v1alpha1.DefaultRolloutUniqueLabelKey: podHash,
-			},
+			Labels:          newExperimentLabels,
 			Annotations: map[string]string{
 				annotations.RevisionAnnotation: revision,
 			},

rollout/experiment_test.go

@@ -57,6 +57,7 @@ func TestRolloutCreateExperiment(t *testing.T) {
 	f.run(getKey(r2, t))
 	createdEx := f.getCreatedExperiment(createExIndex)
 	assert.Equal(t, createdEx.Name, ex.Name)
+	assert.Equal(t, createdEx.ObjectMeta.Labels, ex.ObjectMeta.Labels)
 	assert.Equal(t, createdEx.Spec.Analyses[0].TemplateName, at.Name)
 	assert.Equal(t, createdEx.Spec.Analyses[0].Name, "test")
 	patch := f.getPatchedRollout(patchIndex)