Add new node type selector to databricks_node_type data source (#1957)

* add new node type selector to `databricks_node_type` data source

* address code review comments

Author: alexott

clusters/data_node_type.go

@@ -24,6 +24,7 @@ type NodeTypeRequest struct {
 	IsIOCacheEnabled      bool   `json:"is_io_cache_enabled,omitempty"`
 	SupportPortForwarding bool   `json:"support_port_forwarding,omitempty"`
 	VCPU                  bool   `json:"vcpu,omitempty"`
+	Fleet                 bool   `json:"fleet,omitempty"`
 }
 
 // NodeTypeList contains a list of node types
@@ -133,9 +134,12 @@ func (a ClustersAPI) GetSmallestNodeType(r NodeTypeRequest) string {
 	// error is explicitly ingored here, because Azure returns
 	// apparently too big of a JSON for Go to parse
 	if len(list.NodeTypes) == 0 {
-		if r.VCPU {
+		if r.VCPU && a.client.IsAws() {
 			return "vcpu-worker"
 		}
+		if r.Fleet && a.client.IsAws() {
+			return "md-fleet.xlarge"
+		}
 		return a.defaultSmallestNodeType()
 	}
 	list.Sort()
@@ -144,10 +148,10 @@ func (a ClustersAPI) GetSmallestNodeType(r NodeTypeRequest) string {
 			continue
 		}
 		gbs := (nt.MemoryMB / 1024)
-		if r.VCPU && !strings.HasPrefix(nt.NodeTypeID, "vcpu") {
+		if r.VCPU != strings.HasPrefix(nt.NodeTypeID, "vcpu") {
 			continue
 		}
-		if !r.VCPU && strings.HasPrefix(nt.NodeTypeID, "vcpu") {
+		if r.Fleet != strings.Contains(nt.NodeTypeID, "-fleet.") {
 			continue
 		}
 		if r.MinMemoryGB > 0 && gbs < r.MinMemoryGB {

clusters/data_node_type_test.go

@@ -23,6 +23,12 @@ func TestNodeType(t *testing.T) {
 							MemoryMB:       0,
 							NumCores:       0,
 						},
+						{
+							NodeTypeID:     "m-fleet.xlarge",
+							InstanceTypeID: "m-fleet.xlarge",
+							MemoryMB:       16384,
+							NumCores:       4,
+						},
 						{
 							NodeTypeID:     "Random_05",
 							InstanceTypeID: "Random_05",
@@ -183,7 +189,7 @@ func TestNodeTypeVCPU(t *testing.T) {
 							NodeTypeID:     "Random_05",
 							InstanceTypeID: "Random_05",
 							MemoryMB:       1024,
-							NumCores:       32,
+							NumCores:       2,
 							NodeInstanceType: &NodeInstanceType{
 								LocalDisks:      3,
 								LocalDiskSizeGB: 100,
@@ -193,7 +199,7 @@ func TestNodeTypeVCPU(t *testing.T) {
 							NodeTypeID:     "vcpu-worker",
 							InstanceTypeID: "vcpu-worker",
 							MemoryMB:       0,
-							NumCores:       0,
+							NumCores:       4,
 						},
 					},
 				},
@@ -223,6 +229,12 @@ func TestSmallestNodeTypeClouds(t *testing.T) {
 			Host: "foo.gcp.databricks.com",
 		},
 	}.defaultSmallestNodeType())
+
+	assert.Equal(t, "i3.xlarge", ClustersAPI{
+		client: &common.DatabricksClient{
+			Host: "foo.cloud.databricks.com",
+		},
+	}.defaultSmallestNodeType())
 }
 
 func TestNodeTypeCategoryNotAvailable(t *testing.T) {
@@ -304,3 +316,112 @@ func TestNodeTypeShouldBeSkipped(t *testing.T) {
 	assert.Equal(t, true, toBeSkipped.shouldBeSkipped())
 	assert.Equal(t, false, NodeType{}.shouldBeSkipped())
 }
+
+func TestNodeTypeFleet(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:       4,
+						},
+						{
+							NodeTypeID:     "m-fleet.xlarge",
+							InstanceTypeID: "m-fleet.xlarge",
+							MemoryMB:       16384,
+							NumCores:       4,
+						},
+						{
+							NodeTypeID:     "m-fleet.2xlarge",
+							InstanceTypeID: "m-fleet.2xlarge",
+							MemoryMB:       32768,
+							NumCores:       8,
+						},
+					},
+				},
+			},
+		},
+		Read:        true,
+		Resource:    DataSourceNodeType(),
+		NonWritable: true,
+		State: map[string]any{
+			"fleet":     true,
+			"min_cores": 8,
+		},
+		ID: ".",
+	}.Apply(t)
+	assert.NoError(t, err)
+	assert.Equal(t, "m-fleet.2xlarge", d.Id())
+}
+
+func TestNodeTypeEmptyList(t *testing.T) {
+	d, err := qa.ResourceFixture{
+		Fixtures: []qa.HTTPFixture{
+			{
+				Method:       "GET",
+				ReuseRequest: true,
+				Resource:     "/api/2.0/clusters/list-node-types",
+				Response:     NodeTypeList{},
+			},
+		},
+		Read:        true,
+		Resource:    DataSourceNodeType(),
+		NonWritable: true,
+		Azure:       true,
+		State:       map[string]any{},
+		ID:          ".",
+	}.Apply(t)
+	assert.NoError(t, err)
+	assert.Equal(t, "Standard_D3_v2", d.Id())
+}
+
+func TestNodeTypeFleetEmptyList(t *testing.T) {
+	d, err := qa.ResourceFixture{
+		Fixtures: []qa.HTTPFixture{
+			{
+				Method:       "GET",
+				ReuseRequest: true,
+				Resource:     "/api/2.0/clusters/list-node-types",
+				Response:     NodeTypeList{},
+			},
+		},
+		Read:        true,
+		Resource:    DataSourceNodeType(),
+		NonWritable: true,
+		State: map[string]any{
+			"fleet": true,
+		},
+		ID: ".",
+	}.Apply(t)
+	assert.NoError(t, err)
+	assert.Equal(t, "md-fleet.xlarge", d.Id())
+}
+
+func TestNodeTypeVCPUEmptyList(t *testing.T) {
+	d, err := qa.ResourceFixture{
+		Fixtures: []qa.HTTPFixture{
+			{
+				Method:       "GET",
+				ReuseRequest: true,
+				Resource:     "/api/2.0/clusters/list-node-types",
+				Response:     NodeTypeList{},
+			},
+		},
+		Read:        true,
+		Resource:    DataSourceNodeType(),
+		NonWritable: true,
+		State: map[string]any{
+			"vcpu": true,
+		},
+		ID: ".",
+	}.Apply(t)
+	assert.NoError(t, err)
+	assert.Equal(t, "vcpu-worker", d.Id())
+}