From b6ad60f900ef814f7c5cdb9e06df3d923b312d0a Mon Sep 17 00:00:00 2001
From: nikchern <nikchern@gmail.com>
Date: Tue, 12 Sep 2023 21:42:21 -0700
Subject: [PATCH] PLT-663: Fix machine pool handler. (#342) (#344)

---
 spectrocloud/cluster_node_common.go      | 29 +++++++++++++--
 spectrocloud/cluster_node_common_test.go | 47 ++++++++++++++++++++++++
 2 files changed, 73 insertions(+), 3 deletions(-)
 create mode 100644 spectrocloud/cluster_node_common_test.go

diff --git a/spectrocloud/cluster_node_common.go b/spectrocloud/cluster_node_common.go
index e0ce2a3b..68fad29f 100644
--- a/spectrocloud/cluster_node_common.go
+++ b/spectrocloud/cluster_node_common.go
@@ -2,12 +2,14 @@ package spectrocloud
 
 import (
 	"context"
+	"fmt"
+	"log"
+	"time"
+
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/retry"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
 	"github.com/spectrocloud/hapi/models"
 	"github.com/spectrocloud/palette-sdk-go/client"
-	"log"
-	"time"
 )
 
 var NodeMaintenanceLifecycleStates = []string{
@@ -84,7 +86,12 @@ func resourceNodeAction(c *client.V1Client, ctx context.Context, newMachinePool
 func flattenNodeMaintenanceStatus(c *client.V1Client, d *schema.ResourceData, fn GetNodeStatusMap, mPools []interface{}, cloudConfigId string, ClusterContext string) ([]interface{}, error) {
 	_, n := d.GetChange("machine_pool")
 	nsMap := make(map[string]interface{})
-	for _, mp := range n.(*schema.Set).List() {
+	machinePoolsList, i, err := getMachinePoolList(n)
+	if err != nil {
+		return i, err
+	}
+
+	for _, mp := range machinePoolsList {
 		machinePool := mp.(map[string]interface{})
 		nsMap[machinePool["name"].(string)] = machinePool
 	}
@@ -117,3 +124,19 @@ func flattenNodeMaintenanceStatus(c *client.V1Client, d *schema.ResourceData, fn
 	}
 	return mPools, nil
 }
+
+func getMachinePoolList(n interface{}) ([]interface{}, []interface{}, error) {
+	var machinePoolsList []interface{}
+
+	// Check if n is of type *schema.Set
+	if set, ok := n.(*schema.Set); ok {
+		machinePoolsList = set.List()
+	} else if list, ok := n.([]interface{}); ok {
+		// If n is already a slice of interfaces
+		machinePoolsList = list
+	} else {
+		// Handle error: n is neither *schema.Set nor []interface{}
+		return nil, nil, fmt.Errorf("unexpected type for n: %T", n)
+	}
+	return machinePoolsList, nil, nil
+}
diff --git a/spectrocloud/cluster_node_common_test.go b/spectrocloud/cluster_node_common_test.go
new file mode 100644
index 00000000..a99597ec
--- /dev/null
+++ b/spectrocloud/cluster_node_common_test.go
@@ -0,0 +1,47 @@
+package spectrocloud
+
+import (
+	"testing"
+
+	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGetMachinePoolList(t *testing.T) {
+	tests := []struct {
+		name    string
+		input   interface{}
+		want    []interface{}
+		wantErr bool
+	}{
+		{
+			name:    "Handle *schema.Set",
+			input:   schema.NewSet(schema.HashString, []interface{}{"a", "b"}),
+			want:    []interface{}{"a", "b"},
+			wantErr: false,
+		},
+		{
+			name:    "Handle []interface{}",
+			input:   []interface{}{"a", "b"},
+			want:    []interface{}{"a", "b"},
+			wantErr: false,
+		},
+		{
+			name:    "Handle unexpected type",
+			input:   "unexpected",
+			want:    nil,
+			wantErr: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, _, err := getMachinePoolList(tt.input)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("getMachinePoolList() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			assert.ElementsMatch(t, tt.want, got)
+		})
+	}
+}