From 38498b1981efd130d4224c98b40a2b2dfa77d060 Mon Sep 17 00:00:00 2001
From: katexochen <49727155+katexochen@users.noreply.github.com>
Date: Thu, 6 Oct 2022 12:14:56 +0200
Subject: [PATCH] Update operator role identification

---
 .../internal/azure/client/scalinggroup.go          | 14 ++++++++++++--
 .../internal/azure/client/scalinggroup_test.go     |  4 ++--
 2 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/operators/constellation-node-operator/internal/azure/client/scalinggroup.go b/operators/constellation-node-operator/internal/azure/client/scalinggroup.go
index 286885b0e..7ece313a0 100644
--- a/operators/constellation-node-operator/internal/azure/client/scalinggroup.go
+++ b/operators/constellation-node-operator/internal/azure/client/scalinggroup.go
@@ -89,15 +89,25 @@ func (c *Client) ListScalingGroups(ctx context.Context, uid string) (controlPlan
 		if err != nil {
 			return nil, nil, fmt.Errorf("getting scaling group name: %w", err)
 		}
-		if scaleSet == "constellation-scale-set-controlplanes-"+uid {
+		if isControlPlaneInstanceGroup(scaleSet) {
 			controlPlaneGroupIDs = append(controlPlaneGroupIDs, scaleSetID)
-		} else if strings.HasPrefix(scaleSet, "constellation-scale-set-workers-"+uid) {
+		} else if isWorkerInstanceGroup(scaleSet) {
 			workerGroupIDs = append(workerGroupIDs, scaleSetID)
 		}
 	}
 	return controlPlaneGroupIDs, workerGroupIDs, nil
 }
 
+// isControlPlaneInstanceGroup returns true if the instance group is a control plane instance group.
+func isControlPlaneInstanceGroup(instanceGroupName string) bool {
+	return strings.Contains(instanceGroupName, "control-plane")
+}
+
+// isWorkerInstanceGroup returns true if the instance group is a worker instance group.
+func isWorkerInstanceGroup(instanceGroupName string) bool {
+	return strings.Contains(instanceGroupName, "worker")
+}
+
 func imageReferenceFromImage(img string) *armcompute.ImageReference {
 	ref := &armcompute.ImageReference{}
 
diff --git a/operators/constellation-node-operator/internal/azure/client/scalinggroup_test.go b/operators/constellation-node-operator/internal/azure/client/scalinggroup_test.go
index fb2985691..0eda6dff4 100644
--- a/operators/constellation-node-operator/internal/azure/client/scalinggroup_test.go
+++ b/operators/constellation-node-operator/internal/azure/client/scalinggroup_test.go
@@ -192,9 +192,9 @@ func TestListScalingGroups(t *testing.T) {
 	}{
 		"listing control-plane works": {
 			scaleSet: armcomputev2.VirtualMachineScaleSet{
-				ID: to.Ptr("/subscriptions/subscription-id/resourceGroups/resource-group/providers/Microsoft.Compute/virtualMachineScaleSets/constellation-scale-set-controlplanes-uid"),
+				ID: to.Ptr("/subscriptions/subscription-id/resourceGroups/resource-group/providers/Microsoft.Compute/virtualMachineScaleSets/constellation-scale-set-control-planes-uid"),
 			},
-			wantControlPlanes: []string{"/subscriptions/subscription-id/resourceGroups/resource-group/providers/Microsoft.Compute/virtualMachineScaleSets/constellation-scale-set-controlplanes-uid"},
+			wantControlPlanes: []string{"/subscriptions/subscription-id/resourceGroups/resource-group/providers/Microsoft.Compute/virtualMachineScaleSets/constellation-scale-set-control-planes-uid"},
 		},
 		"listing worker works": {
 			scaleSet: armcomputev2.VirtualMachineScaleSet{