From b7adbc2b63f99ee2ad763ddb35f07844b9b6f500 Mon Sep 17 00:00:00 2001
From: Muhammad Hamza Zaib <hamzazaib3202@gmail.com>
Date: Thu, 25 Mar 2021 14:58:03 +0100
Subject: [PATCH] Add labels to deployment, statefulset and service accounts

---
 templates/injector-deployment.yaml     |  5 ++++-
 templates/injector-serviceaccount.yaml |  3 +++
 templates/server-serviceaccount.yaml   |  3 +++
 templates/server-statefulset.yaml      |  3 +++
 test/unit/injector-deployment.bats     | 11 +++++++++--
 test/unit/injector-serviceaccount.bats | 13 +++++++++++++
 test/unit/server-serviceaccount.bats   | 15 ++++++++++++++-
 test/unit/server-statefulset.bats      | 11 +++++++++--
 8 files changed, 58 insertions(+), 6 deletions(-)

diff --git a/templates/injector-deployment.yaml b/templates/injector-deployment.yaml
index 95e2c2da8..ae4823c69 100644
--- a/templates/injector-deployment.yaml
+++ b/templates/injector-deployment.yaml
@@ -10,6 +10,9 @@ metadata:
     app.kubernetes.io/instance: {{ .Release.Name }}
     app.kubernetes.io/managed-by: {{ .Release.Service }}
     component: webhook
+    {{- if .Values.injector.extraLabels -}}
+      {{- toYaml .Values.injector.extraLabels | nindent 4 -}}
+    {{- end }}
 spec:
   replicas: {{ .Values.injector.replicas }}
   selector:
@@ -24,7 +27,7 @@ spec:
         app.kubernetes.io/name: {{ template "vault.name" . }}-agent-injector
         app.kubernetes.io/instance: {{ .Release.Name }}
         component: webhook
-        {{- if  .Values.injector.extraLabels -}}
+        {{- if .Values.injector.extraLabels -}}
           {{- toYaml .Values.injector.extraLabels | nindent 8 -}}
         {{- end -}}
       {{ template "injector.annotations" . }}
diff --git a/templates/injector-serviceaccount.yaml b/templates/injector-serviceaccount.yaml
index a28d38fab..d826128af 100644
--- a/templates/injector-serviceaccount.yaml
+++ b/templates/injector-serviceaccount.yaml
@@ -8,4 +8,7 @@ metadata:
     app.kubernetes.io/name: {{ include "vault.name" . }}-agent-injector
     app.kubernetes.io/instance: {{ .Release.Name }}
     app.kubernetes.io/managed-by: {{ .Release.Service }}
+    {{- if .Values.injector.extraLabels -}}
+      {{- toYaml .Values.injector.extraLabels | nindent 4 -}}
+    {{- end -}}
 {{ end }}
diff --git a/templates/server-serviceaccount.yaml b/templates/server-serviceaccount.yaml
index 925b166bb..13e6cf758 100644
--- a/templates/server-serviceaccount.yaml
+++ b/templates/server-serviceaccount.yaml
@@ -11,6 +11,9 @@ metadata:
     app.kubernetes.io/name: {{ include "vault.name" . }}
     app.kubernetes.io/instance: {{ .Release.Name }}
     app.kubernetes.io/managed-by: {{ .Release.Service }}
+    {{- if .Values.server.extraLabels -}}
+      {{- toYaml .Values.server.extraLabels | nindent 4 -}}
+    {{- end -}}
   {{ template "vault.serviceAccount.annotations" . }}
 {{ end }}
 {{ end }}
diff --git a/templates/server-statefulset.yaml b/templates/server-statefulset.yaml
index cbcda967d..ecbd9fb09 100644
--- a/templates/server-statefulset.yaml
+++ b/templates/server-statefulset.yaml
@@ -11,6 +11,9 @@ metadata:
     app.kubernetes.io/name: {{ include "vault.name" . }}
     app.kubernetes.io/instance: {{ .Release.Name }}
     app.kubernetes.io/managed-by: {{ .Release.Service }}
+    {{- if .Values.server.extraLabels -}}
+      {{- toYaml .Values.server.extraLabels | nindent 4 -}}
+    {{- end -}}
   {{- template "vault.statefulSet.annotations" . }}
 spec:
   serviceName: {{ template "vault.fullname" . }}-internal
diff --git a/test/unit/injector-deployment.bats b/test/unit/injector-deployment.bats
index 740c17aae..14a071352 100755
--- a/test/unit/injector-deployment.bats
+++ b/test/unit/injector-deployment.bats
@@ -565,12 +565,19 @@ load _helpers
 
 @test "injector/deployment: specify extraLabels" {
   cd `chart_dir`
-  local actual=$(helm template \
+  local object=$(helm template \
       --show-only templates/injector-deployment.yaml \
       --set 'injector.extraLabels.foo=bar' \
       . | tee /dev/stderr |
+      yq -r '.' | tee /dev/stderr)
+
+  local value=$(echo $object |
       yq -r '.spec.template.metadata.labels.foo' | tee /dev/stderr)
-  [ "${actual}" = "bar" ]
+  [ "${value}" = "bar" ]
+
+  local value=$(echo $object |
+      yq -r '.metadata.labels.foo' | tee /dev/stderr)
+  [ "${value}" = "bar" ]
 }
 
 #--------------------------------------------------------------------
diff --git a/test/unit/injector-serviceaccount.bats b/test/unit/injector-serviceaccount.bats
index 1055d906c..fc84026eb 100755
--- a/test/unit/injector-serviceaccount.bats
+++ b/test/unit/injector-serviceaccount.bats
@@ -20,3 +20,16 @@ load _helpers
       yq 'length > 0' | tee /dev/stderr)
   [ "${actual}" = "false" ]
 }
+
+#--------------------------------------------------------------------
+# extra labels
+
+@test "injector/ServiceAccount: specify extraLabels" {
+  cd `chart_dir`
+  local actual=$(helm template \
+      --show-only templates/injector-serviceaccount.yaml \
+      --set 'injector.extraLabels.foo=bar' \
+      . | tee /dev/stderr |
+      yq -r '.metadata.labels.foo' | tee /dev/stderr)
+  [ "${actual}" = "bar" ]
+}
diff --git a/test/unit/server-serviceaccount.bats b/test/unit/server-serviceaccount.bats
index 29e18b56f..7610b7680 100755
--- a/test/unit/server-serviceaccount.bats
+++ b/test/unit/server-serviceaccount.bats
@@ -110,10 +110,23 @@ load _helpers
   [ "${actual}" = "false" ]
 
   local actual=$( (helm template \
-      --show-only templates/server-service.yaml  \
+      --show-only templates/server-serviceaccount.yaml  \
       --set 'server.standalone.enabled=true' \
       --set 'injector.externalVaultAddr=http://vault-outside' \
       . || echo "---") | tee /dev/stderr |
       yq 'length > 0' | tee /dev/stderr)
   [ "${actual}" = "false" ]
 }
+
+#--------------------------------------------------------------------
+# extra labels
+
+@test "server/ServiceAccount: specify extraLabels" {
+  cd `chart_dir`
+  local actual=$(helm template \
+      --show-only templates/server-serviceaccount.yaml \
+      --set 'server.extraLabels.foo=bar' \
+      . | tee /dev/stderr |
+      yq -r '.metadata.labels.foo' | tee /dev/stderr)
+  [ "${actual}" = "bar" ]
+}
diff --git a/test/unit/server-statefulset.bats b/test/unit/server-statefulset.bats
index 65d6d9688..3d9907b81 100755
--- a/test/unit/server-statefulset.bats
+++ b/test/unit/server-statefulset.bats
@@ -1024,12 +1024,19 @@ load _helpers
 
 @test "server/standalone-StatefulSet: specify extraLabels" {
   cd `chart_dir`
-  local actual=$(helm template \
+  local object=$(helm template \
       --show-only templates/server-statefulset.yaml \
       --set 'server.extraLabels.foo=bar' \
       . | tee /dev/stderr |
+      yq -r '.' | tee /dev/stderr)
+
+  local value=$(echo $object |
       yq -r '.spec.template.metadata.labels.foo' | tee /dev/stderr)
-  [ "${actual}" = "bar" ]
+  [ "${value}" = "bar" ]
+
+  local value=$(echo $object |
+      yq -r '.metadata.labels.foo' | tee /dev/stderr)
+  [ "${value}" = "bar" ]
 }
 
 # extra annotations