add template function

Signed-off-by: Mahmoud Gaballah <[email protected]>
zalando-incubator · Aug 17, 2023 · 4e22ef2 · 4e22ef2
1 parent 3e4bf00
commit 4e22ef2
Show file tree

Hide file tree

Showing 2 changed files with 161 additions and 16 deletions.
diff --git a/provisioner/template.go b/provisioner/template.go
@@ -131,18 +131,19 @@ func renderTemplate(context *templateContext, file string) (string, error) {
 		"nodeCIDRMaxNodes": func(maskSize int64, reserved int64) (int64, error) {
 			return nodeCIDRMaxNodes(16, maskSize, reserved)
 		},
-		"nodeCIDRMaxNodesPodCIDR":       nodeCIDRMaxNodes,
-		"nodeCIDRMaxPods":               nodeCIDRMaxPods,
-		"parseInt64":                    parseInt64,
-		"generateJWKSDocument":          generateJWKSDocument,
-		"generateOIDCDiscoveryDocument": generateOIDCDiscoveryDocument,
-		"kubernetesSizeToKiloBytes":     kubernetesSizeToKiloBytes,
-		"indexedList":                   indexedList,
-		"zoneDistributedNodePoolGroups": zoneDistributedNodePoolGroups,
-		"certificateExpiry":             certificateExpiry,
-		"sumQuantities":                 sumQuantities,
-		"awsValidID":                    awsValidID,
-		"indent":                        sprig.GenericFuncMap()["indent"],
+		"nodeCIDRMaxNodesPodCIDR":               nodeCIDRMaxNodes,
+		"nodeCIDRMaxPods":                       nodeCIDRMaxPods,
+		"parseInt64":                            parseInt64,
+		"generateJWKSDocument":                  generateJWKSDocument,
+		"generateOIDCDiscoveryDocument":         generateOIDCDiscoveryDocument,
+		"kubernetesSizeToKiloBytes":             kubernetesSizeToKiloBytes,
+		"indexedList":                           indexedList,
+		"zoneDistributedNodePoolGroups":         zoneDistributedNodePoolGroups,
+		"nodeLifeCycleProviderPerNodePoolGroup": nodeLifeCycleProviderPerNodePoolGroup,
+		"certificateExpiry":                     certificateExpiry,
+		"sumQuantities":                         sumQuantities,
+		"awsValidID":                            awsValidID,
+		"indent":                                sprig.GenericFuncMap()["indent"],
 	}
 
 	content, ok := context.fileData[file]
@@ -671,6 +672,17 @@ func poolsDistributed(dedicated string, pools []*api.NodePool) bool {
 //
 // The default pool is represented with an empty string as the key.
 func zoneDistributedNodePoolGroups(nodePools []*api.NodePool) map[string]bool {
+	poolGroups := groupNodePoolsByPurpose(nodePools)
+	result := make(map[string]bool)
+	for group, pools := range poolGroups {
+		if poolsDistributed(group, pools) {
+			result[group] = true
+		}
+	}
+	return result
+}
+
+func groupNodePoolsByPurpose(nodePools []*api.NodePool) map[string][]*api.NodePool {
 	poolGroups := make(map[string][]*api.NodePool)
 
 	for _, pool := range nodePools {
@@ -686,13 +698,34 @@ func zoneDistributedNodePoolGroups(nodePools []*api.NodePool) map[string]bool {
 		}
 	}
 
-	result := make(map[string]bool)
+	return poolGroups
+}
+
+// nodeLifeCycleProviderPerNodePoolGroup groups node-pools by the dedicated label, each node-pool-group is mapped to
+// the provider suitable for its node-pool profile
+// in case all pools of a group do not share the same profile, the group is mapped to empty string
+func nodeLifeCycleProviderPerNodePoolGroup(nodePools []*api.NodePool) map[string]string {
+	providers := map[string]string{
+		"worker-karpenter": "karpenter",
+		"worker-combined":  "zalando",
+		"worker-splitaz":   "zalando",
+	}
+	poolGroups := groupNodePoolsByPurpose(nodePools)
+	provider := make(map[string]string)
 	for group, pools := range poolGroups {
-		if poolsDistributed(group, pools) {
-			result[group] = true
+		if group == "" {
+			group = "default"
+		}
+		groupProfile := pools[0].Profile
+		for _, pool := range pools[1:] {
+			if pool.Profile != groupProfile {
+				groupProfile = ""
+				break
+			}
 		}
+		provider[group] = providers[groupProfile]
 	}
-	return result
+	return provider
 }
 
 // certificateExpiry returns the notAfter timestamp of a PEM-encoded certificate in the RFC3339 format

diff --git a/provisioner/template_test.go b/provisioner/template_test.go
@@ -1052,3 +1052,115 @@ func TestAWSValidID(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, "aws__12345678910__eu-central-1__kube-1", result)
 }
+
+func TestNodePoolGroupsProfile(t *testing.T) {
+	for _, tc := range []struct {
+		name     string
+		input    []*api.NodePool
+		expected map[string]string
+	}{
+		{
+			name: "application-1 dedicated pools share the same profile",
+			input: []*api.NodePool{
+				{
+					Name:    "example-1",
+					Profile: "worker-combined",
+					ConfigItems: map[string]string{
+						"labels": "dedicated=application-1",
+					},
+				},
+				{
+					Name:    "example-2",
+					Profile: "worker-combined",
+					ConfigItems: map[string]string{
+						"labels": "dedicated=application-1",
+					},
+				},
+			},
+			expected: map[string]string{
+				"application-1": "zalando",
+			},
+		},
+		{
+			name: "application-1 dedicated pools share the same profile, which is un-known",
+			input: []*api.NodePool{
+				{
+					Name:    "example-1",
+					Profile: "worker-unknown",
+					ConfigItems: map[string]string{
+						"labels": "dedicated=application-1",
+					},
+				},
+				{
+					Name:    "example-2",
+					Profile: "worker-unknown",
+					ConfigItems: map[string]string{
+						"labels": "dedicated=application-1",
+					},
+				},
+			},
+			expected: map[string]string{
+				"application-1": "",
+			},
+		},
+		{
+			name: "application-2 dedicated pools do not share the same profile",
+			input: []*api.NodePool{
+				{
+					Name:    "example-1",
+					Profile: "profile-2",
+					ConfigItems: map[string]string{
+						"labels": "dedicated=application-2",
+					},
+				},
+				{
+					Name:    "example-2",
+					Profile: "profile-1",
+					ConfigItems: map[string]string{
+						"labels": "dedicated=application-2",
+					},
+				},
+			},
+			expected: map[string]string{
+				"application-2": "",
+			},
+		},
+		{
+			name: "default pools share the same profile",
+			input: []*api.NodePool{
+				{
+					Name:    "example-1",
+					Profile: "worker-karpenter",
+				},
+				{
+					Name:    "example-2",
+					Profile: "worker-karpenter",
+				},
+			},
+			expected: map[string]string{
+				"default": "karpenter",
+			},
+		},
+		{
+			name: "default pools do not share the same profile",
+			input: []*api.NodePool{
+				{
+					Name:    "example-1",
+					Profile: "worker-karpenter",
+				},
+				{
+					Name:    "example-2",
+					Profile: "worker-combined",
+				},
+			},
+			expected: map[string]string{
+				"default": "",
+			},
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			output := nodeLifeCycleProviderPerNodePoolGroup(tc.input)
+			require.Equal(t, tc.expected, output)
+		})
+	}
+}