From 5631be6e79e208dceac8f314d16429e632f55442 Mon Sep 17 00:00:00 2001
From: Yong Wen Chua <lawliet89@users.noreply.github.com>
Date: Thu, 30 Apr 2020 12:23:22 +0800
Subject: [PATCH] Support setting priorityClassName on pods

---
 templates/injector-deployment.yaml |  3 +++
 templates/server-statefulset.yaml  |  3 +++
 test/unit/injector-deployment.bats | 22 ++++++++++++++++++++++
 test/unit/server-statefulset.bats  | 26 ++++++++++++++++++++++++--
 values.yaml                        |  6 ++++++
 5 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/templates/injector-deployment.yaml b/templates/injector-deployment.yaml
index 423372616..1c5b951b4 100644
--- a/templates/injector-deployment.yaml
+++ b/templates/injector-deployment.yaml
@@ -27,6 +27,9 @@ spec:
       {{ template "injector.affinity" . }}
       {{ template "injector.tolerations" . }}
       {{ template "injector.nodeselector" . }}
+      {{- if .Values.injector.priorityClassName }}
+      priorityClassName: {{ .Values.injector.priorityClassName }}
+      {{- end }}
       serviceAccountName: "{{ template "vault.fullname" . }}-agent-injector"
       securityContext:
         runAsNonRoot: true
diff --git a/templates/server-statefulset.yaml b/templates/server-statefulset.yaml
index 545b3d64f..3f4070961 100644
--- a/templates/server-statefulset.yaml
+++ b/templates/server-statefulset.yaml
@@ -37,6 +37,9 @@ spec:
       {{ template "vault.affinity" . }}
       {{ template "vault.tolerations" . }}
       {{ template "vault.nodeselector" . }}
+      {{- if .Values.server.priorityClassName }}
+      priorityClassName: {{ .Values.server.priorityClassName }}
+      {{- end }}
       terminationGracePeriodSeconds: 10
       serviceAccountName: {{ template "vault.fullname" . }}
       {{ if  .Values.server.shareProcessNamespace }}
diff --git a/test/unit/injector-deployment.bats b/test/unit/injector-deployment.bats
index 033ce7c4c..bd3f63a0b 100755
--- a/test/unit/injector-deployment.bats
+++ b/test/unit/injector-deployment.bats
@@ -425,3 +425,25 @@ load _helpers
       yq -r '.spec.template.spec.nodeSelector' | tee /dev/stderr)
   [ "${actual}" = "testing" ]
 }
+
+#--------------------------------------------------------------------
+# priorityClassName
+
+@test "injector/deployment: priorityClassName not set by default" {
+  cd `chart_dir`
+  local actual=$(helm template \
+      --show-only templates/injector-deployment.yaml  \
+      . | tee /dev/stderr |
+      yq '.spec.template.spec | .priorityClassName? == null' | tee /dev/stderr)
+  [ "${actual}" = "true" ]
+}
+
+@test "injector/deployment: priorityClassName can be set" {
+  cd `chart_dir`
+  local actual=$(helm template \
+      --show-only templates/injector-deployment.yaml  \
+      --set 'injector.priorityClassName=armaggeddon' \
+      . | tee /dev/stderr |
+      yq '.spec.template.spec | .priorityClassName == "armaggeddon"' | tee /dev/stderr)
+  [ "${actual}" = "true" ]
+}
diff --git a/test/unit/server-statefulset.bats b/test/unit/server-statefulset.bats
index 8e801190b..3fa7ba46a 100755
--- a/test/unit/server-statefulset.bats
+++ b/test/unit/server-statefulset.bats
@@ -711,7 +711,7 @@ load _helpers
       . | tee /dev/stderr |
       yq -r '.spec.template.spec.shareProcessNamespace' | tee /dev/stderr)
 
-  [ "${actual}" = "null" ]  
+  [ "${actual}" = "null" ]
 }
 
 @test "server/standalone-StatefulSet: shareProcessNamespace enabled" {
@@ -724,7 +724,7 @@ load _helpers
       . | tee /dev/stderr |
       yq -r '.spec.template.spec.shareProcessNamespace' | tee /dev/stderr)
 
-  [ "${actual}" = "true" ]  
+  [ "${actual}" = "true" ]
 }
 
 # extra labels
@@ -958,3 +958,25 @@ load _helpers
       yq -r '.spec.template.metadata.annotations["vaultIsAwesome"]' | tee /dev/stderr)
   [ "${actual}" = "true" ]
 }
+
+#--------------------------------------------------------------------
+# priorityClassName
+
+@test "server/standalone-StatefulSet: priorityClassName not set by default" {
+  cd `chart_dir`
+  local actual=$(helm template \
+      --show-only templates/server-statefulset.yaml  \
+      . | tee /dev/stderr |
+      yq '.spec.template.spec | .priorityClassName? == null' | tee /dev/stderr)
+  [ "${actual}" = "true" ]
+}
+
+@test "server/standalone-StatefulSet: priorityClassName can be set" {
+  cd `chart_dir`
+  local actual=$(helm template \
+      --show-only templates/server-statefulset.yaml  \
+      --set 'server.priorityClassName=armaggeddon' \
+      . | tee /dev/stderr |
+      yq '.spec.template.spec | .priorityClassName == "armaggeddon"' | tee /dev/stderr)
+  [ "${actual}" = "true" ]
+}
diff --git a/values.yaml b/values.yaml
index 305da7ba7..2385dccef 100644
--- a/values.yaml
+++ b/values.yaml
@@ -103,6 +103,9 @@ injector:
   #   beta.kubernetes.io/arch: amd64
   nodeSelector: null
 
+  # Priority class for injector pods
+  priorityClassName: ""
+
 server:
   # Resource requests, limits, etc. for the server cluster placement. This
   # should map directly to the value of the resources field for a PodSpec.
@@ -227,6 +230,9 @@ server:
   #   beta.kubernetes.io/arch: amd64
   nodeSelector: {}
 
+  # Priority class for server pods
+  priorityClassName: ""
+
   # Extra labels to attach to the server pods
   # This should be a multi-line string mapping directly to the a map of
   # the labels to apply to the server pods