Fixed databricks_node_type by skipping nodes that aren't available in subscription/region (#1534)

Author: alexott

clusters/data_node_type.go

@@ -50,6 +50,11 @@ func (l *NodeTypeList) Sort() {
 	})
 }
 
+const (
+	CloudProviderNodeStatusNotEnabled           = "NotEnabledOnSubscription"
+	CloudProviderNodeStatusNotAvailableInRegion = "NotAvailableInRegion"
+)
+
 // ClusterCloudProviderNodeInfo encapsulates the existing quota available from the cloud service provider.
 type ClusterCloudProviderNodeInfo struct {
 	Status             []string `json:"status,omitempty"`
@@ -89,6 +94,19 @@ type NodeType struct {
 	Graviton              bool                          `json:"is_graviton,omitempty"`
 }
 
+func (nt NodeType) shouldBeSkipped() bool {
+	if nt.NodeInfo == nil {
+		return false
+	}
+	for _, st := range nt.NodeInfo.Status {
+		switch st {
+		case CloudProviderNodeStatusNotAvailableInRegion, CloudProviderNodeStatusNotEnabled:
+			return true
+		}
+	}
+	return false
+}
+
 func (a ClustersAPI) defaultSmallestNodeType() string {
 	if a.client.IsAzure() {
 		return "Standard_D3_v2"
@@ -117,6 +135,9 @@ func (a ClustersAPI) GetSmallestNodeType(r NodeTypeRequest) string {
 	}
 	list.Sort()
 	for _, nt := range list.NodeTypes {
+		if nt.shouldBeSkipped() {
+			continue
+		}
 		gbs := (nt.MemoryMB / 1024)
 		if r.VCPU && !strings.HasPrefix(nt.NodeTypeID, "vcpu") {
 			continue

clusters/data_node_type_test.go

@@ -224,3 +224,83 @@ func TestSmallestNodeTypeClouds(t *testing.T) {
 		},
 	}.defaultSmallestNodeType())
 }
+
+func TestNodeTypeCategoryNotAvailable(t *testing.T) {
+	d, err := qa.ResourceFixture{
+		Fixtures: []qa.HTTPFixture{
+			{
+				Method:       "GET",
+				ReuseRequest: true,
+				Resource:     "/api/2.0/clusters/list-node-types",
+				Response: NodeTypeList{
+					[]NodeType{
+						{
+							NodeTypeID:     "Random_05",
+							InstanceTypeID: "Random_05",
+							MemoryMB:       1024,
+							NumCores:       32,
+							NodeInstanceType: &NodeInstanceType{
+								LocalDisks:      3,
+								LocalDiskSizeGB: 100,
+							},
+						},
+						{
+							NodeTypeID:     "Random_01",
+							InstanceTypeID: "Random_01",
+							MemoryMB:       8192,
+							NumCores:       8,
+							NodeInstanceType: &NodeInstanceType{
+								InstanceTypeID: "_",
+							},
+							Category: "Memory Optimized",
+						},
+						{
+							NodeTypeID:     "Random_02",
+							InstanceTypeID: "Random_02",
+							MemoryMB:       8192,
+							NumCores:       8,
+							NumGPUs:        2,
+							Category:       "Storage Optimized",
+							NodeInfo: &ClusterCloudProviderNodeInfo{
+								Status: []string{CloudProviderNodeStatusNotAvailableInRegion, CloudProviderNodeStatusNotEnabled},
+							},
+						},
+						{
+							NodeTypeID:     "Random_03",
+							InstanceTypeID: "Random_03",
+							MemoryMB:       8192,
+							NumCores:       8,
+							NumGPUs:        2,
+							Category:       "Storage Optimized",
+						},
+					},
+				},
+			},
+		},
+		Read:        true,
+		Resource:    DataSourceNodeType(),
+		NonWritable: true,
+		State: map[string]any{
+			"category": "Storage optimized",
+		},
+		ID: ".",
+	}.Apply(t)
+	assert.NoError(t, err)
+	assert.Equal(t, "Random_03", d.Id())
+}
+
+func TestNodeTypeShouldBeSkipped(t *testing.T) {
+	toBeSkipped := NodeType{
+		NodeTypeID:     "Random_02",
+		InstanceTypeID: "Random_02",
+		MemoryMB:       8192,
+		NumCores:       8,
+		NumGPUs:        2,
+		Category:       "Storage Optimized",
+		NodeInfo: &ClusterCloudProviderNodeInfo{
+			Status: []string{CloudProviderNodeStatusNotAvailableInRegion, CloudProviderNodeStatusNotEnabled},
+		},
+	}
+	assert.Equal(t, true, toBeSkipped.shouldBeSkipped())
+	assert.Equal(t, false, NodeType{}.shouldBeSkipped())
+}