Fix databricks_mws_network_connectivity_config

databricks · Nov 14, 2024 · b19259f · b19259f
1 parent 3f1e078
commit b19259f
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 30 deletions.
diff --git a/internal/providers/sdkv2/sdkv2.go b/internal/providers/sdkv2/sdkv2.go
@@ -103,7 +103,7 @@ func DatabricksProvider(sdkV2Fallbacks ...pluginfw.SdkV2FallbackOption) *schema.
 		"databricks_mlflow_model":                         mlflow.DataSourceModel().ToResource(),
 		"databricks_mlflow_models":                        mlflow.DataSourceModels().ToResource(),
 		"databricks_mws_credentials":                      mws.DataSourceMwsCredentials().ToResource(),
-    "databricks_mws_network_connectivity_config":      mws.DataSourceMwsNetworkConnectivityConfig().ToResource(),
+		"databricks_mws_network_connectivity_config":      mws.DataSourceMwsNetworkConnectivityConfig().ToResource(),
 		"databricks_mws_workspaces":                       mws.DataSourceMwsWorkspaces().ToResource(),
 		"databricks_node_type":                            clusters.DataSourceNodeType().ToResource(),
 		"databricks_notebook":                             workspace.DataSourceNotebook().ToResource(),

diff --git a/mws/data_mws_network_connectivity_config.go b/mws/data_mws_network_connectivity_config.go
@@ -2,33 +2,37 @@ package mws
 
 import (
 	"context"
+	"fmt"
+
 	"github.com/databricks/databricks-sdk-go"
 	"github.com/databricks/databricks-sdk-go/service/settings"
 	"github.com/databricks/terraform-provider-databricks/common"
 )
 
 func DataSourceMwsNetworkConnectivityConfig() common.Resource {
-	type MwsNetworkConnectivity struct {
+	type mwsNetworkConnectivityConfiguration struct {
 		settings.NetworkConnectivityConfiguration
 	}
 
-	type MwsNetworkConnectivityParams struct {
-		Name   string `json:"name" tf:"computed,optional"`
-		Region string `json:"region" tf:"computed,optional"`
+	type mwsNetworkConnectivityConfigurationParams struct {
+		Name string `json:"name" tf:"computed,optional"`
 	}
 
-	return common.AccountDataWithParams(func(ctx context.Context, data MwsNetworkConnectivityParams, a *databricks.AccountClient) (*MwsNetworkConnectivity, error) {
-		nccList := a.NetworkConnectivity.ListNetworkConnectivityConfigurations(ctx, settings.ListNetworkConnectivityConfigurationsRequest{})
-		for nccList.HasNext(ctx) {
-			ncc, err := nccList.Next(ctx)
+	return common.AccountDataWithParams(func(ctx context.Context, data mwsNetworkConnectivityConfigurationParams, a *databricks.AccountClient) (*mwsNetworkConnectivityConfiguration, error) {
+		if data.Name == "" {
+			return nil, fmt.Errorf("`name` should be provided")
+		}
 
-			if err != nil {
-				return nil, err
-			}
-			if data.Region == ncc.Region || data.Name == ncc.Name {
-				return &MwsNetworkConnectivity{NetworkConnectivityConfiguration: ncc}, nil
-			}
+		list, err := a.NetworkConnectivity.ListNetworkConnectivityConfigurationsAll(ctx, settings.ListNetworkConnectivityConfigurationsRequest{})
+		if err != nil {
+			return nil, err
 		}
+
+        for _, ncc := range list {
+            if data.Name == ncc.Name {
+				return &mwsNetworkConnectivityConfiguration{NetworkConnectivityConfiguration: ncc}, nil
+			}
+        }
 		return nil, nil
 	})
 }
diff --git a/mws/data_mws_network_connectivity_config_test.go b/mws/data_mws_network_connectivity_config_test.go
@@ -4,14 +4,14 @@ import (
 	"fmt"
 	"testing"
 
+	"github.com/databricks/databricks-sdk-go/experimental/mocks"
 	"github.com/databricks/databricks-sdk-go/service/settings"
+	"github.com/stretchr/testify/mock"
 
 	"github.com/databricks/terraform-provider-databricks/qa"
 )
 
 func TestDataSourceMwsNetworkConnectivityConfig(t *testing.T) {
-	id := "xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx"
-
 	var ncc = settings.NetworkConnectivityConfiguration{
 		AccountId:    "abc",
 		CreationTime: 0,
@@ -20,34 +20,33 @@ func TestDataSourceMwsNetworkConnectivityConfig(t *testing.T) {
 			TargetRules:  nil,
 		},
 		Name:                        "def",
-		NetworkConnectivityConfigId: "",
+		NetworkConnectivityConfigId: "xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx",
 		Region:                      "us-east-1",
 		UpdatedTime:                 0,
 	}
 
 	qa.ResourceFixture{
-		Fixtures: []qa.HTTPFixture{
-			{
-				Method:   "GET",
-				Resource: fmt.Sprintf("/api/2.0/accounts/%s/network-connectivity-configs/%s?", ncc.AccountId, id),
-				Response: ncc,
-			},
-		},
+		MockAccountClientFunc: func(a *mocks.MockAccountClient) {
+            api := a.GetMockNetworkConnectivityAPI().EXPECT()
+			api.ListNetworkConnectivityConfigurationsAll(mock.Anything, settings.ListNetworkConnectivityConfigurationsRequest{}).Return(
+                []settings.NetworkConnectivityConfiguration{ncc}, nil,
+            )
+        },
 		AccountID:   "abc",
 		Read:        true,
 		NonWritable: true,
 		Resource:    DataSourceMwsNetworkConnectivityConfig(),
-		ID:          id,
+		ID:          "_",
 		HCL: fmt.Sprintf(`
-		network_connectivity_config_id = "%s"
-		`, id),
+		    name = "%s"
+		`, ncc.Name),
 	}.ApplyAndExpectData(t, map[string]interface{}{
 		"account_id":                     "abc",
 		"creation_time":                  0,
 		"egress_config":                  []interface{}{map[string]interface{}{"default_rules": []interface{}{}, "target_rules": []interface{}{}}},
-		"id":                             id,
+		"id":                             "_",
 		"name":                           "def",
-		"network_connectivity_config_id": id,
+		"network_connectivity_config_id": "xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx",
 		"region":                         "us-east-1",
 		"updated_time":                   0,
 	})
@@ -60,9 +59,23 @@ func TestDataSourceMwsNetworkConnectivityConfig_AccountID(t *testing.T) {
 		Read:        true,
 		NonWritable: true,
 		ID:          "_",
+        HCL: `
+        name = "def"
+        `,
 	}.ExpectError(t, "invalid Databricks Account configuration")
 }
 
+func TestDataSourceMwsNetworkConnectivityConfig_NoName(t *testing.T) {
+	qa.ResourceFixture{
+		Fixtures:    qa.HTTPFailures,
+		AccountID:   "abc",
+		Resource:    DataSourceMwsNetworkConnectivityConfig(),
+		Read:        true,
+		NonWritable: true,
+		ID:          "_",
+	}.ExpectError(t, "`name` should be provided")
+}
+
 func TestDataSourceMwsNetworkConnectivityConfig_Error(t *testing.T) {
 	qa.ResourceFixture{
 		Fixtures:    qa.HTTPFailures,
@@ -71,5 +84,8 @@ func TestDataSourceMwsNetworkConnectivityConfig_Error(t *testing.T) {
 		Read:        true,
 		NonWritable: true,
 		ID:          "_",
+        HCL: `
+        name = "def"
+        `,
 	}.ExpectError(t, "i'm a teapot")
 }