From 317d658ccc7a18118f98482fd791a26609592e05 Mon Sep 17 00:00:00 2001
From: Bastian Hofmann <mail@bastianhofmann.de>
Date: Thu, 31 Aug 2023 13:45:56 +0200
Subject: [PATCH] Make pod topologySpreadConstraints configurable

Fixes https://github.com/qdrant/qdrant-helm/issues/65
More info https://kubernetes.io/docs/concepts/scheduling-eviction/topology-spread-constraints/
---
 charts/qdrant/templates/statefulset.yaml      |  4 ++
 charts/qdrant/values.yaml                     |  2 +
 ...qdrant_topology-spread-constraints_test.go | 40 +++++++++++++++++++
 3 files changed, 46 insertions(+)
 create mode 100644 test/qdrant_topology-spread-constraints_test.go

diff --git a/charts/qdrant/templates/statefulset.yaml b/charts/qdrant/templates/statefulset.yaml
index 78895ce..1f31edb 100644
--- a/charts/qdrant/templates/statefulset.yaml
+++ b/charts/qdrant/templates/statefulset.yaml
@@ -170,6 +170,10 @@ spec:
       tolerations:
         {{- toYaml . | nindent 8 }}
       {{- end }}
+      {{- with .Values.topologySpreadConstraints}}
+      topologySpreadConstraints:
+        {{- toYaml . | nindent 8 }}
+      {{- end }}
       serviceAccountName: {{ include "qdrant.fullname" . }}
       volumes:
         - name: qdrant-config
diff --git a/charts/qdrant/values.yaml b/charts/qdrant/values.yaml
index da5f353..4cd5e89 100644
--- a/charts/qdrant/values.yaml
+++ b/charts/qdrant/values.yaml
@@ -110,6 +110,8 @@ tolerations: []
 
 affinity: {}
 
+topologySpreadConstraints: []
+
 persistence:
   accessModes: ["ReadWriteOnce"]
   size: 10Gi
diff --git a/test/qdrant_topology-spread-constraints_test.go b/test/qdrant_topology-spread-constraints_test.go
new file mode 100644
index 0000000..c5c3a12
--- /dev/null
+++ b/test/qdrant_topology-spread-constraints_test.go
@@ -0,0 +1,40 @@
+package test
+
+import (
+	"path/filepath"
+	"strings"
+	"testing"
+
+	"github.com/gruntwork-io/terratest/modules/helm"
+	"github.com/gruntwork-io/terratest/modules/k8s"
+	"github.com/gruntwork-io/terratest/modules/logger"
+	"github.com/gruntwork-io/terratest/modules/random"
+	"github.com/stretchr/testify/require"
+	appsv1 "k8s.io/api/apps/v1"
+)
+
+func TestTopologySpreadConstraints(t *testing.T) {
+	t.Parallel()
+
+	helmChartPath, err := filepath.Abs("../charts/qdrant")
+	releaseName := "qdrant"
+	require.NoError(t, err)
+
+	namespaceName := "qdrant-" + strings.ToLower(random.UniqueId())
+	logger.Log(t, "Namespace: %s\n", namespaceName)
+
+	options := &helm.Options{
+		SetJsonValues: map[string]string{
+			"topologySpreadConstraints": `[{"maxSkew":1,"topologyKey":"kubernetes.io/hostname","whenUnsatisfiable":"DoNotSchedule","labelSelector":{"matchLabels":{"app":"qdrant"}}}]`,
+		},
+		KubectlOptions: k8s.NewKubectlOptions("", "", namespaceName),
+	}
+
+	output := helm.RenderTemplate(t, options, helmChartPath, releaseName, []string{"templates/statefulset.yaml"})
+
+	var statefulSet appsv1.StatefulSet
+	helm.UnmarshalK8SYaml(t, output, &statefulSet)
+
+	require.Equal(t, 1, len(statefulSet.Spec.Template.Spec.TopologySpreadConstraints))
+	require.Equal(t, int32(1), statefulSet.Spec.Template.Spec.TopologySpreadConstraints[0].MaxSkew)
+}