feat: introduce readiness and liveness probe feature

cmd/helmify/flags.go

@@ -38,6 +38,7 @@ func ReadFlags() config.Config {
 	flag.BoolVar(&crd, "crd-dir", false, "Enable crd install into 'crds' directory.\nWarning: CRDs placed in 'crds' directory will not be templated by Helm.\nSee https://helm.sh/docs/chart_best_practices/custom_resource_definitions/#some-caveats-and-explanations\nExample: helmify -crd-dir")
 	flag.BoolVar(&result.ImagePullSecrets, "image-pull-secrets", false, "Allows the user to use existing secrets as imagePullSecrets in values.yaml")
 	flag.BoolVar(&result.GenerateDefaults, "generate-defaults", false, "Allows the user to add empty placeholders for tipical customization options in values.yaml. Currently covers: topology constraints, node selectors, tolerances")
+	flag.BoolVar(&result.Probes, "probes", false, "Allows the user to customize liveness and readiness probes")
 
 	flag.Parse()
 	if h || help {

pkg/config/config.go

@@ -26,6 +26,8 @@ type Config struct {
 	// GenerateDefaults enables the generation of empty values placeholders for common customization options of helm chart
 	// current generated values: tolerances, node selectors, topology constraints
 	GenerateDefaults bool
+	// Probes flag if true the probes will be parametrised
+	Probes bool
 }
 
 func (c *Config) Validate() error {

pkg/helm/init.go

@@ -112,6 +112,7 @@ Usage:
         {{- tpl (.value | toYaml) .context }}
     {{- end }}
 {{- end -}}
+
 `
 
 const defaultChartfile = `apiVersion: v2

pkg/processor/deployment/deployment.go

@@ -7,11 +7,11 @@ import (
 	"text/template"
 
 	"github.com/arttor/helmify/pkg/cluster"
+	"github.com/arttor/helmify/pkg/helmify"
 	"github.com/arttor/helmify/pkg/processor"
 	"github.com/arttor/helmify/pkg/processor/constraints"
 	"github.com/arttor/helmify/pkg/processor/imagePullSecrets"
-
-	"github.com/arttor/helmify/pkg/helmify"
+	"github.com/arttor/helmify/pkg/processor/probes"
 	yamlformat "github.com/arttor/helmify/pkg/yaml"
 	"github.com/iancoleman/strcase"
 	"github.com/pkg/errors"
@@ -51,7 +51,6 @@ const selectorTempl = `%[1]s
 const imagePullPolicyTemplate = "{{ .Values.%[1]s.%[2]s.imagePullPolicy }}"
 const envValue = "{{ .Values.%[1]s.%[2]s.%[3]s.%[4]s }}"
 
-
 // New creates processor for k8s Deployment resource.
 func New() helmify.Processor {
 	return &deployment{}
@@ -166,8 +165,15 @@ func (d deployment) Process(appMeta helmify.AppMetadata, obj *unstructured.Unstr
 	}
 
 	spec := constraints.ProcessSpecMap(nameCamel, specMap, &values, appMeta.Config().GenerateDefaults)
-	spec = strings.ReplaceAll(spec, "'", "")
 
+	if appMeta.Config().Probes {
+		spec, err = probes.ProcessSpecMap(nameCamel, specMap, &values)
+		if err != nil {
+			return true, nil, err
+		}
+	}
+
+	spec = strings.ReplaceAll(spec, "'", "")
 	return true, &result{
 		values: values,
 		data: struct {

pkg/processor/probes/probes.go

@@ -0,0 +1,104 @@
+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"
+)
+
+const livenessProbe = "livenessProbe"
+const readinessProbe = "readinessProbe"
+
+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" +
+	"\n{{- end }}"
+
+const readinessProbeTemplate = "\n{{- if .Values.%[1]s.%[2]s.readinessProbe }}\n" +
+	"readinessProbe: {{- include \"tplvalues.render\" (dict \"value\" .Values.%[1]s.%[2]s.readinessProbe \"context\" $) | nindent 10 }}\n" +
+	" {{- else }}\n" +
+	"readinessProbe:\n%[3]s" +
+	"\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) (string, error) {
+
+	cs, _, err := unstructured.NestedSlice(specMap, "containers")
+
+	if err != nil {
+		return "", err
+	}
+
+	strContainers := make([]interface{}, len(cs))
+	for i, c := range cs {
+		castedContainer := c.(map[string]interface{})
+		strContainers[i], err = setProbesTemplates(name, castedContainer, values)
+		if err != nil {
+			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 res, nil
+}
+
+func setProbesTemplates(name string, castedContainer map[string]interface{}, values *helmify.Values) (string, error) {
+
+	var ready, live string
+	var err error
+	if _, defined := castedContainer[livenessProbe]; defined {
+		live, err = setProbe(name, castedContainer, values, livenessProbe)
+		if err != nil {
+			return "", err
+		}
+		delete(castedContainer, livenessProbe)
+	}
+	if _, defined := castedContainer[readinessProbe]; defined {
+		ready, err = setProbe(name, castedContainer, values, readinessProbe)
+		if err != nil {
+			return "", err
+		}
+		delete(castedContainer, readinessProbe)
+	}
+	return setMap(name, castedContainer, live, ready)
+
+}
+
+func setMap(name string, castedContainer map[string]interface{}, live string, ready string) (string, error) {
+	containerName := strcase.ToLowerCamel(castedContainer["name"].(string))
+	content, err := yamlformat.Marshal(castedContainer, 0)
+	if err != nil {
+		return "", err
+	}
+	strContainer := string(content)
+	if live != "" {
+		strContainer = strContainer + fmt.Sprintf(livenessProbeTemplate, name, containerName, live)
+	}
+	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))
+	templatedProbe, err := yamlformat.Marshal(castedContainer[probe], 1)
+	if err != nil {
+		return "", err
+	}
+
+	return templatedProbe, unstructured.SetNestedField(*values, castedContainer[probe], name, containerName, probe)
+
+}

pkg/processor/probes/probes_test.go

@@ -0,0 +1,87 @@
+package probes
+
+import (
+	"testing"
+
+	"github.com/arttor/helmify/pkg/helmify"
+	"github.com/stretchr/testify/require"
+	"sigs.k8s.io/yaml"
+)
+
+func Test_setProbesTemplates(t *testing.T) {
+
+	tests := []struct {
+		name           string
+		deploymentName string
+		container      map[string]interface{}
+		wantMap        string
+		wantValue      string
+		wantErr        bool
+	}{
+		{
+			name:           "no probe no data generated",
+			deploymentName: "test",
+			container: map[string]interface{}{
+				"name": "mycontainer",
+			},
+			wantMap: "",
+			wantErr: false,
+		},
+		{
+			name:           "readinessProbe probe",
+			deploymentName: "test",
+			container: map[string]interface{}{
+				"name": "mycontainer",
+				readinessProbe: map[string]interface{}{
+					"timeoutSeconds": "1",
+					"periodSeconds":  "20",
+				},
+			},
+			wantMap: "\n{{- if .Values.test.mycontainer.readinessProbe }}\n" +
+				"readinessProbe: {{- include \"tplvalues.render\" (dict \"value\" .Values.test.mycontainer.readinessProbe \"context\" $) | nindent 10 }}\n {{- else }}\n" +
+				"readinessProbe:\n" +
+				" periodSeconds: \"20\"\n" +
+				" timeoutSeconds: \"1\"\n" +
+				"{{- end }}",
+			wantValue: "readinessProbe:\n  periodSeconds: \"20\"\n  timeoutSeconds: \"1\"\n",
+			wantErr:   false,
+		},
+		{
+			name:           "add livenessProbe probe",
+			deploymentName: "test",
+			container: map[string]interface{}{
+				"name": "mycontainer",
+				livenessProbe: map[string]interface{}{
+					"timeoutSeconds": "14",
+					"periodSeconds":  "2",
+				},
+			},
+			wantMap: "{{- if .Values.test.mycontainer.livenessProbe }}\n" +
+				"livenessProbe: {{- include \"tplvalues.render\" (dict \"value\" .Values.test.mycontainer.livenessProbe \"context\" $) | nindent 10 }}\n" +
+				" {{- else }}\nlivenessProbe:\n" +
+				" periodSeconds: \"2\"\n" +
+				" timeoutSeconds: \"14\"\n" +
+				"{{- end }}",
+			wantValue: "livenessProbe:\n  periodSeconds: \"2\"\n  timeoutSeconds: \"14\"\n",
+			wantErr:   false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			v := make(helmify.Values)
+			res, err := setProbesTemplates(tt.deploymentName, tt.container, &v)
+			require.True(t, (err != nil) == tt.wantErr)
+
+			require.Contains(t, res, tt.wantMap)
+			if tt.wantValue != "" {
+				val := (v)["test"].(map[string]interface{})["mycontainer"]
+				t.Log("VAL", val)
+				b, err := yaml.Marshal(val)
+				require.Nil(t, err)
+				require.Contains(t, string(b), tt.wantValue)
+			} else {
+				require.Empty(t, v)
+			}
+		})
+	}
+}