From ded6628b971754ae5643ae5a2f768e6d370be26b Mon Sep 17 00:00:00 2001
From: Vladislav Sukhin <vladislav@kubeshop.io>
Date: Fri, 13 Sep 2024 17:06:41 +0300
Subject: [PATCH] fix: dedup code

Signed-off-by: Vladislav Sukhin <vladislav@kubeshop.io>
---
 .../testworkflowprocessor/processor.go        | 33 ++++++++-----------
 .../testworkflowresolver/config.go            | 30 ++++++++++-------
 2 files changed, 32 insertions(+), 31 deletions(-)

diff --git a/pkg/testworkflows/testworkflowprocessor/processor.go b/pkg/testworkflows/testworkflowprocessor/processor.go
index 5f3914c438..67c3b30782 100644
--- a/pkg/testworkflows/testworkflowprocessor/processor.go
+++ b/pkg/testworkflows/testworkflowprocessor/processor.go
@@ -395,25 +395,8 @@ func (p *processor) Bundle(ctx context.Context, workflow *testworkflowsv1.TestWo
 	}
 	jobSpec.Spec.Template = podSpec
 
-	for i := range jobSpec.Spec.Template.Spec.InitContainers {
-		for envName, envSource := range mapEnv {
-			e := corev1.EnvVar{
-				Name:      envName,
-				ValueFrom: envSource.DeepCopy(),
-			}
-			jobSpec.Spec.Template.Spec.InitContainers[i].Env = append(jobSpec.Spec.Template.Spec.InitContainers[i].Env, e)
-		}
-	}
-
-	for i := range jobSpec.Spec.Template.Spec.Containers {
-		for envName, envSource := range mapEnv {
-			e := corev1.EnvVar{
-				Name:      envName,
-				ValueFrom: envSource.DeepCopy(),
-			}
-			jobSpec.Spec.Template.Spec.Containers[i].Env = append(jobSpec.Spec.Template.Spec.Containers[i].Env, e)
-		}
-	}
+	addEnvVarToContainerSpec(mapEnv, jobSpec.Spec.Template.Spec.InitContainers)
+	addEnvVarToContainerSpec(mapEnv, jobSpec.Spec.Template.Spec.Containers)
 
 	// Build running instructions
 	sigSerialized, _ := json.Marshal(sig)
@@ -436,3 +419,15 @@ func (p *processor) Bundle(ctx context.Context, workflow *testworkflowsv1.TestWo
 	}
 	return bundle, nil
 }
+
+func addEnvVarToContainerSpec(mapEnv map[string]corev1.EnvVarSource, containers []corev1.Container) {
+	for i := range containers {
+		for envName, envSource := range mapEnv {
+			e := corev1.EnvVar{
+				Name:      envName,
+				ValueFrom: envSource.DeepCopy(),
+			}
+			containers[i].Env = append(containers[i].Env, e)
+		}
+	}
+}
diff --git a/pkg/testworkflows/testworkflowresolver/config.go b/pkg/testworkflows/testworkflowresolver/config.go
index 1b13e36bfc..3770533f94 100644
--- a/pkg/testworkflows/testworkflowresolver/config.go
+++ b/pkg/testworkflows/testworkflowresolver/config.go
@@ -51,13 +51,9 @@ func createConfigMachine(cfg map[string]intstr.IntOrString, schema map[string]te
 			return nil, errors.Wrap(err, "config."+k)
 		}
 		if schema[k].Sensitive && externalize != nil {
-			envVar, err := externalize(k, expr.String())
+			expr, err = getSecretCallExpression(expr, k, externalize)
 			if err != nil {
-				return nil, errors.Wrap(err, "config."+k)
-			}
-			if envVar.SecretKeyRef != nil {
-				expr = expressions.NewValue(fmt.Sprintf("{{%ssecret(\"%s\", \"%s\")}}", expressions.InternalFnCall,
-					envVar.SecretKeyRef.Name, envVar.SecretKeyRef.Key))
+				return nil, err
 			}
 		}
 		machine.Register("config."+k, expr)
@@ -69,13 +65,9 @@ func createConfigMachine(cfg map[string]intstr.IntOrString, schema map[string]te
 				return nil, errors.Wrap(err, "config."+k)
 			}
 			if schema[k].Sensitive && externalize != nil {
-				envVar, err := externalize(k, expr.String())
+				expr, err = getSecretCallExpression(expr, k, externalize)
 				if err != nil {
-					return nil, errors.Wrap(err, "config."+k)
-				}
-				if envVar.SecretKeyRef != nil {
-					expr = expressions.NewValue(fmt.Sprintf("{{%ssecret(\"%s\", \"%s\")}}", expressions.InternalFnCall,
-						envVar.SecretKeyRef.Name, envVar.SecretKeyRef.Key))
+					return nil, err
 				}
 			}
 			machine.Register("config."+k, expr)
@@ -84,6 +76,20 @@ func createConfigMachine(cfg map[string]intstr.IntOrString, schema map[string]te
 	return machine, nil
 }
 
+func getSecretCallExpression(expr expressions.Expression, k string, externalize func(key, value string) (*corev1.EnvVarSource, error)) (
+	expressions.Expression, error) {
+	envVar, err := externalize(k, expr.String())
+	if err != nil {
+		return nil, errors.Wrap(err, "config."+k)
+	}
+	if envVar.SecretKeyRef != nil {
+		expr = expressions.NewValue(fmt.Sprintf("{{%ssecret(\"%s\", \"%s\")}}", expressions.InternalFnCall,
+			envVar.SecretKeyRef.Name, envVar.SecretKeyRef.Key))
+	}
+
+	return expr, nil
+}
+
 func ApplyWorkflowConfig(t *testworkflowsv1.TestWorkflow, cfg map[string]intstr.IntOrString,
 	externalize func(key, value string) (*corev1.EnvVarSource, error)) (*testworkflowsv1.TestWorkflow, error) {
 	if t == nil {