From 07725af5758bdd48058f51e6926982c4ce3b149c Mon Sep 17 00:00:00 2001
From: Aleksei Gorobets <agorobets@d2iq.com>
Date: Fri, 15 Apr 2022 18:37:33 -0400
Subject: [PATCH 1/2] add handle Labels[label] and Annotations[annotation]

---
 .../v1beta1/experiment/validator/validator.go | 24 +++++++++++++++----
 .../experiment/validator/validator_test.go    | 22 +++++++++++++++++
 2 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/pkg/webhook/v1beta1/experiment/validator/validator.go b/pkg/webhook/v1beta1/experiment/validator/validator.go
index 5011e2e69f9..68f56f100a4 100644
--- a/pkg/webhook/v1beta1/experiment/validator/validator.go
+++ b/pkg/webhook/v1beta1/experiment/validator/validator.go
@@ -311,10 +311,7 @@ func (g *DefaultValidator) validateTrialTemplate(instance *experimentsv1beta1.Ex
 
 		// Check if parameter reference exist in experiment parameters
 		if len(experimentParameterNames) > 0 {
-			// Check if parameter is trial metadata
-			regex := regexp.MustCompile(consts.TrialTemplateMetaReplaceFormatRegex)
-			match := regex.FindStringSubmatch(parameter.Reference)
-			if !(len(match) > 0 && contains(consts.TrialTemplateMetaKeys, match[1])) {
+			if !isMetaKey(parameter.Reference) {
 				if _, ok := experimentParameterNames[parameter.Reference]; !ok {
 					return fmt.Errorf("parameter reference %v does not exist in spec.parameters: %v", parameter.Reference, instance.Spec.Parameters)
 				}
@@ -487,6 +484,25 @@ func (g *DefaultValidator) validateMetricsCollector(inst *experimentsv1beta1.Exp
 	return nil
 }
 
+func isMetaKey(parameter string) bool {
+	// Check if parameter is trial metadata reference as ${trailSpec.Name}, ${trialSpec.Labels[label]}, etc. used for substitution
+	match := regexp.MustCompile(consts.TrialTemplateMetaReplaceFormatRegex).FindStringSubmatch(parameter)
+	isMeta := false
+	if len(match) > 0 {
+		matchedKey := match[1]
+		if contains(consts.TrialTemplateMetaKeys, matchedKey) {
+			isMeta = true
+		} else {
+			// Check if it's Labels[label] or Annotations[annotation]
+			subMatch := regexp.MustCompile(consts.TrialTemplateMetaParseFormatRegex).FindStringSubmatch(matchedKey)
+			if len(subMatch) == 3 && contains(consts.TrialTemplateMetaKeys, subMatch[1]) {
+				isMeta = true
+			}
+		}
+	}
+	return isMeta
+}
+
 func contains(slice []string, item string) bool {
 	for _, s := range slice {
 		if s == item {
diff --git a/pkg/webhook/v1beta1/experiment/validator/validator_test.go b/pkg/webhook/v1beta1/experiment/validator/validator_test.go
index 67ae47e1558..dca7c5a6166 100644
--- a/pkg/webhook/v1beta1/experiment/validator/validator_test.go
+++ b/pkg/webhook/v1beta1/experiment/validator/validator_test.go
@@ -416,6 +416,8 @@ spec:
 	validTemplate5 := p.EXPECT().GetTrialTemplate(gomock.Any()).Return(validJobStr, nil)
 	validTemplate6 := p.EXPECT().GetTrialTemplate(gomock.Any()).Return(validJobStr, nil)
 	validTemplate7 := p.EXPECT().GetTrialTemplate(gomock.Any()).Return(validJobStr, nil)
+	validTemplate8 := p.EXPECT().GetTrialTemplate(gomock.Any()).Return(validJobStr, nil)
+	validTemplate9 := p.EXPECT().GetTrialTemplate(gomock.Any()).Return(validJobStr, nil)
 
 	missedParameterTemplate := p.EXPECT().GetTrialTemplate(gomock.Any()).Return(missedParameterJobStr, nil)
 	oddParameterTemplate := p.EXPECT().GetTrialTemplate(gomock.Any()).Return(oddParameterJobStr, nil)
@@ -433,6 +435,8 @@ spec:
 		validTemplate5,
 		validTemplate6,
 		validTemplate7,
+		validTemplate8,
+		validTemplate9,
 		missedParameterTemplate,
 		oddParameterTemplate,
 		invalidParameterTemplate,
@@ -582,6 +586,24 @@ spec:
 			Err:             false,
 			testDescription: "Trial template contains Trial metadata reference as parameter",
 		},
+		{
+			Instance: func() *experimentsv1beta1.Experiment {
+				i := newFakeInstance()
+				i.Spec.TrialTemplate.TrialParameters[1].Reference = "${trialSpec.Annotations[test-annotation]}"
+				return i
+			}(),
+			Err:             false,
+			testDescription: "Trial template contains Trial metadata reference as parameter",
+		},
+		{
+			Instance: func() *experimentsv1beta1.Experiment {
+				i := newFakeInstance()
+				i.Spec.TrialTemplate.TrialParameters[1].Reference = "${trialSpec.Labels[test-label]}"
+				return i
+			}(),
+			Err:             false,
+			testDescription: "Trial template contains Trial metadata reference as parameter",
+		},
 		// Trial Template doesn't contain parameter from trialParameters
 		// missedParameterTemplate case
 		{

From 72db2a5dd625ecd32910334f5ed233614ae5d2c7 Mon Sep 17 00:00:00 2001
From: Aleksei Gorobets <agorobets@d2iq.com>
Date: Fri, 15 Apr 2022 18:40:29 -0400
Subject: [PATCH 2/2] fix test description

---
 pkg/webhook/v1beta1/experiment/validator/validator_test.go | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pkg/webhook/v1beta1/experiment/validator/validator_test.go b/pkg/webhook/v1beta1/experiment/validator/validator_test.go
index dca7c5a6166..85ea6226e64 100644
--- a/pkg/webhook/v1beta1/experiment/validator/validator_test.go
+++ b/pkg/webhook/v1beta1/experiment/validator/validator_test.go
@@ -593,7 +593,7 @@ spec:
 				return i
 			}(),
 			Err:             false,
-			testDescription: "Trial template contains Trial metadata reference as parameter",
+			testDescription: "Trial template contains Trial annotation reference as parameter",
 		},
 		{
 			Instance: func() *experimentsv1beta1.Experiment {
@@ -602,7 +602,7 @@ spec:
 				return i
 			}(),
 			Err:             false,
-			testDescription: "Trial template contains Trial metadata reference as parameter",
+			testDescription: "Trial template contains Trial's label reference as parameter",
 		},
 		// Trial Template doesn't contain parameter from trialParameters
 		// missedParameterTemplate case