From b00f84794a1d45c2a20bfdf44cd0844cb1e6ce32 Mon Sep 17 00:00:00 2001
From: RealAnna <anna.reale@dynatrace.com>
Date: Mon, 20 Feb 2023 09:15:03 +0100
Subject: [PATCH] feat: introduced liveness and readiness probes

Signed-off-by: RealAnna <anna.reale@dynatrace.com>
---
 pkg/processor/deployment/deployment.go | 16 +++++++-------
 pkg/processor/probes/probes.go         | 29 ++++++++++++++++----------
 2 files changed, 25 insertions(+), 20 deletions(-)

diff --git a/pkg/processor/deployment/deployment.go b/pkg/processor/deployment/deployment.go
index 2e232528..266f5e1e 100644
--- a/pkg/processor/deployment/deployment.go
+++ b/pkg/processor/deployment/deployment.go
@@ -47,7 +47,7 @@ const selectorTempl = `%[1]s
 {{- include "%[2]s.selectorLabels" . | nindent 6 }}
 %[3]s`
 
-const envValue = "{{ .Values.%[1]s.%[2]s.%[3]s.%[4]s }}"
+const envValue = "{{ .Values.%[1]s.%[2]s.%[3]s }}"
 
 // New creates processor for k8s Deployment resource.
 func New() helmify.Processor {
@@ -162,21 +162,19 @@ 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
+	}
+
 	if appMeta.Config().Probes {
-		err = probes.ProcessSpecMap(nameCamel, specMap, &values)
+		spec, err = probes.ProcessSpecMap(nameCamel, specMap, &values)
 		if err != nil {
 			return true, nil, err
 		}
 	}
 
-	spec, err := yamlformat.Marshal(specMap, 6)
-	if err != nil {
-		return true, nil, err
-	}
-
 	spec = strings.ReplaceAll(spec, "'", "")
-	spec = strings.ReplaceAll(spec, "|\n        ", "")
-	spec = strings.ReplaceAll(spec, "|-\n        ", "")
 	return true, &result{
 		values: values,
 		data: struct {
diff --git a/pkg/processor/probes/probes.go b/pkg/processor/probes/probes.go
index 78346ecd..9dfc86d2 100644
--- a/pkg/processor/probes/probes.go
+++ b/pkg/processor/probes/probes.go
@@ -2,18 +2,18 @@ package probes
 
 import (
 	"fmt"
+	"strings"
 
 	"github.com/arttor/helmify/pkg/helmify"
 	yamlformat "github.com/arttor/helmify/pkg/yaml"
 	"github.com/iancoleman/strcase"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
-	"sigs.k8s.io/yaml"
 )
 
 const livenessProbe = "livenessProbe"
 const readinessProbe = "readinessProbe"
 
-const livenessProbeTemplate = "{{- if .Values.%[1]s.%[2]s.livenessProbe }}\n" +
+const livenessProbeTemplate = "\n{{- if .Values.%[1]s.%[2]s.livenessProbe }}\n" +
 	"livenessProbe: {{- include \"tplvalues.render\" (dict \"value\" .Values.%[1]s.%[2]s.livenessProbe \"context\" $) | nindent 10 }}\n" +
 	" {{- else }}\n" +
 	"livenessProbe:\n%[3]s" +
@@ -26,11 +26,12 @@ const readinessProbeTemplate = "\n{{- if .Values.%[1]s.%[2]s.readinessProbe }}\n
 	"\n{{- end }}"
 
 // ProcessSpecMap adds 'probes' to the Containers in specMap, if they are defined
-func ProcessSpecMap(name string, specMap map[string]interface{}, values *helmify.Values) error {
+func ProcessSpecMap(name string, specMap map[string]interface{}, values *helmify.Values) (string, error) {
 
 	cs, _, err := unstructured.NestedSlice(specMap, "containers")
+
 	if err != nil {
-		return err
+		return "", err
 	}
 
 	strContainers := make([]interface{}, len(cs))
@@ -38,12 +39,18 @@ func ProcessSpecMap(name string, specMap map[string]interface{}, values *helmify
 		castedContainer := c.(map[string]interface{})
 		strContainers[i], err = setProbesTemplates(name, castedContainer, values)
 		if err != nil {
-			return err
+			return "", err
 		}
 	}
+	specMap["containers"] = strContainers
+	specs, err := yamlformat.Marshal(specMap, 6)
+	if err != nil {
+		return "", err
+	}
+	res := strings.ReplaceAll(string(specs), "|\n        ", "")
+	res = strings.ReplaceAll(res, "|-\n        ", "")
 
-	return unstructured.SetNestedSlice(specMap, strContainers, "containers")
-
+	return res, nil
 }
 
 func setProbesTemplates(name string, castedContainer map[string]interface{}, values *helmify.Values) (string, error) {
@@ -70,8 +77,7 @@ func setProbesTemplates(name string, castedContainer map[string]interface{}, val
 
 func setMap(name string, castedContainer map[string]interface{}, live string, ready string) (string, error) {
 	containerName := strcase.ToLowerCamel(castedContainer["name"].(string))
-	content, err := yaml.Marshal(castedContainer)
-
+	content, err := yamlformat.Marshal(castedContainer, 0)
 	if err != nil {
 		return "", err
 	}
@@ -82,16 +88,17 @@ func setMap(name string, castedContainer map[string]interface{}, live string, re
 	if ready != "" {
 		strContainer = strContainer + fmt.Sprintf(readinessProbeTemplate, name, containerName, ready)
 	}
+
 	return strContainer, nil
 }
 
 func setProbe(name string, castedContainer map[string]interface{}, values *helmify.Values, probe string) (string, error) {
 	containerName := strcase.ToLowerCamel(castedContainer["name"].(string))
-	live, err := yamlformat.Marshal(castedContainer[probe], 1)
+	templatedProbe, err := yamlformat.Marshal(castedContainer[probe], 1)
 	if err != nil {
 		return "", err
 	}
 
-	return live, unstructured.SetNestedField(*values, castedContainer[probe], name, containerName, probe)
+	return templatedProbe, unstructured.SetNestedField(*values, castedContainer[probe], name, containerName, probe)
 
 }