From 4915534337af31657c2cf37a39254a3425bd21bb Mon Sep 17 00:00:00 2001
From: RealAnna <anna.reale@dynatrace.com>
Date: Mon, 20 Feb 2023 11:05:16 +0100
Subject: [PATCH] feat:  added  constraints

Signed-off-by: RealAnna <anna.reale@dynatrace.com>
---
 pkg/processor/constraints/constraint.go | 34 ++++++++++++-------------
 pkg/processor/deployment/deployment.go  |  5 +---
 2 files changed, 17 insertions(+), 22 deletions(-)

diff --git a/pkg/processor/constraints/constraint.go b/pkg/processor/constraints/constraint.go
index 2865b41e..7e490d36 100644
--- a/pkg/processor/constraints/constraint.go
+++ b/pkg/processor/constraints/constraint.go
@@ -2,7 +2,7 @@ package constraints
 
 import (
 	"github.com/arttor/helmify/pkg/helmify"
-	corev1 "k8s.io/api/core/v1"
+	yamlformat "github.com/arttor/helmify/pkg/yaml"
 )
 
 const tolerations = "tolerations"
@@ -23,26 +23,24 @@ const tolerationsExpression = "{{- if .Values.tolerations }}\n" +
 
 // ProcessSpecMap adds 'topologyConstraints' to the podSpec in specMap, if it doesn't
 // already has one defined.
-func ProcessSpecMap(specMap string, values *helmify.Values, podspec corev1.PodSpec) string {
+func ProcessSpecMap(specMap map[string]interface{}, values *helmify.Values) string {
 
-	(*values)[topology] = podspec.TopologySpreadConstraints
-	(*values)[nodeSelector] = podspec.NodeSelector
-	(*values)[tolerations] = podspec.Tolerations
+	mapConstraint(specMap, topology, []interface{}{}, values)
+	mapConstraint(specMap, tolerations, []interface{}{}, values)
+	mapConstraint(specMap, nodeSelector, map[string]string{}, values)
 
-	tp := (*values)[topology].([]corev1.TopologySpreadConstraint)
-	if len(tp) == 0 {
-		(*values)[topology] = []interface{}{}
+	spec, err := yamlformat.Marshal(specMap, 6)
+	if err != nil {
+		return ""
 	}
-	ns := (*values)[nodeSelector].(map[string]string)
-	if len(ns) == 0 {
-		(*values)[nodeSelector] = map[string]string{}
-	}
-
-	tl := (*values)[tolerations].([]corev1.Toleration)
+	return spec + topologyExpression + nodeSelectorExpression + tolerationsExpression
+}
 
-	if len(tl) == 0 {
-		(*values)[tolerations] = []interface{}{}
+func mapConstraint(specMap map[string]interface{}, constraint string, override interface{}, values *helmify.Values) {
+	if specMap[constraint] != nil {
+		(*values)[constraint] = specMap[constraint].(interface{})
+	} else {
+		(*values)[constraint] = override
 	}
-
-	return specMap + topologyExpression + nodeSelectorExpression + tolerationsExpression
+	delete(specMap, constraint)
 }
diff --git a/pkg/processor/deployment/deployment.go b/pkg/processor/deployment/deployment.go
index 7a309a27..b3e3f209 100644
--- a/pkg/processor/deployment/deployment.go
+++ b/pkg/processor/deployment/deployment.go
@@ -163,10 +163,7 @@ func (d deployment) Process(appMeta helmify.AppMetadata, obj *unstructured.Unstr
 		imagePullSecrets.ProcessSpecMap(specMap, &values)
 	}
 
-	spec, err := yamlformat.Marshal(specMap, 6)
-	if err != nil {
-		return true, nil, err
-	}
+	spec := constraints.ProcessSpecMap(specMap, &values)
 	spec = strings.ReplaceAll(spec, "'", "")
 
 	return true, &result{