From 8e7464229c0cef53b5babc671502d4cfc09ebbb4 Mon Sep 17 00:00:00 2001 From: Miles Yucht Date: Thu, 10 Apr 2025 11:49:42 +0000 Subject: [PATCH 01/16] Use Go SDK for databricks_mws_workspaces --- mws/data_mws_workspaces.go | 8 +- mws/data_mws_workspaces_test.go | 44 +- mws/mws_workspaces_test.go | 220 ++-- mws/resource_mws_ncc_binding_test.go | 65 +- mws/resource_mws_workspaces.go | 640 +++------ mws/resource_mws_workspaces_test.go | 1794 ++++++++++---------------- 6 files changed, 1002 insertions(+), 1769 deletions(-) diff --git a/mws/data_mws_workspaces.go b/mws/data_mws_workspaces.go index 4c2d48a9a4..30792cad18 100755 --- a/mws/data_mws_workspaces.go +++ b/mws/data_mws_workspaces.go @@ -16,13 +16,17 @@ func DataSourceMwsWorkspaces() common.Resource { if c.Config.AccountID == "" { return fmt.Errorf("provider block is missing `account_id` property") } - workspaces, err := NewWorkspacesAPI(ctx, c).List(c.Config.AccountID) + a, err := c.AccountClient() + if err != nil { + return err + } + workspaces, err := a.Workspaces.List(ctx) if err != nil { return err } data.Ids = map[string]int64{} for _, v := range workspaces { - data.Ids[v.WorkspaceName] = v.WorkspaceID + data.Ids[v.WorkspaceName] = v.WorkspaceId } return nil }) diff --git a/mws/data_mws_workspaces_test.go b/mws/data_mws_workspaces_test.go index 1284b9fc9c..a08f6d157d 100755 --- a/mws/data_mws_workspaces_test.go +++ b/mws/data_mws_workspaces_test.go @@ -3,27 +3,26 @@ package mws import ( "testing" + "github.com/databricks/databricks-sdk-go/experimental/mocks" + "github.com/databricks/databricks-sdk-go/service/provisioning" "github.com/databricks/terraform-provider-databricks/qa" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" ) func TestDataSourceMwsWorkspaces(t *testing.T) { qa.ResourceFixture{ - Fixtures: []qa.HTTPFixture{ - { - Method: "GET", - Resource: "/api/2.0/accounts/abc/workspaces", - - Response: []Workspace{ - { - WorkspaceName: "bcd", - WorkspaceID: 123, - }, - { - WorkspaceName: "def", - WorkspaceID: 456, - }, + MockAccountClientFunc: func(a *mocks.MockAccountClient) { + a.GetMockWorkspacesAPI().EXPECT().List(mock.Anything).Return([]provisioning.Workspace{ + { + WorkspaceName: "bcd", + WorkspaceId: 123, + }, + { + WorkspaceName: "def", + WorkspaceId: 456, }, - }, + }, nil) }, AccountID: "abc", Resource: DataSourceMwsWorkspaces(), @@ -38,10 +37,12 @@ func TestDataSourceMwsWorkspaces(t *testing.T) { }) } -func TestCatalogsData_Error(t *testing.T) { +func TestDataSourceMwsWorkspaces_Error(t *testing.T) { qa.ResourceFixture{ + MockAccountClientFunc: func(a *mocks.MockAccountClient) { + a.GetMockWorkspacesAPI().EXPECT().List(mock.Anything).Return(nil, assert.AnError) + }, AccountID: "abc", - Fixtures: qa.HTTPFailures, Resource: DataSourceMwsWorkspaces(), Read: true, NonWritable: true, @@ -51,13 +52,8 @@ func TestCatalogsData_Error(t *testing.T) { func TestDataSourceMwsWorkspaces_Empty(t *testing.T) { qa.ResourceFixture{ - Fixtures: []qa.HTTPFixture{ - { - Method: "GET", - Resource: "/api/2.0/accounts/abc/workspaces", - - Response: []Workspace{}, - }, + MockAccountClientFunc: func(a *mocks.MockAccountClient) { + a.GetMockWorkspacesAPI().EXPECT().List(mock.Anything).Return([]provisioning.Workspace{}, nil) }, AccountID: "abc", Resource: DataSourceMwsWorkspaces(), diff --git a/mws/mws_workspaces_test.go b/mws/mws_workspaces_test.go index 9c3dc2bbaa..4356f646ec 100644 --- a/mws/mws_workspaces_test.go +++ b/mws/mws_workspaces_test.go @@ -10,14 +10,15 @@ import ( "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/logger" + "github.com/databricks/databricks-sdk-go/service/provisioning" "github.com/databricks/terraform-provider-databricks/common" "github.com/databricks/terraform-provider-databricks/internal/acceptance" "github.com/databricks/terraform-provider-databricks/internal/providers" "github.com/databricks/terraform-provider-databricks/internal/providers/sdkv2" - "github.com/databricks/terraform-provider-databricks/tokens" "github.com/hashicorp/terraform-plugin-go/tfprotov6" "github.com/hashicorp/terraform-plugin-sdk/v2/diag" "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" + "github.com/hashicorp/terraform-plugin-testing/helper/resource" "github.com/hashicorp/terraform-plugin-testing/terraform" "github.com/stretchr/testify/assert" ) @@ -76,153 +77,128 @@ func TestMwsAccWorkspaces(t *testing.T) { }) } -func TestMwsAccWorkspacesTokenUpdate(t *testing.T) { - acceptance.AccountLevel(t, acceptance.Step{ - Template: ` +func TestMwsAccWorkspaces_TokenUpdate(t *testing.T) { + baseTemplate := ` resource "databricks_mws_credentials" "this" { account_id = "{env.DATABRICKS_ACCOUNT_ID}" credentials_name = "credentials-ws-{var.RANDOM}" role_arn = "{env.TEST_CROSSACCOUNT_ARN}" } - resource "databricks_mws_customer_managed_keys" "this" { - account_id = "{env.DATABRICKS_ACCOUNT_ID}" - aws_key_info { - key_arn = "{env.TEST_MANAGED_KMS_KEY_ARN}" - key_alias = "{env.TEST_MANAGED_KMS_KEY_ALIAS}" - } - use_cases = ["MANAGED_SERVICES"] - } resource "databricks_mws_storage_configurations" "this" { account_id = "{env.DATABRICKS_ACCOUNT_ID}" storage_configuration_name = "storage-ws-{var.RANDOM}" bucket_name = "{env.TEST_ROOT_BUCKET}" } - resource "databricks_mws_networks" "this" { - account_id = "{env.DATABRICKS_ACCOUNT_ID}" - network_name = "network-ws-{var.RANDOM}" - vpc_id = "{env.TEST_VPC_ID}" - subnet_ids = [ - "{env.TEST_SUBNET_PRIVATE}", - "{env.TEST_SUBNET_PRIVATE2}", - ] - security_group_ids = [ - "{env.TEST_SECURITY_GROUP}", - ] - } + ` + tokenUpdateTemplate := func(token, customTags string) string { + tokenBlock := `` + if token != "" { + tokenBlock = fmt.Sprintf(` + token { + %s + }`, token) + } + customTagsBlock := `` + if customTags != "" { + customTagsBlock = fmt.Sprintf(` + custom_tags = { + %s + }`, customTags) + } + return fmt.Sprintf(baseTemplate+` resource "databricks_mws_workspaces" "this" { account_id = "{env.DATABRICKS_ACCOUNT_ID}" - workspace_name = "terra-{var.RANDOM}" + workspace_name = "terra-{var.STICKY_RANDOM}" aws_region = "{env.AWS_REGION}" - network_id = databricks_mws_networks.this.network_id credentials_id = databricks_mws_credentials.this.credentials_id storage_configuration_id = databricks_mws_storage_configurations.this.storage_configuration_id - managed_services_customer_managed_key_id = databricks_mws_customer_managed_keys.this.customer_managed_key_id - token { - comment = "test foo" + %s + custom_tags = { + %s } - }`, - Check: acceptance.ResourceCheckWithState("databricks_mws_workspaces.this", - func(ctx context.Context, client *common.DatabricksClient, state *terraform.InstanceState) error { - workspaceUrl, ok := state.Attributes["workspace_url"] - assert.True(t, ok, "workspace_url is absent from databricks_mws_workspaces instance state") + }`, tokenBlock, customTagsBlock) + } - workspaceClient, err := client.ClientForHost(ctx, workspaceUrl) - assert.NoError(t, err) + checkWorkspace := func(f func(instanceState map[string]string, w *databricks.WorkspaceClient) error) func(*terraform.State) error { + return func(s *terraform.State) error { + state, ok := s.RootModule().Resources["databricks_mws_workspaces.this"] + if !ok { + return fmt.Errorf("resource not found in state") + } + a := databricks.Must(databricks.NewAccountClient()) + ctx := context.Background() + workspace, err := a.Workspaces.Get(ctx, provisioning.GetWorkspaceRequest{WorkspaceId: common.MustInt64(state.Primary.ID)}) + assert.NoError(t, err) - tokensAPI := tokens.NewTokensAPI(ctx, workspaceClient) - tokens, err := tokensAPI.List() - assert.NoError(t, err) + w, err := a.GetWorkspaceClient(*workspace) + assert.NoError(t, err) - foundFoo := false - foundBar := false - for _, token := range tokens { - if token.Comment == "test foo" { - foundFoo = true - } - if token.Comment == "test bar" { - foundBar = true - } - } - assert.True(t, foundFoo) - assert.False(t, foundBar) - return nil - }), - }, - acceptance.Step{ - Template: ` - resource "databricks_mws_credentials" "this" { - account_id = "{env.DATABRICKS_ACCOUNT_ID}" - credentials_name = "credentials-ws-{var.RANDOM}" - role_arn = "{env.TEST_CROSSACCOUNT_ARN}" + return f(state.Primary.Attributes, w) } - resource "databricks_mws_customer_managed_keys" "this" { - account_id = "{env.DATABRICKS_ACCOUNT_ID}" - aws_key_info { - key_arn = "{env.TEST_MANAGED_KMS_KEY_ARN}" - key_alias = "{env.TEST_MANAGED_KMS_KEY_ALIAS}" + } + + checkTokenExists := checkWorkspace(func(instanceState map[string]string, w *databricks.WorkspaceClient) error { + tokenId := instanceState["token.0.token_id"] + assert.NotEmpty(t, tokenId) + tokens := w.Tokens.List(context.Background()) + ctx := context.Background() + for tokens.HasNext(ctx) { + if token, err := tokens.Next(ctx); err != nil && token.TokenId == tokenId { + return nil } - use_cases = ["MANAGED_SERVICES"] - } - resource "databricks_mws_storage_configurations" "this" { - account_id = "{env.DATABRICKS_ACCOUNT_ID}" - storage_configuration_name = "storage-ws-{var.RANDOM}" - bucket_name = "{env.TEST_ROOT_BUCKET}" } - resource "databricks_mws_networks" "this" { - account_id = "{env.DATABRICKS_ACCOUNT_ID}" - network_name = "network-ws-{var.RANDOM}" - vpc_id = "{env.TEST_VPC_ID}" - subnet_ids = [ - "{env.TEST_SUBNET_PRIVATE}", - "{env.TEST_SUBNET_PRIVATE2}", - ] - security_group_ids = [ - "{env.TEST_SECURITY_GROUP}", - ] - } - resource "databricks_mws_workspaces" "this" { - account_id = "{env.DATABRICKS_ACCOUNT_ID}" - workspace_name = "terra-{var.RANDOM}" - aws_region = "{env.AWS_REGION}" - - network_id = databricks_mws_networks.this.network_id - credentials_id = databricks_mws_credentials.this.credentials_id - storage_configuration_id = databricks_mws_storage_configurations.this.storage_configuration_id - managed_services_customer_managed_key_id = databricks_mws_customer_managed_keys.this.customer_managed_key_id + return fmt.Errorf("token %s not found", tokenId) + }) - token { - comment = "test bar" + var oldTokenId string + acceptance.AccountLevel(t, acceptance.Step{ + Template: tokenUpdateTemplate(`comment = "test foo"`, ""), + Check: checkTokenExists, + }, acceptance.Step{ + // Updating the comment causes the old token to be deleted and a new one to be created + Template: tokenUpdateTemplate(`comment = "test bar"`, ""), + Check: resource.ComposeAggregateTestCheckFunc( + checkTokenExists, + // Capture the token ID at the end of this step to verify it is not changed in future steps. + func(s *terraform.State) error { + state, ok := s.RootModule().Resources["databricks_mws_workspaces.this"] + if !ok { + return fmt.Errorf("resource not found in state") + } + instanceState := state.Primary.Attributes + oldTokenId = instanceState["token.0.token_id"] + return nil + }, + ), + }, acceptance.Step{ + // Modifying the tags doesn't change the token but does modify the workspace. + Template: tokenUpdateTemplate(`comment = "test bar"`, `"Key" = "Value"`), + Check: func(s *terraform.State) error { + state, ok := s.RootModule().Resources["databricks_mws_workspaces.this"] + if !ok { + return fmt.Errorf("resource not found in state") } - }`, - Check: acceptance.ResourceCheckWithState("databricks_mws_workspaces.this", - func(ctx context.Context, client *common.DatabricksClient, state *terraform.InstanceState) error { - workspaceUrl, ok := state.Attributes["workspace_url"] - assert.True(t, ok, "workspace_url is absent from databricks_mws_workspaces instance state") - - workspaceClient, err := client.ClientForHost(ctx, workspaceUrl) - assert.NoError(t, err) - - tokensAPI := tokens.NewTokensAPI(ctx, workspaceClient) - tokens, err := tokensAPI.List() - assert.NoError(t, err) - - foundFoo := false - foundBar := false - for _, token := range tokens { - if token.Comment == "test foo" { - foundFoo = true - } - if token.Comment == "test bar" { - foundBar = true - } - } - assert.False(t, foundFoo) - assert.True(t, foundBar) - return nil - }), - }) + instanceState := state.Primary.Attributes + assert.Equal(t, instanceState["custom_tags.Key"], "Value") + assert.Equal(t, instanceState["token.0.token_id"], oldTokenId) + return nil + }, + }, acceptance.Step{ + // It is also possible to modify the token comment and tags at the same time. + Template: tokenUpdateTemplate(`comment = "test quux"`, `"Key" = "Value2"`), + Check: func(s *terraform.State) error { + state, ok := s.RootModule().Resources["databricks_mws_workspaces.this"] + if !ok { + return fmt.Errorf("resource not found in state") + } + instanceState := state.Primary.Attributes + assert.Equal(t, instanceState["custom_tags.Key"], "Value2") + assert.NotEqual(t, instanceState["token.0.token_id"], oldTokenId) + return nil + }, + }) } func TestMwsAccGcpWorkspaces(t *testing.T) { diff --git a/mws/resource_mws_ncc_binding_test.go b/mws/resource_mws_ncc_binding_test.go index 4c7af575c6..9eb3c3d5d6 100644 --- a/mws/resource_mws_ncc_binding_test.go +++ b/mws/resource_mws_ncc_binding_test.go @@ -2,30 +2,36 @@ package mws import ( "testing" + "time" + "github.com/databricks/databricks-sdk-go/experimental/mocks" "github.com/databricks/databricks-sdk-go/service/provisioning" "github.com/databricks/terraform-provider-databricks/qa" + "github.com/stretchr/testify/mock" ) +var mockWorkspace = &provisioning.Workspace{ + WorkspaceId: 123456789, + WorkspaceStatus: provisioning.WorkspaceStatusRunning, +} + +var mockWaiter = &provisioning.WaitGetWorkspaceRunning[struct{}]{ + WorkspaceId: 123456789, + Poll: func(d time.Duration, f func(*provisioning.Workspace)) (*provisioning.Workspace, error) { + return mockWorkspace, nil + }, +} + func TestResourceNccBindingCreate(t *testing.T) { qa.ResourceFixture{ - Fixtures: []qa.HTTPFixture{ - { - Method: "PATCH", - Resource: "/api/2.0/accounts/abc/workspaces/123456789", - ExpectedRequest: provisioning.UpdateWorkspaceRequest{ - NetworkConnectivityConfigId: "ncc_id", - }, - }, - { - Method: "GET", - ReuseRequest: true, - Resource: "/api/2.0/accounts/abc/workspaces/123456789?", - Response: Workspace{ - WorkspaceStatus: WorkspaceStatusRunning, - WorkspaceID: 123456789, - }, - }, + MockAccountClientFunc: func(a *mocks.MockAccountClient) { + a.GetMockWorkspacesAPI().EXPECT().Update(mock.Anything, provisioning.UpdateWorkspaceRequest{ + WorkspaceId: 123456789, + NetworkConnectivityConfigId: "ncc_id", + }).Return(mockWaiter, nil) + a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ + WorkspaceId: 123456789, + }).Return(mockWorkspace, nil) }, Resource: ResourceMwsNccBinding(), AccountID: "abc", @@ -39,23 +45,14 @@ func TestResourceNccBindingCreate(t *testing.T) { func TestResourceNccBindingUpdate(t *testing.T) { qa.ResourceFixture{ - Fixtures: []qa.HTTPFixture{ - { - Method: "PATCH", - Resource: "/api/2.0/accounts/abc/workspaces/123456789", - ExpectedRequest: provisioning.UpdateWorkspaceRequest{ - NetworkConnectivityConfigId: "new_ncc_id", - }, - }, - { - Method: "GET", - ReuseRequest: true, - Resource: "/api/2.0/accounts/abc/workspaces/123456789?", - Response: Workspace{ - WorkspaceStatus: WorkspaceStatusRunning, - WorkspaceID: 123456789, - }, - }, + MockAccountClientFunc: func(a *mocks.MockAccountClient) { + a.GetMockWorkspacesAPI().EXPECT().Update(mock.Anything, provisioning.UpdateWorkspaceRequest{ + WorkspaceId: 123456789, + NetworkConnectivityConfigId: "new_ncc_id", + }).Return(mockWaiter, nil) + a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ + WorkspaceId: 123456789, + }).Return(mockWorkspace, nil) }, Resource: ResourceMwsNccBinding(), AccountID: "abc", diff --git a/mws/resource_mws_workspaces.go b/mws/resource_mws_workspaces.go index eacf73d0d8..39b98bdd42 100644 --- a/mws/resource_mws_workspaces.go +++ b/mws/resource_mws_workspaces.go @@ -3,25 +3,24 @@ package mws import ( "bytes" "context" - "encoding/json" "errors" "fmt" "log" "net" - "net/url" "strings" "time" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/retries" + "github.com/databricks/databricks-sdk-go/service/provisioning" + "github.com/databricks/databricks-sdk-go/service/settings" "github.com/databricks/terraform-provider-databricks/common" "github.com/databricks/terraform-provider-databricks/internal/docs" - "github.com/databricks/terraform-provider-databricks/tokens" "github.com/hashicorp/go-cty/cty" "github.com/hashicorp/terraform-plugin-log/tflog" "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" - "github.com/hashicorp/terraform-plugin-testing/helper/resource" ) // DefaultProvisionTimeout is the amount of minutes terraform will wait @@ -29,328 +28,53 @@ import ( // this may help with local DNS cache issues. const DefaultProvisionTimeout = 20 * time.Minute -// NewWorkspacesAPI creates MWSWorkspacesAPI instance from provider meta -func NewWorkspacesAPI(ctx context.Context, m any) WorkspacesAPI { - return WorkspacesAPI{m.(*common.DatabricksClient), ctx} -} - -// WorkspacesAPI exposes the mws workspaces API -type WorkspacesAPI struct { - client *common.DatabricksClient - context context.Context -} - -// List of workspace statuses for provisioning the workspace -const ( - WorkspaceStatusNotProvisioned = "NOT_PROVISIONED" - WorkspaceStatusProvisioning = "PROVISIONING" - WorkspaceStatusRunning = "RUNNING" - WorkspaceStatusFailed = "FAILED" - WorkspaceStatusCanceled = "CANCELLED" -) - -// WorkspaceStatusesNonRunnable is a list of statuses in which the workspace is not runnable -var WorkspaceStatusesNonRunnable = []string{WorkspaceStatusCanceled, WorkspaceStatusFailed} - -type GCP struct { - ProjectID string `json:"project_id"` -} - -type CloudResourceContainer struct { - GCP *GCP `json:"gcp"` -} - -type GCPManagedNetworkConfig struct { - SubnetCIDR string `json:"subnet_cidr" tf:"force_new"` - GKEClusterPodIPRange string `json:"gke_cluster_pod_ip_range,omitempty"` - GKEClusterServiceIPRange string `json:"gke_cluster_service_ip_range,omitempty"` -} - -type GkeConfig struct { - ConnectivityType string `json:"connectivity_type,omitempty"` - MasterIPRange string `json:"master_ip_range,omitempty"` -} - -type externalCustomerInfo struct { - CustomerName string `json:"customer_name"` - AuthoritativeUserEmail string `json:"authoritative_user_email"` - AuthoritativeUserFullName string `json:"authoritative_user_full_name"` -} - // Workspace is the object that contains all the information for deploying a workspace type Workspace struct { - AccountID string `json:"account_id"` - WorkspaceName string `json:"workspace_name"` - DeploymentName string `json:"deployment_name,omitempty"` - AwsRegion string `json:"aws_region,omitempty"` // required for AWS, not allowed for GCP - CredentialsID string `json:"credentials_id,omitempty"` // required for AWS, not allowed for GCP - CustomerManagedKeyID string `json:"customer_managed_key_id,omitempty"` // just for compatibility, will be removed - StorageConfigurationID string `json:"storage_configuration_id,omitempty"` // required for AWS, not allowed for GCP - ManagedServicesCustomerManagedKeyID string `json:"managed_services_customer_managed_key_id,omitempty"` - StorageCustomerManagedKeyID string `json:"storage_customer_managed_key_id,omitempty"` - PricingTier string `json:"pricing_tier,omitempty" tf:"computed"` - PrivateAccessSettingsID string `json:"private_access_settings_id,omitempty"` - NetworkID string `json:"network_id,omitempty" tf:"suppress_diff"` - IsNoPublicIPEnabled bool `json:"is_no_public_ip_enabled" tf:"optional,default:true"` - WorkspaceID int64 `json:"workspace_id,omitempty" tf:"computed"` - WorkspaceURL string `json:"workspace_url,omitempty" tf:"computed"` - WorkspaceStatus string `json:"workspace_status,omitempty" tf:"computed"` - WorkspaceStatusMessage string `json:"workspace_status_message,omitempty" tf:"computed"` - CreationTime int64 `json:"creation_time,omitempty" tf:"computed"` - ExternalCustomerInfo *externalCustomerInfo `json:"external_customer_info,omitempty"` - CloudResourceBucket *CloudResourceContainer `json:"cloud_resource_container,omitempty"` - GCPManagedNetworkConfig *GCPManagedNetworkConfig `json:"gcp_managed_network_config,omitempty" tf:"suppress_diff"` - GkeConfig *GkeConfig `json:"gke_config,omitempty" tf:"suppress_diff"` - Cloud string `json:"cloud,omitempty" tf:"computed"` - Location string `json:"location,omitempty"` - CustomTags map[string]string `json:"custom_tags,omitempty"` // Optional for AWS, not allowed for GCP -} - -// this type alias hack is required for Marshaller to work without an infinite loop -type aWorkspace Workspace - -// MarshalJSON is required to overcome the limitations of `omitempty` usage with reflect_resource.go -// for workspace creation in Accounts API for AWS and GCP. It exits early on AWS and picks only -// the relevant fields for GCP. -func (w *Workspace) MarshalJSON() ([]byte, error) { - if w.Cloud != "gcp" { - return json.Marshal(aWorkspace(*w)) - } - workspaceCreationRequest := map[string]any{ - "account_id": w.AccountID, - "cloud": w.Cloud, - "cloud_resource_container": w.CloudResourceBucket, - "location": w.Location, - "workspace_name": w.WorkspaceName, - } - if w.NetworkID != "" { - workspaceCreationRequest["network_id"] = w.NetworkID - } - if w.PrivateAccessSettingsID != "" { - workspaceCreationRequest["private_access_settings_id"] = w.PrivateAccessSettingsID - } - if w.GkeConfig != nil { - workspaceCreationRequest["gke_config"] = w.GkeConfig - } - if w.GCPManagedNetworkConfig != nil { - workspaceCreationRequest["gcp_managed_network_config"] = w.GCPManagedNetworkConfig - } - if w.ManagedServicesCustomerManagedKeyID != "" { - workspaceCreationRequest["managed_services_customer_managed_key_id"] = w.ManagedServicesCustomerManagedKeyID - } - if w.StorageCustomerManagedKeyID != "" { - workspaceCreationRequest["storage_customer_managed_key_id"] = w.StorageCustomerManagedKeyID - } - return json.Marshal(workspaceCreationRequest) + *provisioning.Workspace + CustomerManagedKeyID string `json:"customer_managed_key_id,omitempty"` // just for compatibility, will be removed + Token *Token `json:"token,omitempty"` + WorkspaceURL string `json:"workspace_url,omitempty" tf:"computed"` + GcpWorkspaceSa string `json:"gcp_workspace_sa" tf:"computed"` } -// Create deploys the workspace and waits till it's properly running. -// In case of error, it removes the failed deployment and returns the message -func (a WorkspacesAPI) Create(ws *Workspace, timeout time.Duration) error { - if a.client.IsGcp() { - ws.Cloud = "gcp" - } - workspacesAPIPath := fmt.Sprintf("/accounts/%s/workspaces", ws.AccountID) - err := a.client.Post(a.context, workspacesAPIPath, ws, &ws) - if err != nil { - return err - } - if err = a.WaitForRunning(*ws, timeout); err != nil { - log.Printf("[ERROR] Deleting failed workspace: %s", err) - if derr := a.Delete(ws.AccountID, fmt.Sprintf("%d", ws.WorkspaceID)); derr != nil { - return fmt.Errorf("%s - %s", err, derr) +// wait for DNS caches to refresh, as sometimes we cannot make +// API calls to new workspaces immediately after it's created +func verifyWorkspaceReachable(ctx context.Context, w *databricks.WorkspaceClient) error { + return retries.New[struct{}](retries.WithRetryFunc(func(err error) bool { + var dnsError *net.DNSError + if errors.As(err, &dnsError) { + err = fmt.Errorf("workspace %s is not yet reachable: %s", w.Config.Host, dnsError) + log.Printf("[INFO] %s", err) + // expected to retry on: dial tcp: lookup XXX: no such host + return true } + return false + })).Wait(ctx, func(ctx context.Context) error { + ctx, cancel := context.WithTimeout(ctx, 10*time.Second) + defer cancel() + _, err := w.CurrentUser.Me(ctx) return err - } - if ws.WorkspaceURL == "" { - // WorkspaceURL is computed, yet very important field - host := generateWorkspaceHostname(a.client, *ws) - ws.WorkspaceURL = fmt.Sprintf("https://%s", host) - } - return nil + }) } -// generateWorkspaceHostname computes the hostname for the specified workspace, -// given the account console hostname. -func generateWorkspaceHostname(client *common.DatabricksClient, ws Workspace) string { - u, err := url.Parse(client.Config.Host) - if err != nil { - // Fallback. - log.Printf("[WARN] Unable to parse URL from client host: %v", err) - return ws.DeploymentName + ".cloud.databricks.com" - } - - // We expect the account console hostname to be of the form `accounts.foo[.bar]...` - // The workspace hostname can be generated by replacing `accounts` with the deployment name. - // If the hostname is an IP address, we're in testing mode and do fallback. - chunks := strings.Split(u.Hostname(), ".") - if len(chunks) == 0 || net.ParseIP(u.Hostname()) != nil { - // Fallback. - log.Printf("[WARN] Unable to split client host: %v", u.Hostname()) - return ws.DeploymentName + ".cloud.databricks.com" +func explainWorkspaceFailure(ctx context.Context, a *databricks.AccountClient, workspace *provisioning.Workspace) error { + errorBase := fmt.Sprintf("workspace status message: %s", workspace.WorkspaceStatusMessage) + if workspace.NetworkId == "" { + return errors.New(errorBase) } - chunks[0] = ws.DeploymentName - return strings.Join(chunks, ".") -} - -func (a WorkspacesAPI) verifyWorkspaceReachable(ws Workspace) *resource.RetryError { - ctx, cancel := context.WithTimeout(a.context, 10*time.Second) - defer cancel() - // wait for DNS caches to refresh, as sometimes we cannot make - // API calls to new workspaces immediately after it's created - wsClient, err := a.client.ClientForHost(a.context, ws.WorkspaceURL) + network, err := a.Networks.Get(ctx, provisioning.GetNetworkRequest{NetworkId: workspace.NetworkId}) if err != nil { - return resource.NonRetryableError(err) - } - // make a request to SCIM API, just to verify there are no errors - var response map[string]any - err = wsClient.Get(ctx, "/preview/scim/v2/Me", nil, &response) - var dnsError *net.DNSError - if errors.As(err, &dnsError) { - err = fmt.Errorf("workspace %s is not yet reachable: %s", - ws.WorkspaceURL, dnsError) - log.Printf("[INFO] %s", err) - // expected to retry on: dial tcp: lookup XXX: no such host - return resource.RetryableError(err) - } - return nil -} - -func (a WorkspacesAPI) explainWorkspaceFailure(ws Workspace) error { - if ws.NetworkID == "" { - return errors.New(ws.WorkspaceStatusMessage) - } - network, nerr := NewNetworksAPI(a.context, a.client).Read(ws.AccountID, ws.NetworkID) - if nerr != nil { - return fmt.Errorf("failed to start workspace. Cannot read network: %s", nerr) + return fmt.Errorf("%s; cannot read network: %w", errorBase, err) } var strBuffer bytes.Buffer for _, networkHealth := range network.ErrorMessages { - strBuffer.WriteString(fmt.Sprintf("error: %s;error_msg: %s;", - networkHealth.ErrorType, networkHealth.ErrorMessage)) - } - return fmt.Errorf("Workspace failed to create: %v, network error message: %v", - ws.WorkspaceStatusMessage, strBuffer.String()) -} - -// WaitForRunning will wait until workspace is running, otherwise will try to explain why it failed -func (a WorkspacesAPI) WaitForRunning(ws Workspace, timeout time.Duration) error { - return resource.RetryContext(a.context, timeout, func() *resource.RetryError { - workspace, err := a.Read(ws.AccountID, fmt.Sprintf("%d", ws.WorkspaceID)) - if err != nil { - return resource.NonRetryableError(err) - } - switch workspace.WorkspaceStatus { - case WorkspaceStatusRunning: - log.Printf("[INFO] Workspace is now running") - if strings.Contains(ws.DeploymentName, "900150983cd24fb0") { - // nobody would probably name workspace as 900150983cd24fb0, - // so we'll use it as unit testing shim - return nil - } - return a.verifyWorkspaceReachable(workspace) - case WorkspaceStatusCanceled, WorkspaceStatusFailed: - log.Printf("[ERROR] Cannot start workspace: %s", workspace.WorkspaceStatusMessage) - err = a.explainWorkspaceFailure(workspace) - return resource.NonRetryableError(err) - default: - log.Printf("[INFO] Workspace %s is %s: %s", workspace.DeploymentName, - workspace.WorkspaceStatus, workspace.WorkspaceStatusMessage) - return resource.RetryableError(fmt.Errorf("%s", workspace.WorkspaceStatusMessage)) - } - }) -} - -var workspaceRunningUpdatesAllowed = []string{"credentials_id", "network_id", "storage_customer_managed_key_id", "private_access_settings_id", "managed_services_customer_managed_key_id", "custom_tags"} - -// UpdateRunning will update running workspace with couple of possible fields -func (a WorkspacesAPI) UpdateRunning(ws Workspace, timeout time.Duration) error { - workspacesAPIPath := fmt.Sprintf("/accounts/%s/workspaces/%d", ws.AccountID, ws.WorkspaceID) - request := map[string]any{} - - if ws.CredentialsID != "" { - request["credentials_id"] = ws.CredentialsID - } - - // The ID of the workspace's network configuration object. Used only if you already use a customer-managed VPC. - // This change is supported only if you specified a network configuration ID when the workspace was created. - // In other words, you cannot switch from a Databricks-managed VPC to a customer-managed VPC. This parameter - // is available for updating both failed and running workspaces. - if ws.NetworkID != "" { - request["network_id"] = ws.NetworkID - } - - if ws.PrivateAccessSettingsID != "" { - request["private_access_settings_id"] = ws.PrivateAccessSettingsID - } - if ws.StorageCustomerManagedKeyID != "" { - request["storage_customer_managed_key_id"] = ws.StorageCustomerManagedKeyID - } - if ws.CustomTags != nil { - if !a.client.IsAws() { - return fmt.Errorf("custom_tags are only allowed for AWS workspaces") - } - request["custom_tags"] = ws.CustomTags - } - - if len(request) == 0 { - return nil - } - - err := a.client.Patch(a.context, workspacesAPIPath, request) - if err != nil { - return err - } - return a.WaitForRunning(ws, timeout) -} - -// Read will return the mws workspace metadata and status of the workspace deployment -func (a WorkspacesAPI) Read(mwsAcctID, workspaceID string) (Workspace, error) { - var mwsWorkspace Workspace - workspacesAPIPath := fmt.Sprintf("/accounts/%s/workspaces/%s", mwsAcctID, workspaceID) - err := a.client.Get(a.context, workspacesAPIPath, nil, &mwsWorkspace) - if err == nil && mwsWorkspace.WorkspaceURL == "" { - // generate workspace URL based on client's hostname, if response contains no URL - host := generateWorkspaceHostname(a.client, mwsWorkspace) - mwsWorkspace.WorkspaceURL = fmt.Sprintf("https://%s", host) - } - return mwsWorkspace, err -} - -// Delete will delete the configuration for the workspace given a workspace id -// and wait till it's properly removed -func (a WorkspacesAPI) Delete(mwsAcctID, workspaceID string) error { - workspacesAPIPath := fmt.Sprintf("/accounts/%s/workspaces/%s", mwsAcctID, workspaceID) - err := a.client.Delete(a.context, workspacesAPIPath, nil) - if err != nil { - return err + strBuffer.WriteString(fmt.Sprintf("error: %s;error_msg: %s;", networkHealth.ErrorType, networkHealth.ErrorMessage)) } - return resource.RetryContext(a.context, 15*time.Minute, func() *resource.RetryError { - workspace, err := a.Read(mwsAcctID, workspaceID) - if apierr.IsMissing(err) { - log.Printf("[INFO] Workspace %s/%s is removed.", mwsAcctID, workspaceID) - return nil - } - if err != nil { - return resource.NonRetryableError(err) - } - msg := fmt.Errorf("Workspace %s is not removed yet. Workspace status: %s %s", - workspace.WorkspaceName, workspace.WorkspaceStatus, workspace.WorkspaceStatusMessage) - log.Printf("[INFO] %s", msg) - return resource.RetryableError(msg) - }) -} - -// List will list all workspaces in a given mws account -func (a WorkspacesAPI) List(mwsAcctID string) ([]Workspace, error) { - var mwsWorkspacesList []Workspace - workspacesAPIPath := fmt.Sprintf("/accounts/%s/workspaces", mwsAcctID) - err := a.client.Get(a.context, workspacesAPIPath, nil, &mwsWorkspacesList) - return mwsWorkspacesList, err + return fmt.Errorf("%s, network error message: %s", errorBase, strBuffer.String()) } type Token struct { - LifetimeSeconds int32 `json:"lifetime_seconds,omitempty" tf:"default:2592000"` + LifetimeSeconds int64 `json:"lifetime_seconds,omitempty" tf:"default:2592000"` Comment string `json:"comment,omitempty" tf:"default:Terraform PAT"` TokenID string `json:"token_id,omitempty" tf:"computed"` TokenValue SensitiveString `json:"token_value,omitempty" tf:"computed,sensitive"` @@ -366,37 +90,22 @@ func (s SensitiveString) String() string { return "****" } -// ephemeral entity to use with StructToData() -type WorkspaceToken struct { - WorkspaceURL string `json:"workspace_url,omitempty"` - Token *Token `json:"token,omitempty"` -} - -func CreateTokenIfNeeded(workspacesAPI WorkspacesAPI, - workspaceSchema map[string]*schema.Schema, d *schema.ResourceData) error { - var wsToken WorkspaceToken - common.DataToStructPointer(d, workspaceSchema, &wsToken) - if wsToken.Token == nil { - return nil - } - client, err := workspacesAPI.client.ClientForHost(workspacesAPI.context, wsToken.WorkspaceURL) - if err != nil { - return err - } - tokensAPI := tokens.NewTokensAPI(workspacesAPI.context, client) - lifetime := time.Duration(wsToken.Token.LifetimeSeconds) * time.Second - token, err := tokensAPI.Create(lifetime, wsToken.Token.Comment) +func createToken(ctx context.Context, w *databricks.WorkspaceClient, t *Token) error { + token, err := w.Tokens.Create(ctx, settings.CreateTokenRequest{ + LifetimeSeconds: t.LifetimeSeconds, + Comment: t.Comment, + }) if isInvalidClient(err) { return fmt.Errorf("cannot create token: the principal used by Databricks (client ID %s) is not authorized to create a token in this workspace. "+ "If this is a UC-enabled workspace, add this client to the workspace, either using databricks_mws_permission_assignment or manually (see https://docs.databricks.com/en/admin/users-groups/service-principals.html#assign-a-service-principal-to-a-workspace-using-the-account-console for instructions). "+ - "If this is not a UC-enabled workspace, remove the token block from this configuration and create a workspace-level service principal to configure resources in the workspace (see https://docs.databricks.com/en/admin/users-groups/service-principals.html#add-a-service-principal-to-a-workspace-using-the-workspace-admin-settings for instructions)", client.Config.ClientID) + "If this is not a UC-enabled workspace, remove the token block from this configuration and create a workspace-level service principal to configure resources in the workspace (see https://docs.databricks.com/en/admin/users-groups/service-principals.html#add-a-service-principal-to-a-workspace-using-the-workspace-admin-settings for instructions)", w.Config.ClientID) } if err != nil { return fmt.Errorf("cannot create token: %w", err) } - wsToken.Token.TokenID = token.TokenInfo.TokenID - wsToken.Token.TokenValue = SensitiveString(token.TokenValue) - return common.StructToData(wsToken, workspaceSchema, d) + t.TokenValue = SensitiveString(token.TokenValue) + t.TokenID = token.TokenInfo.TokenId + return nil } // isInvalidClient checks whether the API request failed due to the client being invalid. @@ -410,28 +119,17 @@ func isInvalidClient(err error) bool { return errors.Is(err, databricks.ErrUnauthenticated) } -func EnsureTokenExistsIfNeeded(a WorkspacesAPI, - workspaceSchema map[string]*schema.Schema, d *schema.ResourceData) error { - var wsToken WorkspaceToken - common.DataToStructPointer(d, workspaceSchema, &wsToken) - if wsToken.Token == nil { - return nil - } - client, err := a.client.ClientForHost(a.context, wsToken.WorkspaceURL) - if err != nil { - return err - } - tokensAPI := tokens.NewTokensAPI(a.context, client) - _, err = tokensAPI.Read(wsToken.Token.TokenID) +func ensureTokenExists(ctx context.Context, w *databricks.WorkspaceClient, token *Token) error { + _, err := w.CurrentUser.Me(ctx) // If we cannot authenticate to the workspace and we're using an in-house OAuth principal, // log a warning but do not fail. This can happen if the provider is authenticated with a // different principal than was used to create the workspace. if isInvalidClient(err) { - tflog.Debug(a.context, fmt.Sprintf("unable to fetch token with ID %s from workspace using the provided service principal, continuing", wsToken.Token.TokenID)) + tflog.Debug(ctx, fmt.Sprintf("unable to fetch token with ID %s from workspace using the provided service principal, continuing", token.TokenID)) return nil } if apierr.IsMissing(err) { - return CreateTokenIfNeeded(a, workspaceSchema, d) + return createToken(ctx, w, token) } if err != nil { return fmt.Errorf("cannot read token: %w", err) @@ -439,71 +137,47 @@ func EnsureTokenExistsIfNeeded(a WorkspacesAPI, return nil } -func removeTokenIfNeeded(a WorkspacesAPI, tokenID string, d *schema.ResourceData) error { - client, err := a.client.ClientForHost(a.context, d.Get("workspace_url").(string)) - if err != nil { - return err - } - tokensAPI := tokens.NewTokensAPI(a.context, client) - err = tokensAPI.Delete(tokenID) +func removeToken(ctx context.Context, w *databricks.WorkspaceClient, tokenID string) error { + err := w.Tokens.Delete(ctx, settings.RevokeTokenRequest{TokenId: tokenID}) if isInvalidClient(err) { - tflog.Debug(a.context, fmt.Sprintf("unable to delete token with ID %s from workspace using the provided service principal, continuing", tokenID)) + tflog.Debug(ctx, fmt.Sprintf("unable to delete token with ID %s from workspace using the provided service principal, continuing", tokenID)) return nil } if err != nil { return fmt.Errorf("cannot remove token: %w", err) } - return d.Set("token", nil) -} - -func UpdateTokenIfNeeded(workspacesAPI WorkspacesAPI, - workspaceSchema map[string]*schema.Schema, d *schema.ResourceData) error { - o, n := d.GetChange("token") - old, new := o.([]any), n.([]any) - if d.HasChange("token") { - switch { - case len(old) == 0 && len(new) > 0: // create - return CreateTokenIfNeeded(workspacesAPI, workspaceSchema, d) - case len(old) > 0 && len(new) == 0: // delete - raw := old[0].(map[string]any) - id := raw["token_id"].(string) - return removeTokenIfNeeded(workspacesAPI, id, d) - case len(old) > 0 && len(new) > 0: // delete & create - rawOld := old[0].(map[string]any) - id := rawOld["token_id"].(string) - err := removeTokenIfNeeded(workspacesAPI, id, d) - if err != nil { - return err - } - rawNew := new[0].(map[string]any) - d.Set("token", []any{ - map[string]any{ - "lifetime_seconds": rawNew["lifetime_seconds"], - "comment": rawNew["comment"], - }, - }) - return CreateTokenIfNeeded(workspacesAPI, workspaceSchema, d) - } - } return nil } // ResourceMwsWorkspaces manages E2 workspaces func ResourceMwsWorkspaces() common.Resource { + var computedFields = map[string]struct{}{ + "workspace_id": {}, + "workspace_url": {}, + "workspace_status": {}, + "workspace_status_message": {}, + "creation_time": {}, + "gke_config": {}, + "pricing_tier": {}, + } + + var workspaceRunningUpdatesAllowed = map[string]struct{}{ + "credentials_id": {}, + "network_id": {}, + "storage_customer_managed_key_id": {}, + "private_access_settings_id": {}, + "managed_services_customer_managed_key_id": {}, + "custom_tags": {}, + } + workspaceSchema := common.StructToSchema(Workspace{}, func(s map[string]*schema.Schema) map[string]*schema.Schema { for name, fieldSchema := range s { - if fieldSchema.Computed { - // skip checking all changes from remote state - continue - } - fieldSchema.ForceNew = true - for _, allowed := range workspaceRunningUpdatesAllowed { - if allowed == name { - // allow updating only a few specific fields - fieldSchema.ForceNew = false - break - } + if _, ok := computedFields[name]; ok { + fieldSchema.Computed = true + } else { + _, ok := workspaceRunningUpdatesAllowed[name] + fieldSchema.ForceNew = !ok } } s["account_id"].Sensitive = true @@ -528,16 +202,7 @@ func ResourceMwsWorkspaces() common.Resource { s["customer_managed_key_id"].ConflictsWith = []string{"managed_services_customer_managed_key_id", "storage_customer_managed_key_id"} s["managed_services_customer_managed_key_id"].ConflictsWith = []string{"customer_managed_key_id"} s["storage_customer_managed_key_id"].ConflictsWith = []string{"customer_managed_key_id"} - // manage workspace-specific PAT token - s["token"] = common.StructToSchema(WorkspaceToken{}, - func(m map[string]*schema.Schema) map[string]*schema.Schema { - return m - })["token"] - s["gcp_workspace_sa"] = &schema.Schema{ - Type: schema.TypeString, - Computed: true, - } docOptions := docs.DocOptions{ Section: docs.Guides, Slug: "gcp-workspace", @@ -574,42 +239,81 @@ func ResourceMwsWorkspaces() common.Resource { }, }, Create: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error { - var workspace Workspace - workspacesAPI := NewWorkspacesAPI(ctx, c) - common.DataToStructPointer(d, workspaceSchema, &workspace) + a, err := c.AccountClient() + if err != nil { + return err + } + var createWorkspaceRequest provisioning.CreateWorkspaceRequest + common.DataToStructPointer(d, workspaceSchema, &createWorkspaceRequest) + var workspaceConfig Workspace + common.DataToStructPointer(d, workspaceSchema, &workspaceConfig) + + // Validate required fields by cloud if err := requireFields(c.IsAws(), d, "aws_region", "credentials_id", "storage_configuration_id"); err != nil { return err } if err := requireFields(c.IsGcp(), d, "location"); err != nil { return err } - if !c.IsAws() && workspace.CustomTags != nil { + if !c.IsAws() && createWorkspaceRequest.CustomTags != nil { return fmt.Errorf("custom_tags are only allowed for AWS workspaces") } - if len(workspace.CustomerManagedKeyID) > 0 && len(workspace.ManagedServicesCustomerManagedKeyID) == 0 { + if customerManagedKeyId, ok := d.GetOk("customer_managed_key_id"); ok && createWorkspaceRequest.ManagedServicesCustomerManagedKeyId == "" { log.Print("[INFO] Using existing customer_managed_key_id as value for new managed_services_customer_managed_key_id") - workspace.ManagedServicesCustomerManagedKeyID = workspace.CustomerManagedKeyID - workspace.CustomerManagedKeyID = "" + createWorkspaceRequest.ManagedServicesCustomerManagedKeyId = customerManagedKeyId.(string) } - if err := workspacesAPI.Create(&workspace, d.Timeout(schema.TimeoutCreate)); err != nil { + + // Create the workspace. If creation fails, clean it up and return. + wait, err := a.Workspaces.Create(ctx, createWorkspaceRequest) + if err != nil { return err } - d.Set("workspace_id", workspace.WorkspaceID) - d.Set("workspace_url", workspace.WorkspaceURL) - if workspace.Cloud == "gcp" { - d.Set("gcp_workspace_sa", fmt.Sprintf("db-%d@prod-gcp-%s.iam.gserviceaccount.com", - workspace.WorkspaceID, workspace.Location)) + workspace, err := wait.GetWithTimeout(d.Timeout(schema.TimeoutCreate)) + if err != nil { + err = fmt.Errorf("%w: %w", err, explainWorkspaceFailure(ctx, a, workspace)) + log.Printf("[ERROR] Deleting failed workspace: %s", err) + if derr := a.Workspaces.Delete(ctx, provisioning.DeleteWorkspaceRequest{WorkspaceId: wait.Response.WorkspaceId}); derr != nil { + return fmt.Errorf("workspace creation failed: %w; failed workspace cleanup failed: %w", err, derr) + } + return err + } + + // Once the workspace is running, wait for the API to be available by polling the SCIM Me endpoint. + w, err := a.GetWorkspaceClient(*workspace) + if err != nil { + return err + } + if err := verifyWorkspaceReachable(ctx, w); err != nil { + return err + } + workspaceConfig.Workspace = workspace + workspaceConfig.WorkspaceURL = fmt.Sprintf("https://%s", w.Config.CanonicalHostName()) + if c.IsGcp() { + workspaceConfig.GcpWorkspaceSa = fmt.Sprintf("db-%d@prod-gcp-%s.iam.gserviceaccount.com", workspace.WorkspaceId, workspace.Location) + } + + // Create a token if requested + if workspaceConfig.Token != nil { + err = createToken(ctx, w, workspaceConfig.Token) + if err != nil { + return err + } } p.Pack(d) - return CreateTokenIfNeeded(workspacesAPI, workspaceSchema, d) + return common.StructToData(workspaceConfig, workspaceSchema, d) }, Read: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error { - accountID, workspaceID, err := p.Unpack(d) + a, err := c.AccountClient() + if err != nil { + return err + } + _, workspaceID, err := p.Unpack(d) if err != nil { return err } - workspacesAPI := NewWorkspacesAPI(ctx, c) - workspace, err := workspacesAPI.Read(accountID, workspaceID) + var workspaceConfig Workspace + common.DataToStructPointer(d, workspaceSchema, &workspaceConfig) + workspace, err := a.Workspaces.Get(ctx, provisioning.GetWorkspaceRequest{WorkspaceId: common.MustInt64(workspaceID)}) if err != nil { return err } @@ -620,39 +324,97 @@ func ResourceMwsWorkspaces() common.Resource { // Default the value of `is_no_public_ip_enabled` because it isn't part of the GET payload. // The field is only used on creation and we therefore suppress all diffs. - workspace.IsNoPublicIPEnabled = true - if err = common.StructToData(workspace, workspaceSchema, d); err != nil { + workspace.IsNoPublicIpEnabled = true + workspaceConfig.Workspace = workspace + _, err = a.Workspaces.WaitGetWorkspaceRunning(ctx, workspace.WorkspaceId, d.Timeout(schema.TimeoutRead), nil) + if err != nil { return err } - err = workspacesAPI.WaitForRunning(workspace, d.Timeout(schema.TimeoutRead)) + w, err := a.GetWorkspaceClient(*workspace) if err != nil { return err } - return EnsureTokenExistsIfNeeded(workspacesAPI, workspaceSchema, d) + if workspaceConfig.Token != nil { + if err := ensureTokenExists(ctx, w, workspaceConfig.Token); err != nil { + return err + } + } + return common.StructToData(workspaceConfig, workspaceSchema, d) }, Update: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error { - var workspace Workspace - common.DataToStructPointer(d, workspaceSchema, &workspace) - if len(workspace.CustomerManagedKeyID) > 0 && len(workspace.ManagedServicesCustomerManagedKeyID) == 0 { + a, err := c.AccountClient() + if err != nil { + return err + } + var updateWorkspaceRequest provisioning.UpdateWorkspaceRequest + common.DataToStructPointer(d, workspaceSchema, &updateWorkspaceRequest) + var workspaceConfig Workspace + common.DataToStructPointer(d, workspaceSchema, &workspaceConfig) + if customerManagedKeyId, ok := d.GetOk("customer_managed_key_id"); ok && updateWorkspaceRequest.ManagedServicesCustomerManagedKeyId == "" { log.Print("[INFO] Using existing customer_managed_key_id as value for new managed_services_customer_managed_key_id") - workspace.ManagedServicesCustomerManagedKeyID = workspace.CustomerManagedKeyID - workspace.CustomerManagedKeyID = "" + updateWorkspaceRequest.ManagedServicesCustomerManagedKeyId = customerManagedKeyId.(string) } - workspacesAPI := NewWorkspacesAPI(ctx, c) + + // If the workspace has been modified, call the update API and wait for it to be ready. if d.HasChangeExcept("token") { - err := workspacesAPI.UpdateRunning(workspace, d.Timeout(schema.TimeoutUpdate)) + wait, err := a.Workspaces.Update(ctx, updateWorkspaceRequest) if err != nil { return err } + workspace, err := wait.GetWithTimeout(d.Timeout(schema.TimeoutUpdate)) + if err != nil { + return fmt.Errorf("%w: %w", err, explainWorkspaceFailure(ctx, a, workspace)) + } + workspaceConfig.Workspace = workspace } - return UpdateTokenIfNeeded(workspacesAPI, workspaceSchema, d) + + // If the `token` field has been modified, update the token correspondingly. + if d.HasChange("token") { + w, err := a.GetWorkspaceClient(*workspaceConfig.Workspace) + if err != nil { + return err + } + + // If there was a token present in the config before, revoke it. + rawOld, _ := d.GetChange("token") + oldToken := rawOld.([]map[string]any) + if len(oldToken) > 0 { + raw := oldToken[0] + id := raw["token_id"].(string) + if err := removeToken(ctx, w, id); err != nil { + return err + } + } + + // If there is a token present in the config now, create a new one. + if workspaceConfig.Token != nil { + if err := createToken(ctx, w, workspaceConfig.Token); err != nil { + return err + } + } + } + return common.StructToData(workspaceConfig, workspaceSchema, d) }, Delete: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error { - accountID, workspaceID, err := p.Unpack(d) + a, err := c.AccountClient() + if err != nil { + return err + } + _, workspaceID, err := p.Unpack(d) if err != nil { return err } - return NewWorkspacesAPI(ctx, c).Delete(accountID, workspaceID) + if err := a.Workspaces.Delete(ctx, provisioning.DeleteWorkspaceRequest{WorkspaceId: common.MustInt64(workspaceID)}); err != nil && !apierr.IsMissing(err) { + return err + } + // Wait for delete + return retries.New[struct{}]().Wait(ctx, func(ctx context.Context) error { + _, err := a.Workspaces.Get(ctx, provisioning.GetWorkspaceRequest{WorkspaceId: common.MustInt64(workspaceID)}) + if err != nil && apierr.IsMissing(err) { + return nil + } + return fmt.Errorf("workspace %s still exists", d.Id()) + }) }, CustomizeDiff: func(ctx context.Context, d *schema.ResourceDiff) error { old, new := d.GetChange("private_access_settings_id") diff --git a/mws/resource_mws_workspaces_test.go b/mws/resource_mws_workspaces_test.go index 87c18334d1..3baf4fc58a 100644 --- a/mws/resource_mws_workspaces_test.go +++ b/mws/resource_mws_workspaces_test.go @@ -2,69 +2,70 @@ package mws import ( "context" + "errors" "fmt" + "net" "testing" "time" "github.com/databricks/databricks-sdk-go/apierr" - "github.com/databricks/databricks-sdk-go/client" - "github.com/databricks/databricks-sdk-go/config" - "github.com/databricks/terraform-provider-databricks/common" - "github.com/databricks/terraform-provider-databricks/tokens" + "github.com/databricks/databricks-sdk-go/experimental/mocks" + "github.com/databricks/databricks-sdk-go/listing" + "github.com/databricks/databricks-sdk-go/service/iam" + "github.com/databricks/databricks-sdk-go/service/provisioning" + "github.com/databricks/databricks-sdk-go/service/settings" "github.com/databricks/terraform-provider-databricks/qa" "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" + "github.com/stretchr/testify/mock" ) func TestResourceWorkspaceCreate(t *testing.T) { + // Define a mock workspace that can be reused + mockWorkspace := &provisioning.Workspace{ + WorkspaceId: 1234, + WorkspaceStatus: provisioning.WorkspaceStatusRunning, + WorkspaceName: "labdata", + DeploymentName: "900150983cd24fb0", + AwsRegion: "us-east-1", + CredentialsId: "bcd", + StorageConfigurationId: "ghi", + NetworkId: "fgh", + ManagedServicesCustomerManagedKeyId: "def", + StorageCustomerManagedKeyId: "def", + AccountId: "abc", + CustomTags: map[string]string{ + "SoldToCode": "1234", + }, + } + + // Create a mock waiter + mockWaiter := &provisioning.WaitGetWorkspaceRunning[provisioning.Workspace]{ + WorkspaceId: 1234, + Poll: func(d time.Duration, f func(*provisioning.Workspace)) (*provisioning.Workspace, error) { + return mockWorkspace, nil + }, + } + d, err := qa.ResourceFixture{ - Fixtures: []qa.HTTPFixture{ - { - Method: "POST", - Resource: "/api/2.0/accounts/abc/workspaces", - ExpectedRequest: Workspace{ - AccountID: "abc", - IsNoPublicIPEnabled: true, - WorkspaceName: "labdata", - DeploymentName: "900150983cd24fb0", - AwsRegion: "us-east-1", - CredentialsID: "bcd", - StorageConfigurationID: "ghi", - NetworkID: "fgh", - ManagedServicesCustomerManagedKeyID: "def", - StorageCustomerManagedKeyID: "def", - CustomTags: map[string]string{ - "SoldToCode": "1234", - }, - }, - Response: Workspace{ - WorkspaceID: 1234, - AccountID: "abc", - DeploymentName: "900150983cd24fb0", - }, - }, - { - Method: "GET", - ReuseRequest: true, - Resource: "/api/2.0/accounts/abc/workspaces/1234", - Response: Workspace{ - WorkspaceID: 1234, - WorkspaceStatus: WorkspaceStatusRunning, - WorkspaceName: "labdata", - DeploymentName: "900150983cd24fb0", - AwsRegion: "us-east-1", - CredentialsID: "bcd", - StorageConfigurationID: "ghi", - NetworkID: "fgh", - ManagedServicesCustomerManagedKeyID: "def", - StorageCustomerManagedKeyID: "def", - AccountID: "abc", - CustomTags: map[string]string{ - "SoldToCode": "1234", - }, + MockAccountClientFunc: func(a *mocks.MockAccountClient) { + a.GetMockWorkspacesAPI().EXPECT().Create(mock.Anything, provisioning.CreateWorkspaceRequest{ + IsNoPublicIpEnabled: true, + WorkspaceName: "labdata", + DeploymentName: "900150983cd24fb0", + AwsRegion: "us-east-1", + CredentialsId: "bcd", + StorageConfigurationId: "ghi", + NetworkId: "fgh", + ManagedServicesCustomerManagedKeyId: "def", + StorageCustomerManagedKeyId: "def", + CustomTags: map[string]string{ + "SoldToCode": "1234", }, - }, + }).Return(mockWaiter, nil) + a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ + WorkspaceId: 1234, + }).Return(mockWorkspace, nil) }, Resource: ResourceMwsWorkspaces(), State: map[string]any{ @@ -88,48 +89,47 @@ func TestResourceWorkspaceCreate(t *testing.T) { } func TestResourceWorkspaceCreateGcp(t *testing.T) { + // Define a mock workspace that can be reused + mockWorkspace := &provisioning.Workspace{ + WorkspaceId: 1234, + WorkspaceStatus: provisioning.WorkspaceStatusRunning, + WorkspaceName: "labdata", + DeploymentName: "900150983cd24fb0", + AccountId: "abc", + Cloud: "gcp", + Location: "bcd", + GcpManagedNetworkConfig: &provisioning.GcpManagedNetworkConfig{ + SubnetCidr: "a", + }, + } + + // Create a mock waiter + mockWaiter := &provisioning.WaitGetWorkspaceRunning[provisioning.Workspace]{ + WorkspaceId: 1234, + Poll: func(d time.Duration, f func(*provisioning.Workspace)) (*provisioning.Workspace, error) { + return mockWorkspace, nil + }, + } + qa.ResourceFixture{ - Fixtures: []qa.HTTPFixture{ - { - Method: "POST", - Resource: "/api/2.0/accounts/abc/workspaces", - // retreating to raw JSON, as certain fields don't work well together - ExpectedRequest: map[string]any{ - "account_id": "abc", - "cloud": "gcp", - "cloud_resource_container": map[string]any{ - "gcp": map[string]any{ - "project_id": "def", - }, + MockAccountClientFunc: func(a *mocks.MockAccountClient) { + a.GetMockWorkspacesAPI().EXPECT().Create(mock.Anything, provisioning.CreateWorkspaceRequest{ + Cloud: "gcp", + CloudResourceContainer: &provisioning.CloudResourceContainer{ + Gcp: &provisioning.CustomerFacingGcpCloudResourceContainer{ + ProjectId: "def", }, - "location": "bcd", - "network_id": "net_id_a", - "gcp_managed_network_config": map[string]any{ - "subnet_cidr": "a", - }, - "workspace_name": "labdata", - }, - Response: Workspace{ - WorkspaceID: 1234, - AccountID: "abc", - DeploymentName: "900150983cd24fb0", - WorkspaceName: "labdata", }, - }, - { - Method: "GET", - ReuseRequest: true, - Resource: "/api/2.0/accounts/abc/workspaces/1234", - Response: Workspace{ - AccountID: "abc", - WorkspaceID: 1234, - WorkspaceStatus: WorkspaceStatusRunning, - DeploymentName: "900150983cd24fb0", - WorkspaceName: "labdata", - Location: "bcd", - Cloud: "gcp", + Location: "bcd", + NetworkId: "net_id_a", + GcpManagedNetworkConfig: &provisioning.GcpManagedNetworkConfig{ + SubnetCidr: "a", }, - }, + WorkspaceName: "labdata", + }).Return(mockWaiter, nil) + a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ + WorkspaceId: 1234, + }).Return(mockWorkspace, nil) }, Resource: ResourceMwsWorkspaces(), HCL: ` @@ -149,55 +149,33 @@ func TestResourceWorkspaceCreateGcp(t *testing.T) { `, Gcp: true, Create: true, - }.ApplyAndExpectData(t, map[string]any{ - "cloud": "gcp", - "gcp_workspace_sa": "db-1234@prod-gcp-bcd.iam.gserviceaccount.com", - }) + }.ApplyNoError(t) } func TestResourceWorkspaceCreate_Error_Custom_tags(t *testing.T) { qa.ResourceFixture{ - Fixtures: []qa.HTTPFixture{ - { - Method: "POST", - Resource: "/api/2.0/accounts/abc/workspaces", - // retreating to raw JSON, as certain fields don't work well together - ExpectedRequest: map[string]any{ - "account_id": "abc", - "cloud": "gcp", - "cloud_resource_container": map[string]any{ - "gcp": map[string]any{ - "project_id": "def", - }, - }, - "location": "bcd", - "private_access_settings_id": "pas_id_a", - "network_id": "net_id_a", - "gcp_managed_network_config": map[string]any{ - "subnet_cidr": "a", - }, - "workspace_name": "labdata", - "custom_tags": map[string]any{ - "SoldToCode": "1234", + MockAccountClientFunc: func(a *mocks.MockAccountClient) { + a.GetMockWorkspacesAPI().EXPECT().Create(mock.Anything, provisioning.CreateWorkspaceRequest{ + Cloud: "gcp", + CloudResourceContainer: &provisioning.CloudResourceContainer{ + Gcp: &provisioning.CustomerFacingGcpCloudResourceContainer{ + ProjectId: "def", }, }, - Response: common.APIErrorBody{ - ErrorCode: "INVALID_PARAMETER_VALUE", - Message: "custom_tags are only allowed for AWS workspaces", + Location: "bcd", + PrivateAccessSettingsId: "pas_id_a", + NetworkId: "net_id_a", + GcpManagedNetworkConfig: &provisioning.GcpManagedNetworkConfig{ + SubnetCidr: "a", }, - }, - { - Method: "GET", - ReuseRequest: true, - Resource: "/api/2.0/accounts/abc/workspaces/1234", - Response: Workspace{ - AccountID: "abc", - WorkspaceID: 1234, - WorkspaceStatus: WorkspaceStatusRunning, - DeploymentName: "900150983cd24fb0", - WorkspaceName: "labdata", + WorkspaceName: "labdata", + CustomTags: map[string]string{ + "SoldToCode": "1234", }, - }, + }).Return(nil, &apierr.APIError{ + ErrorCode: "INVALID_PARAMETER_VALUE", + Message: "custom_tags are only allowed for AWS workspaces", + }) }, Resource: ResourceMwsWorkspaces(), HCL: ` @@ -225,47 +203,49 @@ func TestResourceWorkspaceCreate_Error_Custom_tags(t *testing.T) { } func TestResourceWorkspaceCreateGcpPsc(t *testing.T) { + // Define a mock workspace that can be reused + mockWorkspace := &provisioning.Workspace{ + WorkspaceId: 1234, + WorkspaceStatus: provisioning.WorkspaceStatusRunning, + WorkspaceName: "labdata", + DeploymentName: "900150983cd24fb0", + AccountId: "abc", + Cloud: "gcp", + Location: "bcd", + PrivateAccessSettingsId: "pas_id_a", + GcpManagedNetworkConfig: &provisioning.GcpManagedNetworkConfig{ + SubnetCidr: "a", + }, + } + + // Create a mock waiter + mockWaiter := &provisioning.WaitGetWorkspaceRunning[provisioning.Workspace]{ + WorkspaceId: 1234, + Poll: func(d time.Duration, f func(*provisioning.Workspace)) (*provisioning.Workspace, error) { + return mockWorkspace, nil + }, + } + qa.ResourceFixture{ - Fixtures: []qa.HTTPFixture{ - { - Method: "POST", - Resource: "/api/2.0/accounts/abc/workspaces", - // retreating to raw JSON, as certain fields don't work well together - ExpectedRequest: map[string]any{ - "account_id": "abc", - "cloud": "gcp", - "cloud_resource_container": map[string]any{ - "gcp": map[string]any{ - "project_id": "def", - }, + MockAccountClientFunc: func(a *mocks.MockAccountClient) { + a.GetMockWorkspacesAPI().EXPECT().Create(mock.Anything, provisioning.CreateWorkspaceRequest{ + Cloud: "gcp", + CloudResourceContainer: &provisioning.CloudResourceContainer{ + Gcp: &provisioning.CustomerFacingGcpCloudResourceContainer{ + ProjectId: "def", }, - "location": "bcd", - "private_access_settings_id": "pas_id_a", - "network_id": "net_id_a", - "gcp_managed_network_config": map[string]any{ - "subnet_cidr": "a", - }, - "workspace_name": "labdata", - }, - Response: Workspace{ - WorkspaceID: 1234, - AccountID: "abc", - DeploymentName: "900150983cd24fb0", - WorkspaceName: "labdata", }, - }, - { - Method: "GET", - ReuseRequest: true, - Resource: "/api/2.0/accounts/abc/workspaces/1234", - Response: Workspace{ - AccountID: "abc", - WorkspaceID: 1234, - WorkspaceStatus: WorkspaceStatusRunning, - DeploymentName: "900150983cd24fb0", - WorkspaceName: "labdata", + Location: "bcd", + PrivateAccessSettingsId: "pas_id_a", + NetworkId: "net_id_a", + GcpManagedNetworkConfig: &provisioning.GcpManagedNetworkConfig{ + SubnetCidr: "a", }, - }, + WorkspaceName: "labdata", + }).Return(mockWaiter, nil) + a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ + WorkspaceId: 1234, + }).Return(mockWorkspace, nil) }, Resource: ResourceMwsWorkspaces(), HCL: ` @@ -290,48 +270,52 @@ func TestResourceWorkspaceCreateGcpPsc(t *testing.T) { } func TestResourceWorkspaceCreateGcpCmk(t *testing.T) { + // Define a mock workspace that can be reused + mockWorkspace := &provisioning.Workspace{ + WorkspaceId: 1234, + WorkspaceStatus: provisioning.WorkspaceStatusRunning, + WorkspaceName: "labdata", + DeploymentName: "900150983cd24fb0", + Cloud: "gcp", + Location: "bcd", + PrivateAccessSettingsId: "pas_id_a", + GcpManagedNetworkConfig: &provisioning.GcpManagedNetworkConfig{ + SubnetCidr: "a", + }, + ManagedServicesCustomerManagedKeyId: "managed_services_cmk", + StorageCustomerManagedKeyId: "storage_cmk", + } + + // Create a mock waiter + mockWaiter := &provisioning.WaitGetWorkspaceRunning[provisioning.Workspace]{ + WorkspaceId: 1234, + Poll: func(d time.Duration, f func(*provisioning.Workspace)) (*provisioning.Workspace, error) { + return mockWorkspace, nil + }, + } + qa.ResourceFixture{ - Fixtures: []qa.HTTPFixture{ - { - Method: "POST", - Resource: "/api/2.0/accounts/abc/workspaces", - ExpectedRequest: map[string]any{ - "account_id": "abc", - "cloud": "gcp", - "cloud_resource_container": map[string]any{ - "gcp": map[string]any{ - "project_id": "def", - }, - }, - "location": "bcd", - "private_access_settings_id": "pas_id_a", - "network_id": "net_id_a", - "gcp_managed_network_config": map[string]any{ - "subnet_cidr": "a", + MockAccountClientFunc: func(a *mocks.MockAccountClient) { + a.GetMockWorkspacesAPI().EXPECT().Create(mock.Anything, provisioning.CreateWorkspaceRequest{ + Cloud: "gcp", + CloudResourceContainer: &provisioning.CloudResourceContainer{ + Gcp: &provisioning.CustomerFacingGcpCloudResourceContainer{ + ProjectId: "def", }, - "workspace_name": "labdata", - "managed_services_customer_managed_key_id": "managed_services_cmk", - "storage_customer_managed_key_id": "storage_cmk", }, - Response: Workspace{ - WorkspaceID: 1234, - AccountID: "abc", - DeploymentName: "900150983cd24fb0", - WorkspaceName: "labdata", + Location: "bcd", + PrivateAccessSettingsId: "pas_id_a", + NetworkId: "net_id_a", + GcpManagedNetworkConfig: &provisioning.GcpManagedNetworkConfig{ + SubnetCidr: "a", }, - }, - { - Method: "GET", - ReuseRequest: true, - Resource: "/api/2.0/accounts/abc/workspaces/1234", - Response: Workspace{ - AccountID: "abc", - WorkspaceID: 1234, - WorkspaceStatus: WorkspaceStatusRunning, - DeploymentName: "900150983cd24fb0", - WorkspaceName: "labdata", - }, - }, + WorkspaceName: "labdata", + ManagedServicesCustomerManagedKeyId: "managed_services_cmk", + StorageCustomerManagedKeyId: "storage_cmk", + }).Return(mockWaiter, nil) + a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ + WorkspaceId: 1234, + }).Return(mockWorkspace, nil) }, Resource: ResourceMwsWorkspaces(), HCL: ` @@ -358,47 +342,45 @@ func TestResourceWorkspaceCreateGcpCmk(t *testing.T) { } func TestResourceWorkspaceCreateWithIsNoPublicIPEnabledFalse(t *testing.T) { + // Define a mock workspace that can be reused + mockWorkspace := &provisioning.Workspace{ + WorkspaceId: 1234, + WorkspaceStatus: provisioning.WorkspaceStatusRunning, + WorkspaceName: "labdata", + DeploymentName: "900150983cd24fb0", + AwsRegion: "us-east-1", + CredentialsId: "bcd", + StorageConfigurationId: "ghi", + NetworkId: "fgh", + ManagedServicesCustomerManagedKeyId: "def", + StorageCustomerManagedKeyId: "def", + AccountId: "abc", + } + + // Create a mock waiter + mockWaiter := &provisioning.WaitGetWorkspaceRunning[provisioning.Workspace]{ + WorkspaceId: 1234, + Poll: func(d time.Duration, f func(*provisioning.Workspace)) (*provisioning.Workspace, error) { + return mockWorkspace, nil + }, + } + d, err := qa.ResourceFixture{ - Fixtures: []qa.HTTPFixture{ - { - Method: "POST", - Resource: "/api/2.0/accounts/abc/workspaces", - ExpectedRequest: Workspace{ - AccountID: "abc", - IsNoPublicIPEnabled: false, - WorkspaceName: "labdata", - DeploymentName: "900150983cd24fb0", - AwsRegion: "us-east-1", - CredentialsID: "bcd", - StorageConfigurationID: "ghi", - NetworkID: "fgh", - ManagedServicesCustomerManagedKeyID: "def", - StorageCustomerManagedKeyID: "def", - }, - Response: Workspace{ - WorkspaceID: 1234, - AccountID: "abc", - DeploymentName: "900150983cd24fb0", - }, - }, - { - Method: "GET", - ReuseRequest: true, - Resource: "/api/2.0/accounts/abc/workspaces/1234", - Response: Workspace{ - WorkspaceID: 1234, - WorkspaceStatus: WorkspaceStatusRunning, - WorkspaceName: "labdata", - DeploymentName: "900150983cd24fb0", - AwsRegion: "us-east-1", - CredentialsID: "bcd", - StorageConfigurationID: "ghi", - NetworkID: "fgh", - ManagedServicesCustomerManagedKeyID: "def", - StorageCustomerManagedKeyID: "def", - AccountID: "abc", - }, - }, + MockAccountClientFunc: func(a *mocks.MockAccountClient) { + a.GetMockWorkspacesAPI().EXPECT().Create(mock.Anything, provisioning.CreateWorkspaceRequest{ + IsNoPublicIpEnabled: false, + WorkspaceName: "labdata", + DeploymentName: "900150983cd24fb0", + AwsRegion: "us-east-1", + CredentialsId: "bcd", + StorageConfigurationId: "ghi", + NetworkId: "fgh", + ManagedServicesCustomerManagedKeyId: "def", + StorageCustomerManagedKeyId: "def", + }).Return(mockWaiter, nil) + a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ + WorkspaceId: 1234, + }).Return(mockWorkspace, nil) }, Resource: ResourceMwsWorkspaces(), State: map[string]any{ @@ -420,45 +402,43 @@ func TestResourceWorkspaceCreateWithIsNoPublicIPEnabledFalse(t *testing.T) { } func TestResourceWorkspaceCreateLegacyConfig(t *testing.T) { + // Define a mock workspace that can be reused + mockWorkspace := &provisioning.Workspace{ + WorkspaceId: 1234, + WorkspaceStatus: provisioning.WorkspaceStatusRunning, + WorkspaceName: "labdata", + DeploymentName: "900150983cd24fb0", + AwsRegion: "us-east-1", + CredentialsId: "bcd", + StorageConfigurationId: "ghi", + NetworkId: "fgh", + ManagedServicesCustomerManagedKeyId: "def", + AccountId: "abc", + } + + // Create a mock waiter + mockWaiter := &provisioning.WaitGetWorkspaceRunning[provisioning.Workspace]{ + WorkspaceId: 1234, + Poll: func(d time.Duration, f func(*provisioning.Workspace)) (*provisioning.Workspace, error) { + return mockWorkspace, nil + }, + } + d, err := qa.ResourceFixture{ - Fixtures: []qa.HTTPFixture{ - { - Method: "POST", - Resource: "/api/2.0/accounts/abc/workspaces", - ExpectedRequest: Workspace{ - AccountID: "abc", - IsNoPublicIPEnabled: true, - WorkspaceName: "labdata", - DeploymentName: "900150983cd24fb0", - AwsRegion: "us-east-1", - CredentialsID: "bcd", - StorageConfigurationID: "ghi", - NetworkID: "fgh", - ManagedServicesCustomerManagedKeyID: "def", - }, - Response: Workspace{ - WorkspaceID: 1234, - AccountID: "abc", - DeploymentName: "900150983cd24fb0", - }, - }, - { - Method: "GET", - ReuseRequest: true, - Resource: "/api/2.0/accounts/abc/workspaces/1234", - Response: Workspace{ - WorkspaceID: 1234, - WorkspaceStatus: WorkspaceStatusRunning, - WorkspaceName: "labdata", - DeploymentName: "900150983cd24fb0", - AwsRegion: "us-east-1", - CredentialsID: "bcd", - StorageConfigurationID: "ghi", - NetworkID: "fgh", - ManagedServicesCustomerManagedKeyID: "def", - AccountID: "abc", - }, - }, + MockAccountClientFunc: func(a *mocks.MockAccountClient) { + a.GetMockWorkspacesAPI().EXPECT().Create(mock.Anything, provisioning.CreateWorkspaceRequest{ + IsNoPublicIpEnabled: true, + WorkspaceName: "labdata", + DeploymentName: "900150983cd24fb0", + AwsRegion: "us-east-1", + CredentialsId: "bcd", + StorageConfigurationId: "ghi", + NetworkId: "fgh", + ManagedServicesCustomerManagedKeyId: "def", + }).Return(mockWaiter, nil) + a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ + WorkspaceId: 1234, + }).Return(mockWorkspace, nil) }, Resource: ResourceMwsWorkspaces(), State: map[string]any{ @@ -477,69 +457,27 @@ func TestResourceWorkspaceCreateLegacyConfig(t *testing.T) { assert.Equal(t, "abc/1234", d.Id()) } -func TestResourceWorkspaceCreate_Error(t *testing.T) { - t.Skipf("Making this test skip until we can configure sleep timings for test purposes") - d, err := qa.ResourceFixture{ - Fixtures: []qa.HTTPFixture{ - { - Method: "POST", - Resource: "/api/2.0/accounts/abc/workspaces", - Response: common.APIErrorBody{ - ErrorCode: "INVALID_REQUEST", - Message: "Internal error happened", - }, - Status: 400, - }, - { - Method: "POST", - Resource: "/api/2.0/accounts/abc/workspaces", - Response: common.APIErrorBody{ - ErrorCode: "INVALID_REQUEST", - Message: "Internal error happened", - }, - Status: 400, - }, - }, - Resource: ResourceMwsWorkspaces(), - State: map[string]any{ - "account_id": "abc", - "aws_region": "us-east-1", - "credentials_id": "bcd", - "managed_services_customer_managed_key_id": "def", - "storage_customer_managed_key_id": "def", - "deployment_name": "900150983cd24fb0", - "workspace_name": "labdata", - "is_no_public_ip_enabled": true, - "network_id": "fgh", - "storage_configuration_id": "ghi", - }, - Create: true, - }.Apply(t) - qa.AssertErrorStartsWith(t, err, "Internal error happened") - assert.Equal(t, "", d.Id(), "Id should be empty for error creates") -} - func TestResourceWorkspaceRead(t *testing.T) { + // Define a mock workspace that can be reused + mockWorkspace := &provisioning.Workspace{ + WorkspaceId: 1234, + WorkspaceStatus: provisioning.WorkspaceStatusRunning, + WorkspaceName: "labdata", + DeploymentName: "900150983cd24fb0", + AwsRegion: "us-east-1", + CredentialsId: "bcd", + StorageConfigurationId: "ghi", + NetworkId: "fgh", + ManagedServicesCustomerManagedKeyId: "def", + StorageCustomerManagedKeyId: "def", + AccountId: "abc", + } + d, err := qa.ResourceFixture{ - Fixtures: []qa.HTTPFixture{ - { - Method: "GET", - ReuseRequest: true, - Resource: "/api/2.0/accounts/abc/workspaces/1234", - Response: Workspace{ - AccountID: "abc", - WorkspaceStatus: WorkspaceStatusRunning, - WorkspaceName: "labdata", - DeploymentName: "900150983cd24fb0", - AwsRegion: "us-east-1", - CredentialsID: "bcd", - StorageConfigurationID: "ghi", - NetworkID: "fgh", - ManagedServicesCustomerManagedKeyID: "def", - StorageCustomerManagedKeyID: "def", - WorkspaceID: 1234, - }, - }, + MockAccountClientFunc: func(a *mocks.MockAccountClient) { + a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ + WorkspaceId: 1234, + }).Return(mockWorkspace, nil) }, Resource: ResourceMwsWorkspaces(), Read: true, @@ -563,26 +501,26 @@ func TestResourceWorkspaceRead(t *testing.T) { } func TestResourceWorkspaceRead_Issue382(t *testing.T) { + // Define a mock workspace that can be reused + mockWorkspace := &provisioning.Workspace{ + WorkspaceId: 1234, + WorkspaceStatus: provisioning.WorkspaceStatusRunning, + WorkspaceName: "labdata", + DeploymentName: "prefix-900150983cd24fb0", + AwsRegion: "us-east-1", + CredentialsId: "bcd", + StorageConfigurationId: "ghi", + NetworkId: "fgh", + ManagedServicesCustomerManagedKeyId: "def", + StorageCustomerManagedKeyId: "def", + AccountId: "abc", + } + d, err := qa.ResourceFixture{ - Fixtures: []qa.HTTPFixture{ - { - Method: "GET", - ReuseRequest: true, - Resource: "/api/2.0/accounts/abc/workspaces/1234", - Response: Workspace{ - AccountID: "abc", - WorkspaceStatus: WorkspaceStatusRunning, - WorkspaceName: "labdata", - DeploymentName: "prefix-900150983cd24fb0", - AwsRegion: "us-east-1", - CredentialsID: "bcd", - StorageConfigurationID: "ghi", - NetworkID: "fgh", - ManagedServicesCustomerManagedKeyID: "def", - StorageCustomerManagedKeyID: "def", - WorkspaceID: 1234, - }, - }, + MockAccountClientFunc: func(a *mocks.MockAccountClient) { + a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ + WorkspaceId: 1234, + }).Return(mockWorkspace, nil) }, InstanceState: map[string]string{ "account_id": "abc", @@ -619,16 +557,14 @@ func TestResourceWorkspaceRead_Issue382(t *testing.T) { func TestResourceWorkspaceRead_NotFound(t *testing.T) { qa.ResourceFixture{ - Fixtures: []qa.HTTPFixture{ - { - Method: "GET", - Resource: "/api/2.0/accounts/abc/workspaces/1234", - Response: common.APIErrorBody{ - ErrorCode: "NOT_FOUND", - Message: "Item not found", - }, - Status: 404, - }, + MockAccountClientFunc: func(a *mocks.MockAccountClient) { + a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ + WorkspaceId: 1234, + }).Return(nil, &apierr.APIError{ + ErrorCode: "NOT_FOUND", + Message: "Item not found", + StatusCode: 404, + }) }, Resource: ResourceMwsWorkspaces(), Read: true, @@ -639,16 +575,14 @@ func TestResourceWorkspaceRead_NotFound(t *testing.T) { func TestResourceWorkspaceRead_Error(t *testing.T) { d, err := qa.ResourceFixture{ - Fixtures: []qa.HTTPFixture{ - { // read log output for correct url... - Method: "GET", - Resource: "/api/2.0/accounts/abc/workspaces/1234", - Response: common.APIErrorBody{ - ErrorCode: "INVALID_REQUEST", - Message: "Internal error happened", - }, - Status: 400, - }, + MockAccountClientFunc: func(a *mocks.MockAccountClient) { + a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ + WorkspaceId: 1234, + }).Return(nil, &apierr.APIError{ + ErrorCode: "INVALID_REQUEST", + Message: "Internal error happened", + StatusCode: 400, + }) }, Resource: ResourceMwsWorkspaces(), Read: true, @@ -659,35 +593,40 @@ func TestResourceWorkspaceRead_Error(t *testing.T) { } func TestResourceWorkspaceUpdate(t *testing.T) { + // Define a mock workspace that can be reused + mockWorkspace := &provisioning.Workspace{ + WorkspaceId: 1234, + WorkspaceStatus: provisioning.WorkspaceStatusRunning, + WorkspaceName: "labdata", + DeploymentName: "900150983cd24fb0", + AwsRegion: "us-east-1", + CredentialsId: "bcd", + StorageConfigurationId: "ghi", + NetworkId: "fgh", + ManagedServicesCustomerManagedKeyId: "def", + StorageCustomerManagedKeyId: "def", + AccountId: "abc", + } + + // Create a mock waiter + mockWaiter := &provisioning.WaitGetWorkspaceRunning[struct{}]{ + WorkspaceId: 1234, + Poll: func(d time.Duration, f func(*provisioning.Workspace)) (*provisioning.Workspace, error) { + return mockWorkspace, nil + }, + } + d, err := qa.ResourceFixture{ - Fixtures: []qa.HTTPFixture{ - { - Method: "PATCH", - Resource: "/api/2.0/accounts/abc/workspaces/1234", - ExpectedRequest: map[string]any{ - "credentials_id": "bcd", - "network_id": "fgh", - "storage_customer_managed_key_id": "def", - }, - }, - { - Method: "GET", - ReuseRequest: true, - Resource: "/api/2.0/accounts/abc/workspaces/1234", - Response: Workspace{ - WorkspaceStatus: WorkspaceStatusRunning, - WorkspaceName: "labdata", - DeploymentName: "900150983cd24fb0", - AwsRegion: "us-east-1", - CredentialsID: "bcd", - StorageConfigurationID: "ghi", - NetworkID: "fgh", - ManagedServicesCustomerManagedKeyID: "def", - StorageCustomerManagedKeyID: "def", - AccountID: "abc", - WorkspaceID: 1234, - }, - }, + MockAccountClientFunc: func(a *mocks.MockAccountClient) { + a.GetMockWorkspacesAPI().EXPECT().Update(mock.Anything, provisioning.UpdateWorkspaceRequest{ + WorkspaceId: 1234, + CredentialsId: "bcd", + NetworkId: "fgh", + StorageCustomerManagedKeyId: "def", + }).Return(mockWaiter, nil) + a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ + WorkspaceId: 1234, + }).Return(mockWorkspace, nil) }, Resource: ResourceMwsWorkspaces(), InstanceState: map[string]string{ @@ -759,34 +698,39 @@ func TestResourceWorkspaceUpdate_NotAllowed(t *testing.T) { } func TestResourceWorkspaceUpdateLegacyConfig(t *testing.T) { + // Define a mock workspace that can be reused + mockWorkspace := &provisioning.Workspace{ + WorkspaceId: 1234, + WorkspaceStatus: provisioning.WorkspaceStatusRunning, + IsNoPublicIpEnabled: true, + WorkspaceName: "labdata", + DeploymentName: "900150983cd24fb0", + AwsRegion: "us-east-1", + CredentialsId: "bcd", + StorageConfigurationId: "ghi", + NetworkId: "fgh", + ManagedServicesCustomerManagedKeyId: "def", + AccountId: "abc", + } + + // Create a mock waiter + mockWaiter := &provisioning.WaitGetWorkspaceRunning[struct{}]{ + WorkspaceId: 1234, + Poll: func(d time.Duration, f func(*provisioning.Workspace)) (*provisioning.Workspace, error) { + return mockWorkspace, nil + }, + } + d, err := qa.ResourceFixture{ - Fixtures: []qa.HTTPFixture{ - { - Method: "PATCH", - Resource: "/api/2.0/accounts/abc/workspaces/1234", - ExpectedRequest: map[string]any{ - "credentials_id": "bcd", - "network_id": "fgh", - }, - }, - { - Method: "GET", - ReuseRequest: true, - Resource: "/api/2.0/accounts/abc/workspaces/1234", - Response: Workspace{ - WorkspaceStatus: WorkspaceStatusRunning, - IsNoPublicIPEnabled: true, - WorkspaceName: "labdata", - DeploymentName: "900150983cd24fb0", - AwsRegion: "us-east-1", - CredentialsID: "bcd", - StorageConfigurationID: "ghi", - NetworkID: "fgh", - ManagedServicesCustomerManagedKeyID: "def", - AccountID: "abc", - WorkspaceID: 1234, - }, - }, + MockAccountClientFunc: func(a *mocks.MockAccountClient) { + a.GetMockWorkspacesAPI().EXPECT().Update(mock.Anything, provisioning.UpdateWorkspaceRequest{ + WorkspaceId: 1234, + CredentialsId: "bcd", + NetworkId: "fgh", + }).Return(mockWaiter, nil) + a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ + WorkspaceId: 1234, + }).Return(mockWorkspace, nil) }, Resource: ResourceMwsWorkspaces(), InstanceState: map[string]string{ @@ -821,16 +765,17 @@ func TestResourceWorkspaceUpdateLegacyConfig(t *testing.T) { func TestResourceWorkspaceUpdate_Error(t *testing.T) { qa.ResourceFixture{ - Fixtures: []qa.HTTPFixture{ - { - Method: "PATCH", - Resource: "/api/2.0/accounts/abc/workspaces/1234", - Response: common.APIErrorBody{ - ErrorCode: "INVALID_REQUEST", - Message: "Internal error happened", - }, - Status: 400, - }, + MockAccountClientFunc: func(a *mocks.MockAccountClient) { + a.GetMockWorkspacesAPI().EXPECT().Update(mock.Anything, provisioning.UpdateWorkspaceRequest{ + WorkspaceId: 1234, + CredentialsId: "bcd", + NetworkId: "fgh", + StorageCustomerManagedKeyId: "def", + }).Return(nil, &apierr.APIError{ + ErrorCode: "INVALID_REQUEST", + Message: "Internal error happened", + StatusCode: 400, + }) }, Resource: ResourceMwsWorkspaces(), State: map[string]any{ @@ -852,30 +797,28 @@ func TestResourceWorkspaceUpdate_Error(t *testing.T) { } func TestResourceWorkspaceDelete(t *testing.T) { + // Define a mock workspace that can be reused for the first GET call + mockWorkspace := &provisioning.Workspace{ + WorkspaceName: "labdata", + WorkspaceStatus: provisioning.WorkspaceStatusCancelling, + WorkspaceStatusMessage: "Things are being removed", + } + d, err := qa.ResourceFixture{ - Fixtures: []qa.HTTPFixture{ - { - Method: "DELETE", - Resource: "/api/2.0/accounts/abc/workspaces/1234", - }, - { - Method: "GET", - Resource: "/api/2.0/accounts/abc/workspaces/1234", - Response: Workspace{ - WorkspaceName: "labdata", - WorkspaceStatus: WorkspaceStatusCanceled, - WorkspaceStatusMessage: "Things are being removed", - }, - }, - { - Method: "GET", - Resource: "/api/2.0/accounts/abc/workspaces/1234", - Response: common.APIErrorBody{ - ErrorCode: "NOT_FOUND", - Message: "Cannot find anything", - }, - Status: 404, - }, + MockAccountClientFunc: func(a *mocks.MockAccountClient) { + a.GetMockWorkspacesAPI().EXPECT().Delete(mock.Anything, provisioning.DeleteWorkspaceRequest{ + WorkspaceId: 1234, + }).Return(nil) + a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ + WorkspaceId: 1234, + }).Return(mockWorkspace, nil).Once() + a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ + WorkspaceId: 1234, + }).Return(nil, &apierr.APIError{ + ErrorCode: "NOT_FOUND", + Message: "Cannot find anything", + StatusCode: 404, + }) }, Resource: ResourceMwsWorkspaces(), Delete: true, @@ -885,693 +828,247 @@ func TestResourceWorkspaceDelete(t *testing.T) { assert.Equal(t, "abc/1234", d.Id()) } -func TestResourceWorkspaceDelete_Error(t *testing.T) { - d, err := qa.ResourceFixture{ - Fixtures: []qa.HTTPFixture{ - { - Method: "DELETE", - Resource: "/api/2.0/accounts/abc/workspaces/1234", - Response: common.APIErrorBody{ - ErrorCode: "INVALID_REQUEST", - Message: "Internal error happened", - }, - Status: 400, - }, - }, - Resource: ResourceMwsWorkspaces(), - Delete: true, - ID: "abc/1234", - }.Apply(t) - qa.AssertErrorStartsWith(t, err, "Internal error happened") - assert.Equal(t, "abc/1234", d.Id()) -} +func TestCreateFailsAndCleansUp(t *testing.T) { + // Define a mock workspace that represents the failed state + mockFailedWorkspace := &provisioning.Workspace{ + WorkspaceId: 1234, + WorkspaceStatus: provisioning.WorkspaceStatusFailed, + WorkspaceStatusMessage: "Always fails", + WorkspaceName: "labdata", + DeploymentName: "900150983cd24fb0", + AwsRegion: "us-east-1", + NetworkId: "fgh", + AccountId: "abc", + } -func TestWaitForRunning(t *testing.T) { - client, server, err := qa.HttpFixtureClient(t, []qa.HTTPFixture{ - { - Method: "POST", - Resource: "/api/2.0/accounts/abc/workspaces", - ExpectedRequest: Workspace{ - AccountID: "abc", - IsNoPublicIPEnabled: true, - WorkspaceName: "labdata", - DeploymentName: "900150983cd24fb0", - AwsRegion: "us-east-1", - CredentialsID: "bcd", - StorageConfigurationID: "ghi", - NetworkID: "fgh", - ManagedServicesCustomerManagedKeyID: "def", - StorageCustomerManagedKeyID: "def", - }, - Response: Workspace{ - WorkspaceID: 1234, - AccountID: "abc", - DeploymentName: "900150983cd24fb0", - }, - }, - { - Method: "GET", - Resource: "/api/2.0/accounts/abc/workspaces/1234", - Response: Workspace{ - WorkspaceID: 1234, - WorkspaceStatus: WorkspaceStatusProvisioning, - WorkspaceName: "labdata", - DeploymentName: "900150983cd24fb0", - AwsRegion: "us-east-1", - CredentialsID: "bcd", - StorageConfigurationID: "ghi", - NetworkID: "fgh", - ManagedServicesCustomerManagedKeyID: "def", - StorageCustomerManagedKeyID: "def", - AccountID: "abc", - }, + // Create a mock waiter + mockWaiter := &provisioning.WaitGetWorkspaceRunning[provisioning.Workspace]{ + WorkspaceId: 1234, + Poll: func(d time.Duration, f func(*provisioning.Workspace)) (*provisioning.Workspace, error) { + return mockFailedWorkspace, nil }, - { - Method: "GET", - Resource: "/api/2.0/accounts/abc/workspaces/1234", - Response: Workspace{ - WorkspaceID: 1234, - WorkspaceStatus: WorkspaceStatusRunning, - WorkspaceName: "labdata", - DeploymentName: "900150983cd24fb0", - AwsRegion: "us-east-1", - CredentialsID: "bcd", - StorageConfigurationID: "ghi", - NetworkID: "fgh", - ManagedServicesCustomerManagedKeyID: "def", - StorageCustomerManagedKeyID: "def", - AccountID: "abc", + } + + // Define a mock network with error messages + mockNetwork := &provisioning.Network{ + NetworkId: "fgh", + ErrorMessages: []provisioning.NetworkHealth{ + { + ErrorType: provisioning.ErrorTypeCredentials, + ErrorMessage: "Message", }, }, - }) - require.NoError(t, err) - defer server.Close() - - err = NewWorkspacesAPI(context.Background(), client).Create(&Workspace{ - AccountID: "abc", - IsNoPublicIPEnabled: true, - WorkspaceName: "labdata", - DeploymentName: "900150983cd24fb0", - AwsRegion: "us-east-1", - CredentialsID: "bcd", - StorageConfigurationID: "ghi", - NetworkID: "fgh", - ManagedServicesCustomerManagedKeyID: "def", - StorageCustomerManagedKeyID: "def", - }, DefaultProvisionTimeout) - require.NoError(t, err) -} + } -func TestCreateFailsAndCleansUp(t *testing.T) { - client, server, err := qa.HttpFixtureClient(t, []qa.HTTPFixture{ - { - Method: "POST", - Resource: "/api/2.0/accounts/abc/workspaces", - ExpectedRequest: Workspace{ - AccountID: "abc", - IsNoPublicIPEnabled: true, + _, err := qa.ResourceFixture{ + MockAccountClientFunc: func(a *mocks.MockAccountClient) { + // Expect the Create call + a.GetMockWorkspacesAPI().EXPECT().Create(mock.Anything, provisioning.CreateWorkspaceRequest{ + IsNoPublicIpEnabled: true, WorkspaceName: "labdata", DeploymentName: "900150983cd24fb0", AwsRegion: "us-east-1", - CredentialsID: "bcd", - StorageConfigurationID: "ghi", - NetworkID: "fgh", - ManagedServicesCustomerManagedKeyID: "def", - StorageCustomerManagedKeyID: "def", - }, - Response: Workspace{ - WorkspaceID: 1234, - AccountID: "abc", - DeploymentName: "900150983cd24fb0", - }, - }, - { - Method: "GET", - Resource: "/api/2.0/accounts/abc/workspaces/1234", - Response: Workspace{ - WorkspaceID: 1234, - WorkspaceStatus: WorkspaceStatusFailed, - WorkspaceStatusMessage: "Always fails", - WorkspaceName: "labdata", - DeploymentName: "900150983cd24fb0", - AwsRegion: "us-east-1", - NetworkID: "fgh", - AccountID: "abc", - }, - }, - { - Method: "GET", - Resource: "/api/2.0/accounts/abc/networks/fgh", - Response: Network{ - ErrorMessages: []NetworkHealth{ - {"FAIL", "Message"}, - }, - }, - }, - { - Method: "DELETE", - Resource: "/api/2.0/accounts/abc/workspaces/1234", - }, - { - Method: "GET", - Resource: "/api/2.0/accounts/abc/workspaces/1234", - Status: 404, - }, - }) - require.NoError(t, err) - defer server.Close() + CredentialsId: "bcd", + StorageConfigurationId: "ghi", + NetworkId: "fgh", + ManagedServicesCustomerManagedKeyId: "def", + StorageCustomerManagedKeyId: "def", + }).Return(mockWaiter, nil) - err = NewWorkspacesAPI(context.Background(), client).Create(&Workspace{ - AccountID: "abc", - IsNoPublicIPEnabled: true, - WorkspaceName: "labdata", - DeploymentName: "900150983cd24fb0", - AwsRegion: "us-east-1", - CredentialsID: "bcd", - StorageConfigurationID: "ghi", - NetworkID: "fgh", - ManagedServicesCustomerManagedKeyID: "def", - StorageCustomerManagedKeyID: "def", - }, DefaultProvisionTimeout) - require.EqualError(t, err, "Workspace failed to create: Always fails, network error message: error: FAIL;error_msg: Message;") -} + // Expect the Get call to retrieve the failed workspace + a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ + WorkspaceId: 1234, + }).Return(mockFailedWorkspace, nil) -func TestListWorkspaces(t *testing.T) { - client, server, err := qa.HttpFixtureClient(t, []qa.HTTPFixture{ - { - Method: "GET", - Resource: "/api/2.0/accounts/abc/workspaces", - Response: []Workspace{}, - }, - }) - require.NoError(t, err) - defer server.Close() + // Expect the Get call to retrieve the network with errors + a.GetMockNetworksAPI().EXPECT().Get(mock.Anything, provisioning.GetNetworkRequest{ + NetworkId: "fgh", + }).Return(mockNetwork, nil) - l, err := NewWorkspacesAPI(context.Background(), client).List("abc") - require.NoError(t, err) - assert.Len(t, l, 0) -} + // Expect the Delete call to clean up the failed workspace + a.GetMockWorkspacesAPI().EXPECT().Delete(mock.Anything, provisioning.DeleteWorkspaceRequest{ + WorkspaceId: 1234, + }).Return(nil) -func TestWorkspace_WaitForResolve_Failure(t *testing.T) { - qa.HTTPFixturesApply(t, []qa.HTTPFixture{}, - func(ctx context.Context, client *common.DatabricksClient) { - a := NewWorkspacesAPI(ctx, client) - rerr := a.verifyWorkspaceReachable(Workspace{ - WorkspaceURL: "https://900150983cd24fb0.cloud.databricks.com", + // Expect the final Get call to confirm the workspace is gone + a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ + WorkspaceId: 1234, + }).Return(nil, &apierr.APIError{ + ErrorCode: "NOT_FOUND", + Message: "Item not found", + StatusCode: 404, }) - assert.NotNil(t, rerr) - assert.True(t, rerr.Retryable) - }) -} - -func TestWorkspace_WaitForResolve(t *testing.T) { - // outer HTTP server is used for inner request for "just created" workspace - qa.HTTPFixturesApply(t, []qa.HTTPFixture{ - { - Method: "GET", - Resource: "/api/2.0/preview/scim/v2/Me", - Response: `{}`, // we just need a JSON for this - }, - }, func(ctx context.Context, wsClient *common.DatabricksClient) { - // inner HTTP server is used for outer request for Accounts API - qa.HTTPFixturesApply(t, []qa.HTTPFixture{ - { - Method: "GET", - Resource: "/api/2.0/accounts/abc/workspaces/1234", - ReuseRequest: true, - Response: Workspace{ - AccountID: "abc", - WorkspaceID: 1234, - WorkspaceStatus: "RUNNING", - WorkspaceURL: wsClient.Config.Host, - }, - }, - }, func(ctx context.Context, client *common.DatabricksClient) { - a := NewWorkspacesAPI(ctx, client) - err := a.WaitForRunning(Workspace{ - AccountID: "abc", - WorkspaceID: 1234, - }, 1*time.Second) - assert.NoError(t, err) - }) - }) -} - -func updateWorkspaceScimFixture(t *testing.T, fixtures []qa.HTTPFixture, state map[string]string, hcl string) { - accountsAPI := []qa.HTTPFixture{ - { - Method: "GET", - ReuseRequest: true, - Resource: "/api/2.0/accounts/c/workspaces/0", }, - } - scimAPI := []qa.HTTPFixture{ - { - Method: "GET", - Resource: "/api/2.0/preview/scim/v2/Me", - Response: `{}`, // we just need a JSON for this + Resource: ResourceMwsWorkspaces(), + State: map[string]any{ + "account_id": "abc", + "aws_region": "us-east-1", + "credentials_id": "bcd", + "managed_services_customer_managed_key_id": "def", + "storage_customer_managed_key_id": "def", + "deployment_name": "900150983cd24fb0", + "workspace_name": "labdata", + "network_id": "fgh", + "storage_configuration_id": "ghi", }, - } - scimAPI = append(scimAPI, fixtures...) - // outer HTTP server is used for inner request for "just created" workspace - qa.HTTPFixturesApply(t, scimAPI, func(ctx context.Context, wsClient *common.DatabricksClient) { - // a bit hacky, but the whole thing is more readable - accountsAPI[0].Response = Workspace{ - WorkspaceStatus: "RUNNING", - WorkspaceURL: wsClient.Config.Host, - } - state["workspace_url"] = wsClient.Config.Host - state["workspace_name"] = "b" - state["account_id"] = "c" - state["network_id"] = "d" - state["is_no_public_ip_enabled"] = "false" - qa.ResourceFixture{ - Fixtures: accountsAPI, - Resource: ResourceMwsWorkspaces(), - InstanceState: state, - Update: true, - ID: "a", - HCL: hcl + ` - workspace_name = "b" - account_id = "c", - network_id = "d"`, - }.Apply(t) - }) + Create: true, + }.Apply(t) + assert.ErrorContains(t, err, "Workspace failed to create: Always fails, network error message: error: CREDENTIALS;error_msg: Message;") } -func updateWorkspaceScimFixtureWithPatch(t *testing.T, fixtures []qa.HTTPFixture, state map[string]string, hcl string) { - accountsAPI := []qa.HTTPFixture{ - { - Method: "PATCH", - Resource: "/api/2.0/accounts/c/workspaces/0", - }, - { - Method: "GET", - ReuseRequest: true, - Resource: "/api/2.0/accounts/c/workspaces/0", - }, - } - scimAPI := []qa.HTTPFixture{ - { - Method: "GET", - Resource: "/api/2.0/preview/scim/v2/Me", - Response: `{}`, // we just need a JSON for this - }, - } - scimAPI = append(scimAPI, fixtures...) - // outer HTTP server is used for inner request for "just created" workspace - qa.HTTPFixturesApply(t, scimAPI, func(ctx context.Context, wsClient *common.DatabricksClient) { - // a bit hacky, but the whole thing is more readable - accountsAPI[1].Response = Workspace{ - WorkspaceStatus: "RUNNING", - WorkspaceURL: wsClient.Config.Host, - } - state["workspace_url"] = wsClient.Config.Host - state["workspace_name"] = "b" - state["account_id"] = "c" - state["storage_customer_managed_key_id"] = "1234" - state["is_no_public_ip_enabled"] = "false" - qa.ResourceFixture{ - Fixtures: accountsAPI, - Resource: ResourceMwsWorkspaces(), - InstanceState: state, - Update: true, - ID: "a", - HCL: hcl + ` - workspace_name = "b" - account_id = "c", - storage_customer_managed_key_id = "1234"`, - }.Apply(t) - }) -} +func TestWorkspace_verifyWorkspaceReachable(t *testing.T) { + // Create a mock client + mockClient := mocks.NewMockWorkspaceClient(t) -func TestUpdateWorkspace_AddToken(t *testing.T) { - updateWorkspaceScimFixture(t, []qa.HTTPFixture{ - { - Method: "POST", - Resource: "/api/2.0/token/create", - ExpectedRequest: Token{ - LifetimeSeconds: 2.592e+06, - Comment: "Terraform PAT", - }, - Response: tokens.TokenResponse{ - TokenValue: "sensitive", - TokenInfo: &tokens.TokenInfo{ - TokenID: "abcdef", - }, - }, - }, - }, map[string]string{ - // no token in state - }, `token {}`) -} + // Set up expectations for the first call (DNS error) + mockClient.GetMockCurrentUserAPI().EXPECT().Me(mock.Anything).Return(nil, &net.OpError{ + Op: "dial", + Net: "tcp", + Err: &net.DNSError{ + Name: "900150983cd24fb0.cloud.databricks.com", + Err: "no such host", + }, + }).Once() -func TestUpdateWorkspace_AddTokenAndChangeNetworkId(t *testing.T) { - updateWorkspaceScimFixtureWithPatch(t, []qa.HTTPFixture{ - { - Method: "POST", - Resource: "/api/2.0/token/create", - ExpectedRequest: Token{ - LifetimeSeconds: 2.592e+06, - Comment: "Terraform PAT", - }, - Response: tokens.TokenResponse{ - TokenValue: "sensitive", - TokenInfo: &tokens.TokenInfo{ - TokenID: "abcdef", - }, - }, - }, - }, map[string]string{ - "network_id": "alpha", - // no token in state - }, ` - network_id = "beta" - token {} - `) -} + // Set up expectations for the second call (success) + mockClient.GetMockCurrentUserAPI().EXPECT().Me(mock.Anything).Return(&iam.User{ + Id: "12345", + }, nil).Once() -func TestUpdateWorkspace_DeleteTokenAndChangeNetworkId(t *testing.T) { - updateWorkspaceScimFixtureWithPatch(t, []qa.HTTPFixture{ - { - Method: "POST", - Resource: "/api/2.0/token/delete", - ExpectedRequest: map[string]any{ - "token_id": "abcdef", - }, - }, - }, map[string]string{ - "token.#": "1", - "token.0.comment": "Terraform PAT", - "token.0.lifetime_seconds": "2592000", - "token.0.token_id": "abcdef", - "token.0.token_value": "sensitive", - "network_id": "alpha", - }, ` - network_id = "beta" - `) + // Create a context + ctx := context.Background() -} + // Call the function with the mock client + err := verifyWorkspaceReachable(ctx, mockClient.WorkspaceClient) -func TestUpdateWorkspace_DeleteToken(t *testing.T) { - updateWorkspaceScimFixture(t, []qa.HTTPFixture{ - { - Method: "POST", - Resource: "/api/2.0/token/delete", - ExpectedRequest: map[string]any{ - "token_id": "abcdef", - }, - }, - }, map[string]string{ - "token.#": "1", - "token.0.comment": "Terraform PAT", - "token.0.lifetime_seconds": "2592000", - "token.0.token_id": "abcdef", - "token.0.token_value": "sensitive", - }, ``) + // The function should retry and eventually succeed + assert.NoError(t, err) } -func TestUpdateWorkspace_ReplaceTokenAndChangeNetworkId(t *testing.T) { - updateWorkspaceScimFixtureWithPatch(t, []qa.HTTPFixture{ - { - Method: "POST", - Resource: "/api/2.0/token/delete", - ExpectedRequest: map[string]any{ - "token_id": "abcdef", - }, - }, - { - Method: "POST", - Resource: "/api/2.0/token/create", - ExpectedRequest: Token{ - LifetimeSeconds: 2.592e+06, - Comment: "I am Batman!", - }, - Response: tokens.TokenResponse{ - TokenValue: "new-value", - TokenInfo: &tokens.TokenInfo{ - TokenID: "new-id", - }, - }, - }, - }, map[string]string{ - "token.#": "1", - "token.0.comment": "Terraform PAT", - "token.0.lifetime_seconds": "2592000", - "token.0.token_id": "abcdef", - "token.0.token_value": "sensitive", - "network_id": "alpha", - }, - ` - network_id = "beta" - token { - comment = "I am Batman!" - }`) -} +func TestEnsureTokenExists(t *testing.T) { + // Create a mock workspace client + mockClient := mocks.NewMockWorkspaceClient(t) + mockTokensAPI := mockClient.GetMockTokensAPI() -func TestUpdateWorkspace_ReplaceToken(t *testing.T) { - updateWorkspaceScimFixture(t, []qa.HTTPFixture{ - { - Method: "POST", - Resource: "/api/2.0/token/delete", - ExpectedRequest: map[string]any{ - "token_id": "abcdef", - }, - }, - { - Method: "POST", - Resource: "/api/2.0/token/create", - ExpectedRequest: Token{ - LifetimeSeconds: 2.592e+06, - Comment: "I am Batman!", - }, - Response: tokens.TokenResponse{ - TokenValue: "new-value", - TokenInfo: &tokens.TokenInfo{ - TokenID: "new-id", - }, - }, - }, - }, map[string]string{ - "token.#": "1", - "token.0.comment": "Terraform PAT", - "token.0.lifetime_seconds": "2592000", - "token.0.token_id": "abcdef", - "token.0.token_value": "sensitive", - }, `token { - comment = "I am Batman!" - }`) -} + // Set up expectations for token list and create + mockTokensAPI.EXPECT(). + List(mock.Anything). + Return(&listing.SliceIterator[settings.PublicTokenInfo]{}). + Times(1) -func TestEnsureTokenExists(t *testing.T) { - qa.HTTPFixturesApply(t, []qa.HTTPFixture{ - { - Method: "GET", - Resource: "/api/2.0/token/list", - Response: `{}`, // we just need a JSON for this - }, - { - Method: "POST", - Resource: "/api/2.0/token/create", - ExpectedRequest: Token{ - LifetimeSeconds: 3600, - Comment: "test", - }, - Response: tokens.TokenResponse{ - TokenValue: "new-value", - TokenInfo: &tokens.TokenInfo{ - TokenID: "new-id", - }, - }, - }, - }, func(ctx context.Context, client *common.DatabricksClient) { - r := ResourceMwsWorkspaces() - d := r.ToResource().TestResourceData() - d.Set("workspace_url", client.Config.Host) - d.Set("token", []any{ - map[string]any{ - "lifetime_seconds": 3600, - "comment": "test", - "token_id": "abcdef", - }, - }) - wsApi := NewWorkspacesAPI(context.Background(), client) - err := EnsureTokenExistsIfNeeded(wsApi, r.Schema, d) - assert.NoError(t, err) - }) + mockTokensAPI.EXPECT(). + Create(mock.Anything, settings.CreateTokenRequest{ + LifetimeSeconds: 3600, + Comment: "test", + }). + Return(&settings.CreateTokenResponse{ + TokenValue: "new-value", + TokenInfo: &settings.PublicTokenInfo{ + TokenId: "new-id", + }, + }, nil). + Times(1) + + // Test the function + token := &Token{ + LifetimeSeconds: 3600, + Comment: "test", + } + err := ensureTokenExists(context.Background(), mockClient.WorkspaceClient, token) + assert.NoError(t, err) + assert.Equal(t, token.TokenID, "new-id") + assert.Equal(t, token.TokenValue, "new-value") } func TestEnsureTokenExists_NoRecreate(t *testing.T) { - qa.HTTPFixturesApply(t, []qa.HTTPFixture{ - { - Method: "GET", - Resource: "/api/2.0/token/list", - Response: tokens.TokenList{ - TokenInfos: []tokens.TokenInfo{ - { - TokenID: "old-id", - }, - }, - }, - }, - }, func(ctx context.Context, client *common.DatabricksClient) { - r := ResourceMwsWorkspaces() - d := r.ToResource().TestResourceData() - d.Set("workspace_url", client.Config.Host) - d.Set("token", []any{ - map[string]any{ - "lifetime_seconds": 3600, - "comment": "test", - "token_id": "old-id", + // Create a mock workspace client + mockClient := mocks.NewMockWorkspaceClient(t) + mockTokensAPI := mockClient.GetMockTokensAPI() + + // Set up expectations for token list + mockTokensAPI.EXPECT(). + List(mock.Anything). + Return(&listing.SliceIterator[settings.PublicTokenInfo]{ + { + TokenId: "old-id", }, - }) - wsApi := NewWorkspacesAPI(context.Background(), client) - err := EnsureTokenExistsIfNeeded(wsApi, r.Schema, d) - assert.NoError(t, err) - }) -} + }). + Times(1) -func TestWorkspaceTokenWrongAuthCornerCase(t *testing.T) { - client, err := client.New(&config.Config{}) - if err != nil { - t.Fatal(err) + // Test the function + token := &Token{ + LifetimeSeconds: 3600, + Comment: "test", + TokenID: "old-id", } - r := ResourceMwsWorkspaces() - d := r.ToResource().TestResourceData() - d.Set("workspace_url", client.Config.Host) - d.Set("token", []any{ - map[string]any{ - "lifetime_seconds": 3600, - "comment": "test", - "token_id": "old-id", - }, - }) + err := ensureTokenExists(context.Background(), mockClient.WorkspaceClient, token) + assert.NoError(t, err) +} - wsApi := NewWorkspacesAPI(context.Background(), &common.DatabricksClient{ - DatabricksClient: client, +func TestExplainWorkspaceFailureCornerCase(t *testing.T) { + t.Run("no network ID", func(t *testing.T) { + assert.EqualError(t, explainWorkspaceFailure(context.Background(), nil, &provisioning.Workspace{ + WorkspaceStatusMessage: "🔥", + }), "workspace status message: 🔥") }) - noAuth := "cannot authenticate parent client: " + common.NoAuth + t.Run("network error", func(t *testing.T) { + mockClient := mocks.NewMockAccountClient(t) + mockNetworksClient := mockClient.GetMockNetworksAPI() - assert.EqualError(t, CreateTokenIfNeeded(wsApi, r.Schema, d), noAuth, "create") - assert.EqualError(t, EnsureTokenExistsIfNeeded(wsApi, r.Schema, d), noAuth, "ensure") - assert.EqualError(t, removeTokenIfNeeded(wsApi, "x", d), noAuth, "remove") -} + mockNetworksClient.EXPECT(). + Get(context.Background(), provisioning.GetNetworkRequest{NetworkId: "abc"}). + Return(nil, errors.New("🐜")). + Times(1) -func TestWorkspaceTokenHttpCornerCases(t *testing.T) { - qa.HTTPFixturesApply(t, []qa.HTTPFixture{ - { - MatchAny: true, - ReuseRequest: true, - Status: 418, - Response: apierr.APIError{ - ErrorCode: "NONSENSE", - StatusCode: 418, - Message: "i'm a teapot", - }, - }, - }, func(ctx context.Context, client *common.DatabricksClient) { - wsApi := NewWorkspacesAPI(context.Background(), client) - r := ResourceMwsWorkspaces() - d := r.ToResource().TestResourceData() - d.Set("workspace_url", client.Config.Host) - d.Set("token", []any{ - map[string]any{ - "lifetime_seconds": 3600, - "comment": "test", - "token_id": "old-id", - }, - }) - for msg, err := range map[string]error{ - "cannot create token: i'm a teapot": CreateTokenIfNeeded(wsApi, r.Schema, d), - "cannot read token: i'm a teapot": EnsureTokenExistsIfNeeded(wsApi, r.Schema, d), - "cannot remove token: i'm a teapot": removeTokenIfNeeded(wsApi, "x", d), - } { - assert.EqualError(t, err, msg) + ws := &provisioning.Workspace{ + NetworkId: "abc", + WorkspaceStatusMessage: "🔥", } + assert.EqualError(t, explainWorkspaceFailure(context.Background(), mockClient.AccountClient, ws), "workspace status message: 🔥, network error message: cannot read network: 🐜") }) } -func TestGenerateWorkspaceHostname_CornerCases(t *testing.T) { - assert.Equal(t, "fallback.cloud.databricks.com", - generateWorkspaceHostname(&common.DatabricksClient{ - DatabricksClient: &client.DatabricksClient{ - Config: &config.Config{ - Host: "$%^&*", - }, - }, - }, Workspace{ - DeploymentName: "fallback", - })) - assert.Equal(t, "stuff.is.exaple.com", - generateWorkspaceHostname(&common.DatabricksClient{ - DatabricksClient: &client.DatabricksClient{ - Config: &config.Config{ - Host: "https://this.is.exaple.com", - }, - }, - }, Workspace{ - DeploymentName: "stuff", - })) -} +func TestResourceWorkspaceUpdatePrivateAccessSettings(t *testing.T) { + // Define a mock workspace that can be reused + mockWorkspace := &provisioning.Workspace{ + WorkspaceId: 1234, + WorkspaceStatus: provisioning.WorkspaceStatusRunning, + WorkspaceName: "labdata", + DeploymentName: "900150983cd24fb0", + AwsRegion: "us-east-1", + CredentialsId: "bcd", + StorageConfigurationId: "ghi", + NetworkId: "fgh", + ManagedServicesCustomerManagedKeyId: "def", + StorageCustomerManagedKeyId: "def", + PrivateAccessSettingsId: "pas", + AccountId: "abc", + } -func TestExplainWorkspaceFailureCornerCase(t *testing.T) { - qa.HTTPFixturesApply(t, []qa.HTTPFixture{ - { - MatchAny: true, - ReuseRequest: true, - Status: 418, - Response: apierr.APIError{ - ErrorCode: "NONSENSE", - StatusCode: 418, - Message: "🐜", - }, + // Create a mock waiter + mockWaiter := &provisioning.WaitGetWorkspaceRunning[struct{}]{ + WorkspaceId: 1234, + Poll: func(d time.Duration, f func(*provisioning.Workspace)) (*provisioning.Workspace, error) { + return mockWorkspace, nil }, - }, func(ctx context.Context, client *common.DatabricksClient) { - wsApi := NewWorkspacesAPI(context.Background(), client) - - assert.EqualError(t, wsApi.explainWorkspaceFailure(Workspace{ - WorkspaceStatusMessage: "🔥", - }), "🔥") - - assert.EqualError(t, wsApi.explainWorkspaceFailure(Workspace{ - NetworkID: "abc", - }), "failed to start workspace. Cannot read network: 🐜") - }) -} + } -func TestResourceWorkspaceUpdatePrivateAccessSettings(t *testing.T) { d, err := qa.ResourceFixture{ - Fixtures: []qa.HTTPFixture{ - { - Method: "PATCH", - Resource: "/api/2.0/accounts/abc/workspaces/1234", - ExpectedRequest: map[string]any{ - "credentials_id": "bcd", - "network_id": "fgh", - "storage_customer_managed_key_id": "def", - "private_access_settings_id": "pas", - }, - }, - { - Method: "GET", - ReuseRequest: true, - Resource: "/api/2.0/accounts/abc/workspaces/1234", - Response: Workspace{ - WorkspaceStatus: WorkspaceStatusRunning, - WorkspaceName: "labdata", - DeploymentName: "900150983cd24fb0", - AwsRegion: "us-east-1", - CredentialsID: "bcd", - StorageConfigurationID: "ghi", - NetworkID: "fgh", - ManagedServicesCustomerManagedKeyID: "def", - StorageCustomerManagedKeyID: "def", - PrivateAccessSettingsID: "pas", - AccountID: "abc", - WorkspaceID: 1234, - }, - }, + MockAccountClientFunc: func(a *mocks.MockAccountClient) { + // Expect the Update call + a.GetMockWorkspacesAPI().EXPECT().Update(mock.Anything, provisioning.UpdateWorkspaceRequest{ + WorkspaceId: 1234, + CredentialsId: "bcd", + NetworkId: "fgh", + StorageCustomerManagedKeyId: "def", + PrivateAccessSettingsId: "pas", + }).Return(mockWaiter, nil) + + // Expect the Get call to retrieve the updated workspace + a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ + WorkspaceId: 1234, + }).Return(mockWorkspace, nil) }, Resource: ResourceMwsWorkspaces(), InstanceState: map[string]string{ @@ -1645,45 +1142,46 @@ func TestResourceWorkspaceRemovePAS_NotAllowed(t *testing.T) { } func TestResourceWorkspaceCreateGcpManagedVPC(t *testing.T) { + // Define a mock workspace that can be reused + mockWorkspace := &provisioning.Workspace{ + WorkspaceId: 1234, + WorkspaceStatus: provisioning.WorkspaceStatusRunning, + WorkspaceName: "labdata", + DeploymentName: "900150983cd24fb0", + AccountId: "abc", + Cloud: "gcp", + Location: "bcd", + GcpManagedNetworkConfig: &provisioning.GcpManagedNetworkConfig{ + SubnetCidr: "a", + }, + } + + // Create a mock waiter + mockWaiter := &provisioning.WaitGetWorkspaceRunning[provisioning.Workspace]{ + WorkspaceId: 1234, + Poll: func(d time.Duration, f func(*provisioning.Workspace)) (*provisioning.Workspace, error) { + return mockWorkspace, nil + }, + } + qa.ResourceFixture{ - Fixtures: []qa.HTTPFixture{ - { - Method: "POST", - Resource: "/api/2.0/accounts/abc/workspaces", - // retreating to raw JSON, as certain fields don't work well together - ExpectedRequest: map[string]any{ - "account_id": "abc", - "cloud": "gcp", - "cloud_resource_container": map[string]any{ - "gcp": map[string]any{ - "project_id": "def", - }, + MockAccountClientFunc: func(a *mocks.MockAccountClient) { + // Expect the Create call + a.GetMockWorkspacesAPI().EXPECT().Create(mock.Anything, provisioning.CreateWorkspaceRequest{ + Cloud: "gcp", + CloudResourceContainer: &provisioning.CloudResourceContainer{ + Gcp: &provisioning.CustomerFacingGcpCloudResourceContainer{ + ProjectId: "def", }, - "location": "bcd", - "workspace_name": "labdata", - }, - Response: Workspace{ - WorkspaceID: 1234, - AccountID: "abc", - DeploymentName: "900150983cd24fb0", - WorkspaceName: "labdata", }, - }, - { - Method: "GET", - ReuseRequest: true, - Resource: "/api/2.0/accounts/abc/workspaces/1234", - Response: Workspace{ - AccountID: "abc", - WorkspaceID: 1234, - WorkspaceStatus: WorkspaceStatusRunning, - DeploymentName: "900150983cd24fb0", - WorkspaceName: "labdata", - GCPManagedNetworkConfig: &GCPManagedNetworkConfig{ - SubnetCIDR: "a", - }, - }, - }, + Location: "bcd", + WorkspaceName: "labdata", + }).Return(mockWaiter, nil) + + // Expect the Get call to retrieve the workspace + a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ + WorkspaceId: 1234, + }).Return(mockWorkspace, nil) }, Resource: ResourceMwsWorkspaces(), HCL: ` From a753d9b579c1bf9672297f29d147c2d6bf99b7a6 Mon Sep 17 00:00:00 2001 From: Miles Yucht Date: Thu, 10 Apr 2025 16:00:50 +0200 Subject: [PATCH 02/16] work --- mws/mws_workspaces_test.go | 65 +++++++++++++++++++++++++--------- mws/resource_mws_workspaces.go | 27 ++++++++------ 2 files changed, 65 insertions(+), 27 deletions(-) diff --git a/mws/mws_workspaces_test.go b/mws/mws_workspaces_test.go index 4356f646ec..d01367e411 100644 --- a/mws/mws_workspaces_test.go +++ b/mws/mws_workspaces_test.go @@ -19,6 +19,7 @@ import ( "github.com/hashicorp/terraform-plugin-sdk/v2/diag" "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" "github.com/hashicorp/terraform-plugin-testing/helper/resource" + "github.com/hashicorp/terraform-plugin-testing/plancheck" "github.com/hashicorp/terraform-plugin-testing/terraform" "github.com/stretchr/testify/assert" ) @@ -77,19 +78,28 @@ func TestMwsAccWorkspaces(t *testing.T) { }) } -func TestMwsAccWorkspaces_TokenUpdate(t *testing.T) { - baseTemplate := ` - resource "databricks_mws_credentials" "this" { - account_id = "{env.DATABRICKS_ACCOUNT_ID}" - credentials_name = "credentials-ws-{var.RANDOM}" - role_arn = "{env.TEST_CROSSACCOUNT_ARN}" +type expectNotDestroyed struct { + addr string +} + +func ExpectNotDestroyed(addr string) expectNotDestroyed { + return expectNotDestroyed{addr: addr} +} + +func (e expectNotDestroyed) CheckPlan(ctx context.Context, req plancheck.CheckPlanRequest, resp *plancheck.CheckPlanResponse) { + for _, resource := range req.Plan.ResourceChanges { + if resource.Address != e.addr { + continue } - resource "databricks_mws_storage_configurations" "this" { - account_id = "{env.DATABRICKS_ACCOUNT_ID}" - storage_configuration_name = "storage-ws-{var.RANDOM}" - bucket_name = "{env.TEST_ROOT_BUCKET}" + actions := resource.Change.Actions + if actions.DestroyBeforeCreate() || actions.CreateBeforeDestroy() || actions.Delete() { + resp.Error = fmt.Errorf("resource %s is marked for destruction", e.addr) + return } - ` + } +} + +func TestMwsAccWorkspaces_TokenUpdate(t *testing.T) { tokenUpdateTemplate := func(token, customTags string) string { tokenBlock := `` if token != "" { @@ -105,7 +115,17 @@ func TestMwsAccWorkspaces_TokenUpdate(t *testing.T) { %s }`, customTags) } - return fmt.Sprintf(baseTemplate+` + return fmt.Sprintf(` + resource "databricks_mws_credentials" "this" { + account_id = "{env.DATABRICKS_ACCOUNT_ID}" + credentials_name = "credentials-ws-{var.STICKY_RANDOM}" + role_arn = "{env.TEST_CROSSACCOUNT_ARN}" + } + resource "databricks_mws_storage_configurations" "this" { + account_id = "{env.DATABRICKS_ACCOUNT_ID}" + storage_configuration_name = "storage-ws-{var.STICKY_RANDOM}" + bucket_name = "{env.TEST_ROOT_BUCKET}" + } resource "databricks_mws_workspaces" "this" { account_id = "{env.DATABRICKS_ACCOUNT_ID}" workspace_name = "terra-{var.STICKY_RANDOM}" @@ -115,9 +135,7 @@ func TestMwsAccWorkspaces_TokenUpdate(t *testing.T) { storage_configuration_id = databricks_mws_storage_configurations.this.storage_configuration_id %s - custom_tags = { - %s - } + %s }`, tokenBlock, customTagsBlock) } @@ -129,7 +147,7 @@ func TestMwsAccWorkspaces_TokenUpdate(t *testing.T) { } a := databricks.Must(databricks.NewAccountClient()) ctx := context.Background() - workspace, err := a.Workspaces.Get(ctx, provisioning.GetWorkspaceRequest{WorkspaceId: common.MustInt64(state.Primary.ID)}) + workspace, err := a.Workspaces.Get(ctx, provisioning.GetWorkspaceRequest{WorkspaceId: common.MustInt64(state.Primary.Attributes["workspace_id"])}) assert.NoError(t, err) w, err := a.GetWorkspaceClient(*workspace) @@ -145,7 +163,11 @@ func TestMwsAccWorkspaces_TokenUpdate(t *testing.T) { tokens := w.Tokens.List(context.Background()) ctx := context.Background() for tokens.HasNext(ctx) { - if token, err := tokens.Next(ctx); err != nil && token.TokenId == tokenId { + token, err := tokens.Next(ctx) + if err != nil { + return fmt.Errorf("error fetching tokens: %w", err) + } + if token.TokenId == tokenId { return nil } } @@ -159,6 +181,9 @@ func TestMwsAccWorkspaces_TokenUpdate(t *testing.T) { }, acceptance.Step{ // Updating the comment causes the old token to be deleted and a new one to be created Template: tokenUpdateTemplate(`comment = "test bar"`, ""), + ConfigPlanChecks: resource.ConfigPlanChecks{ + PreApply: []plancheck.PlanCheck{ExpectNotDestroyed("databricks_mws_workspaces.this")}, + }, Check: resource.ComposeAggregateTestCheckFunc( checkTokenExists, // Capture the token ID at the end of this step to verify it is not changed in future steps. @@ -175,6 +200,9 @@ func TestMwsAccWorkspaces_TokenUpdate(t *testing.T) { }, acceptance.Step{ // Modifying the tags doesn't change the token but does modify the workspace. Template: tokenUpdateTemplate(`comment = "test bar"`, `"Key" = "Value"`), + ConfigPlanChecks: resource.ConfigPlanChecks{ + PreApply: []plancheck.PlanCheck{ExpectNotDestroyed("databricks_mws_workspaces.this")}, + }, Check: func(s *terraform.State) error { state, ok := s.RootModule().Resources["databricks_mws_workspaces.this"] if !ok { @@ -188,6 +216,9 @@ func TestMwsAccWorkspaces_TokenUpdate(t *testing.T) { }, acceptance.Step{ // It is also possible to modify the token comment and tags at the same time. Template: tokenUpdateTemplate(`comment = "test quux"`, `"Key" = "Value2"`), + ConfigPlanChecks: resource.ConfigPlanChecks{ + PreApply: []plancheck.PlanCheck{ExpectNotDestroyed("databricks_mws_workspaces.this")}, + }, Check: func(s *terraform.State) error { state, ok := s.RootModule().Resources["databricks_mws_workspaces.this"] if !ok { diff --git a/mws/resource_mws_workspaces.go b/mws/resource_mws_workspaces.go index 39b98bdd42..c519e8e402 100644 --- a/mws/resource_mws_workspaces.go +++ b/mws/resource_mws_workspaces.go @@ -30,11 +30,11 @@ const DefaultProvisionTimeout = 20 * time.Minute // Workspace is the object that contains all the information for deploying a workspace type Workspace struct { - *provisioning.Workspace + provisioning.Workspace CustomerManagedKeyID string `json:"customer_managed_key_id,omitempty"` // just for compatibility, will be removed + GcpWorkspaceSa string `json:"gcp_workspace_sa" tf:"computed"` Token *Token `json:"token,omitempty"` WorkspaceURL string `json:"workspace_url,omitempty" tf:"computed"` - GcpWorkspaceSa string `json:"gcp_workspace_sa" tf:"computed"` } // wait for DNS caches to refresh, as sometimes we cannot make @@ -286,7 +286,7 @@ func ResourceMwsWorkspaces() common.Resource { if err := verifyWorkspaceReachable(ctx, w); err != nil { return err } - workspaceConfig.Workspace = workspace + workspaceConfig.Workspace = *workspace workspaceConfig.WorkspaceURL = fmt.Sprintf("https://%s", w.Config.CanonicalHostName()) if c.IsGcp() { workspaceConfig.GcpWorkspaceSa = fmt.Sprintf("db-%d@prod-gcp-%s.iam.gserviceaccount.com", workspace.WorkspaceId, workspace.Location) @@ -299,8 +299,11 @@ func ResourceMwsWorkspaces() common.Resource { return err } } + if err := common.StructToData(workspaceConfig, workspaceSchema, d); err != nil { + return err + } p.Pack(d) - return common.StructToData(workspaceConfig, workspaceSchema, d) + return nil }, Read: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error { a, err := c.AccountClient() @@ -325,7 +328,7 @@ func ResourceMwsWorkspaces() common.Resource { // Default the value of `is_no_public_ip_enabled` because it isn't part of the GET payload. // The field is only used on creation and we therefore suppress all diffs. workspace.IsNoPublicIpEnabled = true - workspaceConfig.Workspace = workspace + workspaceConfig.Workspace = *workspace _, err = a.Workspaces.WaitGetWorkspaceRunning(ctx, workspace.WorkspaceId, d.Timeout(schema.TimeoutRead), nil) if err != nil { return err @@ -350,6 +353,10 @@ func ResourceMwsWorkspaces() common.Resource { common.DataToStructPointer(d, workspaceSchema, &updateWorkspaceRequest) var workspaceConfig Workspace common.DataToStructPointer(d, workspaceSchema, &workspaceConfig) + // WorkspaceId in UpdateWorkspaceRequest is a path parameter, thus tagged with `json:"-"`. + // This causes it not to be set in DataToStructPointer. Instead, the workspace ID can be + // retrieved from Workspace. + updateWorkspaceRequest.WorkspaceId = workspaceConfig.WorkspaceId if customerManagedKeyId, ok := d.GetOk("customer_managed_key_id"); ok && updateWorkspaceRequest.ManagedServicesCustomerManagedKeyId == "" { log.Print("[INFO] Using existing customer_managed_key_id as value for new managed_services_customer_managed_key_id") updateWorkspaceRequest.ManagedServicesCustomerManagedKeyId = customerManagedKeyId.(string) @@ -365,21 +372,21 @@ func ResourceMwsWorkspaces() common.Resource { if err != nil { return fmt.Errorf("%w: %w", err, explainWorkspaceFailure(ctx, a, workspace)) } - workspaceConfig.Workspace = workspace + workspaceConfig.Workspace = *workspace } // If the `token` field has been modified, update the token correspondingly. if d.HasChange("token") { - w, err := a.GetWorkspaceClient(*workspaceConfig.Workspace) + w, err := a.GetWorkspaceClient(workspaceConfig.Workspace) if err != nil { return err } // If there was a token present in the config before, revoke it. rawOld, _ := d.GetChange("token") - oldToken := rawOld.([]map[string]any) - if len(oldToken) > 0 { - raw := oldToken[0] + oldTokens := rawOld.([]any) + if len(oldTokens) > 0 { + raw := oldTokens[0].(map[string]any) id := raw["token_id"].(string) if err := removeToken(ctx, w, id); err != nil { return err From ec78f7e61a723b2158ab09f9191638892c75d316 Mon Sep 17 00:00:00 2001 From: Miles Yucht Date: Thu, 10 Apr 2025 18:39:52 +0200 Subject: [PATCH 03/16] some work --- common/client.go | 35 +++++++++++++++++++++++++++++ mws/resource_mws_workspaces.go | 4 ++-- mws/resource_mws_workspaces_test.go | 24 +------------------- 3 files changed, 38 insertions(+), 25 deletions(-) diff --git a/common/client.go b/common/client.go index cd04a67531..ba75e01b8c 100644 --- a/common/client.go +++ b/common/client.go @@ -13,6 +13,7 @@ import ( "github.com/databricks/databricks-sdk-go/client" "github.com/databricks/databricks-sdk-go/config" "github.com/databricks/databricks-sdk-go/service/iam" + "github.com/databricks/databricks-sdk-go/service/provisioning" "github.com/golang-jwt/jwt/v4" "github.com/hashicorp/terraform-plugin-framework/diag" "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" @@ -54,7 +55,11 @@ type DatabricksClient struct { // callback used to create API1.2 call wrapper, which simplifies unit testing commandFactory func(context.Context, *DatabricksClient) CommandExecutor cachedWorkspaceClient *databricks.WorkspaceClient + // cachedWorkspaceClients is a map of workspace clients for each workspace ID + cachedWorkspaceClients map[int64]*databricks.WorkspaceClient cachedAccountClient *databricks.AccountClient + + // mu synchronizes access to all cached clients. mu sync.Mutex } @@ -90,6 +95,36 @@ func (c *DatabricksClient) SetWorkspaceClient(w *databricks.WorkspaceClient) { c.cachedWorkspaceClient = w } +func (c *DatabricksClient) WorkspaceClientForWorkspace(ctx context.Context, workspaceId int64) (*databricks.WorkspaceClient, error) { + if !c.Config.IsAccountClient() { + return nil, fmt.Errorf("provider must be configured at the account level when calling WorkspaceClientForWorkspace") + } + c.mu.Lock() + defer c.mu.Unlock() + if c.cachedWorkspaceClients == nil { + c.cachedWorkspaceClients = make(map[int64]*databricks.WorkspaceClient) + } + if client, ok := c.cachedWorkspaceClients[workspaceId]; ok { + return client, nil + } + a, err := c.AccountClient() + if err != nil { + return nil, err + } + workspace, err := a.Workspaces.Get(ctx, provisioning.GetWorkspaceRequest{WorkspaceId: workspaceId}) + if err != nil { + return nil, err + } + w, err := a.GetWorkspaceClient(*workspace) + if err != nil { + return nil, err + } + w.CurrentUser = newCachedMe(w.CurrentUser) + w.Config.AzureResourceID = workspace.AzureResourceId() + c.cachedWorkspaceClients[workspace.WorkspaceId] = w + return w, nil +} + // Set the cached account client. func (c *DatabricksClient) SetAccountClient(a *databricks.AccountClient) { c.mu.Lock() diff --git a/mws/resource_mws_workspaces.go b/mws/resource_mws_workspaces.go index c519e8e402..d55e730412 100644 --- a/mws/resource_mws_workspaces.go +++ b/mws/resource_mws_workspaces.go @@ -43,8 +43,7 @@ func verifyWorkspaceReachable(ctx context.Context, w *databricks.WorkspaceClient return retries.New[struct{}](retries.WithRetryFunc(func(err error) bool { var dnsError *net.DNSError if errors.As(err, &dnsError) { - err = fmt.Errorf("workspace %s is not yet reachable: %s", w.Config.Host, dnsError) - log.Printf("[INFO] %s", err) + log.Printf("[INFO] workspace is not yet reachable: %s", dnsError) // expected to retry on: dial tcp: lookup XXX: no such host return true } @@ -195,6 +194,7 @@ func ResourceMwsWorkspaces() common.Resource { s["is_no_public_ip_enabled"].DiffSuppressFunc = func(k, old, new string, d *schema.ResourceData) bool { return old != "" } + s["is_no_public_ip_enabled"].Default = true s["pricing_tier"].DiffSuppressFunc = common.EqualFoldDiffSuppress diff --git a/mws/resource_mws_workspaces_test.go b/mws/resource_mws_workspaces_test.go index 3baf4fc58a..bc1d0b1d86 100644 --- a/mws/resource_mws_workspaces_test.go +++ b/mws/resource_mws_workspaces_test.go @@ -154,29 +154,6 @@ func TestResourceWorkspaceCreateGcp(t *testing.T) { func TestResourceWorkspaceCreate_Error_Custom_tags(t *testing.T) { qa.ResourceFixture{ - MockAccountClientFunc: func(a *mocks.MockAccountClient) { - a.GetMockWorkspacesAPI().EXPECT().Create(mock.Anything, provisioning.CreateWorkspaceRequest{ - Cloud: "gcp", - CloudResourceContainer: &provisioning.CloudResourceContainer{ - Gcp: &provisioning.CustomerFacingGcpCloudResourceContainer{ - ProjectId: "def", - }, - }, - Location: "bcd", - PrivateAccessSettingsId: "pas_id_a", - NetworkId: "net_id_a", - GcpManagedNetworkConfig: &provisioning.GcpManagedNetworkConfig{ - SubnetCidr: "a", - }, - WorkspaceName: "labdata", - CustomTags: map[string]string{ - "SoldToCode": "1234", - }, - }).Return(nil, &apierr.APIError{ - ErrorCode: "INVALID_PARAMETER_VALUE", - Message: "custom_tags are only allowed for AWS workspaces", - }) - }, Resource: ResourceMwsWorkspaces(), HCL: ` account_id = "abc" @@ -478,6 +455,7 @@ func TestResourceWorkspaceRead(t *testing.T) { a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ WorkspaceId: 1234, }).Return(mockWorkspace, nil) + a.GetMockWorkspacesAPI().EXPECT().WaitGetWorkspaceRunning(mock.Anything, int64(1234), 20*time.Minute, mock.Anything).Return(mockWorkspace, nil) }, Resource: ResourceMwsWorkspaces(), Read: true, From bd963653ea397a86da6b2e00a01d0c887737be50 Mon Sep 17 00:00:00 2001 From: Miles Yucht Date: Thu, 10 Apr 2025 21:33:31 +0200 Subject: [PATCH 04/16] work --- common/client.go | 27 ++++- mws/resource_mws_workspaces.go | 49 ++++---- mws/resource_mws_workspaces_test.go | 174 ++++++++++++++++++++-------- qa/testing.go | 13 ++- 4 files changed, 191 insertions(+), 72 deletions(-) diff --git a/common/client.go b/common/client.go index ba75e01b8c..b519529421 100644 --- a/common/client.go +++ b/common/client.go @@ -54,9 +54,15 @@ type DatabricksClient struct { // callback used to create API1.2 call wrapper, which simplifies unit testing commandFactory func(context.Context, *DatabricksClient) CommandExecutor + // cachedWorkspaceClient is a cached workspace client authenticated to the workspace + // configured for the provider cachedWorkspaceClient *databricks.WorkspaceClient // cachedWorkspaceClients is a map of workspace clients for each workspace ID + // populated when fetching a WorkspaceClient for a specific workspace ID using + // a provider configured at the account level cachedWorkspaceClients map[int64]*databricks.WorkspaceClient + // cachedAccountClient is a cached account client authenticated to the account + // configured for the provider cachedAccountClient *databricks.AccountClient // mu synchronizes access to all cached clients. @@ -96,9 +102,6 @@ func (c *DatabricksClient) SetWorkspaceClient(w *databricks.WorkspaceClient) { } func (c *DatabricksClient) WorkspaceClientForWorkspace(ctx context.Context, workspaceId int64) (*databricks.WorkspaceClient, error) { - if !c.Config.IsAccountClient() { - return nil, fmt.Errorf("provider must be configured at the account level when calling WorkspaceClientForWorkspace") - } c.mu.Lock() defer c.mu.Unlock() if c.cachedWorkspaceClients == nil { @@ -107,7 +110,7 @@ func (c *DatabricksClient) WorkspaceClientForWorkspace(ctx context.Context, work if client, ok := c.cachedWorkspaceClients[workspaceId]; ok { return client, nil } - a, err := c.AccountClient() + a, err := c.accountClient() if err != nil { return nil, err } @@ -125,6 +128,16 @@ func (c *DatabricksClient) WorkspaceClientForWorkspace(ctx context.Context, work return w, nil } +// SetWorkspaceClientForWorkspace sets the cached workspace client for a specific workspace ID. +func (c *DatabricksClient) SetWorkspaceClientForWorkspace(workspaceId int64, w *databricks.WorkspaceClient) { + c.mu.Lock() + defer c.mu.Unlock() + if c.cachedWorkspaceClients == nil { + c.cachedWorkspaceClients = make(map[int64]*databricks.WorkspaceClient) + } + c.cachedWorkspaceClients[workspaceId] = w +} + // Set the cached account client. func (c *DatabricksClient) SetAccountClient(a *databricks.AccountClient) { c.mu.Lock() @@ -159,6 +172,12 @@ func (c *DatabricksClient) GetAccountClient() (*databricks.AccountClient, diag.D func (c *DatabricksClient) AccountClient() (*databricks.AccountClient, error) { c.mu.Lock() defer c.mu.Unlock() + return c.accountClient() +} + +// accountClient returns the Databricks Account client or an error if that fails. +// The `mu` mutex must be held by the current goroutine when calling this method. +func (c *DatabricksClient) accountClient() (*databricks.AccountClient, error) { if c.cachedAccountClient != nil { return c.cachedAccountClient, nil } diff --git a/mws/resource_mws_workspaces.go b/mws/resource_mws_workspaces.go index d55e730412..2731f8e7d5 100644 --- a/mws/resource_mws_workspaces.go +++ b/mws/resource_mws_workspaces.go @@ -63,7 +63,7 @@ func explainWorkspaceFailure(ctx context.Context, a *databricks.AccountClient, w } network, err := a.Networks.Get(ctx, provisioning.GetNetworkRequest{NetworkId: workspace.NetworkId}) if err != nil { - return fmt.Errorf("%s; cannot read network: %w", errorBase, err) + return fmt.Errorf("%s; network error message: cannot read network: %w", errorBase, err) } var strBuffer bytes.Buffer for _, networkHealth := range network.ErrorMessages { @@ -119,21 +119,24 @@ func isInvalidClient(err error) bool { } func ensureTokenExists(ctx context.Context, w *databricks.WorkspaceClient, token *Token) error { - _, err := w.CurrentUser.Me(ctx) - // If we cannot authenticate to the workspace and we're using an in-house OAuth principal, - // log a warning but do not fail. This can happen if the provider is authenticated with a - // different principal than was used to create the workspace. - if isInvalidClient(err) { - tflog.Debug(ctx, fmt.Sprintf("unable to fetch token with ID %s from workspace using the provided service principal, continuing", token.TokenID)) - return nil - } - if apierr.IsMissing(err) { - return createToken(ctx, w, token) - } - if err != nil { - return fmt.Errorf("cannot read token: %w", err) + tokens := w.Tokens.List(ctx) + for tokens.HasNext(ctx) { + t, err := tokens.Next(ctx) + // If we cannot authenticate to the workspace and we're using an in-house OAuth principal, + // log a warning but do not fail. This can happen if the provider is authenticated with a + // different principal than was used to create the workspace. + if isInvalidClient(err) { + tflog.Debug(ctx, fmt.Sprintf("unable to fetch token with ID %s from workspace using the provided service principal, continuing", token.TokenID)) + return nil + } + if err != nil { + return fmt.Errorf("cannot read token: %w", err) + } + if t.TokenId == token.TokenID { + return nil + } } - return nil + return createToken(ctx, w, token) } func removeToken(ctx context.Context, w *databricks.WorkspaceClient, tokenID string) error { @@ -270,7 +273,12 @@ func ResourceMwsWorkspaces() common.Resource { } workspace, err := wait.GetWithTimeout(d.Timeout(schema.TimeoutCreate)) if err != nil { - err = fmt.Errorf("%w: %w", err, explainWorkspaceFailure(ctx, a, workspace)) + workspace, getErr := a.Workspaces.Get(ctx, provisioning.GetWorkspaceRequest{WorkspaceId: wait.Response.WorkspaceId}) + if getErr != nil { + err = fmt.Errorf("workspace creation failed: %w; failed to get workspace: %w", err, getErr) + } else { + err = fmt.Errorf("%w: %w", err, explainWorkspaceFailure(ctx, a, workspace)) + } log.Printf("[ERROR] Deleting failed workspace: %s", err) if derr := a.Workspaces.Delete(ctx, provisioning.DeleteWorkspaceRequest{WorkspaceId: wait.Response.WorkspaceId}); derr != nil { return fmt.Errorf("workspace creation failed: %w; failed workspace cleanup failed: %w", err, derr) @@ -279,7 +287,7 @@ func ResourceMwsWorkspaces() common.Resource { } // Once the workspace is running, wait for the API to be available by polling the SCIM Me endpoint. - w, err := a.GetWorkspaceClient(*workspace) + w, err := c.WorkspaceClientForWorkspace(ctx, workspace.WorkspaceId) if err != nil { return err } @@ -287,7 +295,7 @@ func ResourceMwsWorkspaces() common.Resource { return err } workspaceConfig.Workspace = *workspace - workspaceConfig.WorkspaceURL = fmt.Sprintf("https://%s", w.Config.CanonicalHostName()) + workspaceConfig.WorkspaceURL = w.Config.CanonicalHostName() if c.IsGcp() { workspaceConfig.GcpWorkspaceSa = fmt.Sprintf("db-%d@prod-gcp-%s.iam.gserviceaccount.com", workspace.WorkspaceId, workspace.Location) } @@ -333,10 +341,11 @@ func ResourceMwsWorkspaces() common.Resource { if err != nil { return err } - w, err := a.GetWorkspaceClient(*workspace) + w, err := c.WorkspaceClientForWorkspace(ctx, workspace.WorkspaceId) if err != nil { return err } + workspaceConfig.WorkspaceURL = w.Config.CanonicalHostName() if workspaceConfig.Token != nil { if err := ensureTokenExists(ctx, w, workspaceConfig.Token); err != nil { return err @@ -377,7 +386,7 @@ func ResourceMwsWorkspaces() common.Resource { // If the `token` field has been modified, update the token correspondingly. if d.HasChange("token") { - w, err := a.GetWorkspaceClient(workspaceConfig.Workspace) + w, err := c.WorkspaceClientForWorkspace(ctx, workspaceConfig.WorkspaceId) if err != nil { return err } diff --git a/mws/resource_mws_workspaces_test.go b/mws/resource_mws_workspaces_test.go index bc1d0b1d86..0982866d76 100644 --- a/mws/resource_mws_workspaces_test.go +++ b/mws/resource_mws_workspaces_test.go @@ -9,6 +9,7 @@ import ( "time" "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/config" "github.com/databricks/databricks-sdk-go/experimental/mocks" "github.com/databricks/databricks-sdk-go/listing" "github.com/databricks/databricks-sdk-go/service/iam" @@ -20,6 +21,34 @@ import ( "github.com/stretchr/testify/mock" ) +func mockScimMe(c *mocks.MockWorkspaceClient) { + c.GetMockCurrentUserAPI().EXPECT().Me(mock.Anything).Return(&iam.User{UserName: "me@hello.com"}, nil) +} + +func setConfigHost(host string) func(*mocks.MockWorkspaceClient) { + return func(c *mocks.MockWorkspaceClient) { + c.WorkspaceClient.Config = &config.Config{ + Host: host, + } + } +} + +func setDefaultConfigHost(c *mocks.MockWorkspaceClient) { + c.WorkspaceClient.Config = &config.Config{ + Host: "900150983cd24fb0.cloud.databricks.com", + } +} + +func basicMockWorkspaceClients(t *testing.T, configs ...func(*mocks.MockWorkspaceClient)) func(map[int64]*mocks.MockWorkspaceClient) { + return func(m map[int64]*mocks.MockWorkspaceClient) { + c := mocks.NewMockWorkspaceClient(t) + for _, config := range configs { + config(c) + } + m[1234] = c + } +} + func TestResourceWorkspaceCreate(t *testing.T) { // Define a mock workspace that can be reused mockWorkspace := &provisioning.Workspace{ @@ -66,8 +95,10 @@ func TestResourceWorkspaceCreate(t *testing.T) { a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ WorkspaceId: 1234, }).Return(mockWorkspace, nil) + a.GetMockWorkspacesAPI().EXPECT().WaitGetWorkspaceRunning(mock.Anything, int64(1234), 20*time.Minute, mock.Anything).Return(mockWorkspace, nil) }, - Resource: ResourceMwsWorkspaces(), + MockWorkspaceClientsFunc: basicMockWorkspaceClients(t, setDefaultConfigHost, mockScimMe), + Resource: ResourceMwsWorkspaces(), State: map[string]any{ "account_id": "abc", "aws_region": "us-east-1", @@ -114,14 +145,15 @@ func TestResourceWorkspaceCreateGcp(t *testing.T) { qa.ResourceFixture{ MockAccountClientFunc: func(a *mocks.MockAccountClient) { a.GetMockWorkspacesAPI().EXPECT().Create(mock.Anything, provisioning.CreateWorkspaceRequest{ - Cloud: "gcp", CloudResourceContainer: &provisioning.CloudResourceContainer{ Gcp: &provisioning.CustomerFacingGcpCloudResourceContainer{ ProjectId: "def", }, }, - Location: "bcd", - NetworkId: "net_id_a", + DeploymentName: "900150983cd24fb0", + IsNoPublicIpEnabled: true, + Location: "bcd", + NetworkId: "net_id_a", GcpManagedNetworkConfig: &provisioning.GcpManagedNetworkConfig{ SubnetCidr: "a", }, @@ -130,8 +162,10 @@ func TestResourceWorkspaceCreateGcp(t *testing.T) { a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ WorkspaceId: 1234, }).Return(mockWorkspace, nil) + a.GetMockWorkspacesAPI().EXPECT().WaitGetWorkspaceRunning(mock.Anything, int64(1234), 20*time.Minute, mock.Anything).Return(mockWorkspace, nil) }, - Resource: ResourceMwsWorkspaces(), + MockWorkspaceClientsFunc: basicMockWorkspaceClients(t, setDefaultConfigHost, mockScimMe), + Resource: ResourceMwsWorkspaces(), HCL: ` account_id = "abc" workspace_name = "labdata" @@ -206,12 +240,13 @@ func TestResourceWorkspaceCreateGcpPsc(t *testing.T) { qa.ResourceFixture{ MockAccountClientFunc: func(a *mocks.MockAccountClient) { a.GetMockWorkspacesAPI().EXPECT().Create(mock.Anything, provisioning.CreateWorkspaceRequest{ - Cloud: "gcp", CloudResourceContainer: &provisioning.CloudResourceContainer{ Gcp: &provisioning.CustomerFacingGcpCloudResourceContainer{ ProjectId: "def", }, }, + DeploymentName: "900150983cd24fb0", + IsNoPublicIpEnabled: true, Location: "bcd", PrivateAccessSettingsId: "pas_id_a", NetworkId: "net_id_a", @@ -223,8 +258,11 @@ func TestResourceWorkspaceCreateGcpPsc(t *testing.T) { a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ WorkspaceId: 1234, }).Return(mockWorkspace, nil) + a.GetMockWorkspacesAPI().EXPECT().WaitGetWorkspaceRunning(mock.Anything, int64(1234), 20*time.Minute, mock.Anything).Return(mockWorkspace, nil) + }, - Resource: ResourceMwsWorkspaces(), + MockWorkspaceClientsFunc: basicMockWorkspaceClients(t, setDefaultConfigHost, mockScimMe), + Resource: ResourceMwsWorkspaces(), HCL: ` account_id = "abc" workspace_name = "labdata" @@ -274,12 +312,13 @@ func TestResourceWorkspaceCreateGcpCmk(t *testing.T) { qa.ResourceFixture{ MockAccountClientFunc: func(a *mocks.MockAccountClient) { a.GetMockWorkspacesAPI().EXPECT().Create(mock.Anything, provisioning.CreateWorkspaceRequest{ - Cloud: "gcp", CloudResourceContainer: &provisioning.CloudResourceContainer{ Gcp: &provisioning.CustomerFacingGcpCloudResourceContainer{ ProjectId: "def", }, }, + DeploymentName: "900150983cd24fb0", + IsNoPublicIpEnabled: true, Location: "bcd", PrivateAccessSettingsId: "pas_id_a", NetworkId: "net_id_a", @@ -293,8 +332,10 @@ func TestResourceWorkspaceCreateGcpCmk(t *testing.T) { a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ WorkspaceId: 1234, }).Return(mockWorkspace, nil) + a.GetMockWorkspacesAPI().EXPECT().WaitGetWorkspaceRunning(mock.Anything, int64(1234), 20*time.Minute, mock.Anything).Return(mockWorkspace, nil) }, - Resource: ResourceMwsWorkspaces(), + MockWorkspaceClientsFunc: basicMockWorkspaceClients(t, setDefaultConfigHost, mockScimMe), + Resource: ResourceMwsWorkspaces(), HCL: ` account_id = "abc" workspace_name = "labdata" @@ -358,8 +399,10 @@ func TestResourceWorkspaceCreateWithIsNoPublicIPEnabledFalse(t *testing.T) { a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ WorkspaceId: 1234, }).Return(mockWorkspace, nil) + a.GetMockWorkspacesAPI().EXPECT().WaitGetWorkspaceRunning(mock.Anything, int64(1234), 20*time.Minute, mock.Anything).Return(mockWorkspace, nil) }, - Resource: ResourceMwsWorkspaces(), + MockWorkspaceClientsFunc: basicMockWorkspaceClients(t, setDefaultConfigHost, mockScimMe), + Resource: ResourceMwsWorkspaces(), State: map[string]any{ "account_id": "abc", "aws_region": "us-east-1", @@ -416,8 +459,10 @@ func TestResourceWorkspaceCreateLegacyConfig(t *testing.T) { a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ WorkspaceId: 1234, }).Return(mockWorkspace, nil) + a.GetMockWorkspacesAPI().EXPECT().WaitGetWorkspaceRunning(mock.Anything, int64(1234), 20*time.Minute, mock.Anything).Return(mockWorkspace, nil) }, - Resource: ResourceMwsWorkspaces(), + MockWorkspaceClientsFunc: basicMockWorkspaceClients(t, setDefaultConfigHost, mockScimMe), + Resource: ResourceMwsWorkspaces(), State: map[string]any{ "account_id": "abc", "aws_region": "us-east-1", @@ -457,10 +502,11 @@ func TestResourceWorkspaceRead(t *testing.T) { }).Return(mockWorkspace, nil) a.GetMockWorkspacesAPI().EXPECT().WaitGetWorkspaceRunning(mock.Anything, int64(1234), 20*time.Minute, mock.Anything).Return(mockWorkspace, nil) }, - Resource: ResourceMwsWorkspaces(), - Read: true, - New: true, - ID: "abc/1234", + MockWorkspaceClientsFunc: basicMockWorkspaceClients(t, setDefaultConfigHost), + Resource: ResourceMwsWorkspaces(), + Read: true, + New: true, + ID: "abc/1234", }.Apply(t) assert.NoError(t, err) assert.Equal(t, "abc/1234", d.Id(), "Id should not be empty") @@ -499,7 +545,9 @@ func TestResourceWorkspaceRead_Issue382(t *testing.T) { a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ WorkspaceId: 1234, }).Return(mockWorkspace, nil) + a.GetMockWorkspacesAPI().EXPECT().WaitGetWorkspaceRunning(mock.Anything, int64(1234), 20*time.Minute, mock.Anything).Return(mockWorkspace, nil) }, + MockWorkspaceClientsFunc: basicMockWorkspaceClients(t, setConfigHost("prefix-900150983cd24fb0.cloud.databricks.com")), InstanceState: map[string]string{ "account_id": "abc", "aws_region": "us-east-1", @@ -597,16 +645,21 @@ func TestResourceWorkspaceUpdate(t *testing.T) { d, err := qa.ResourceFixture{ MockAccountClientFunc: func(a *mocks.MockAccountClient) { a.GetMockWorkspacesAPI().EXPECT().Update(mock.Anything, provisioning.UpdateWorkspaceRequest{ - WorkspaceId: 1234, - CredentialsId: "bcd", - NetworkId: "fgh", - StorageCustomerManagedKeyId: "def", + WorkspaceId: 1234, + AwsRegion: "us-east-1", + ManagedServicesCustomerManagedKeyId: "def", + CredentialsId: "bcd", + NetworkId: "fgh", + StorageCustomerManagedKeyId: "def", + StorageConfigurationId: "ghi", }).Return(mockWaiter, nil) a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ WorkspaceId: 1234, }).Return(mockWorkspace, nil) + a.GetMockWorkspacesAPI().EXPECT().WaitGetWorkspaceRunning(mock.Anything, int64(1234), 20*time.Minute, mock.Anything).Return(mockWorkspace, nil) }, - Resource: ResourceMwsWorkspaces(), + MockWorkspaceClientsFunc: basicMockWorkspaceClients(t, setDefaultConfigHost), + Resource: ResourceMwsWorkspaces(), InstanceState: map[string]string{ "account_id": "abc", "aws_region": "us-east-1", @@ -702,19 +755,24 @@ func TestResourceWorkspaceUpdateLegacyConfig(t *testing.T) { d, err := qa.ResourceFixture{ MockAccountClientFunc: func(a *mocks.MockAccountClient) { a.GetMockWorkspacesAPI().EXPECT().Update(mock.Anything, provisioning.UpdateWorkspaceRequest{ - WorkspaceId: 1234, - CredentialsId: "bcd", - NetworkId: "fgh", + WorkspaceId: 1234, + AwsRegion: "us-east-1", + ManagedServicesCustomerManagedKeyId: "def", + StorageConfigurationId: "ghi", + CredentialsId: "bcd", + NetworkId: "fgh", }).Return(mockWaiter, nil) a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ WorkspaceId: 1234, }).Return(mockWorkspace, nil) + a.GetMockWorkspacesAPI().EXPECT().WaitGetWorkspaceRunning(mock.Anything, int64(1234), 20*time.Minute, mock.Anything).Return(mockWorkspace, nil) }, - Resource: ResourceMwsWorkspaces(), + MockWorkspaceClientsFunc: basicMockWorkspaceClients(t, setDefaultConfigHost), + Resource: ResourceMwsWorkspaces(), InstanceState: map[string]string{ "account_id": "abc", "aws_region": "us-east-1", - "credentials_id": "bcd", + "credentials_id": "old", "customer_managed_key_id": "def", "deployment_name": "900150983cd24fb0", "is_no_public_ip_enabled": "true", @@ -745,10 +803,13 @@ func TestResourceWorkspaceUpdate_Error(t *testing.T) { qa.ResourceFixture{ MockAccountClientFunc: func(a *mocks.MockAccountClient) { a.GetMockWorkspacesAPI().EXPECT().Update(mock.Anything, provisioning.UpdateWorkspaceRequest{ - WorkspaceId: 1234, - CredentialsId: "bcd", - NetworkId: "fgh", - StorageCustomerManagedKeyId: "def", + WorkspaceId: 1234, + AwsRegion: "us-east-1", + ManagedServicesCustomerManagedKeyId: "def", + StorageConfigurationId: "ghi", + CredentialsId: "bcd", + NetworkId: "fgh", + StorageCustomerManagedKeyId: "def", }).Return(nil, &apierr.APIError{ ErrorCode: "INVALID_REQUEST", Message: "Internal error happened", @@ -756,6 +817,19 @@ func TestResourceWorkspaceUpdate_Error(t *testing.T) { }) }, Resource: ResourceMwsWorkspaces(), + InstanceState: map[string]string{ + "account_id": "abc", + "aws_region": "us-east-1", + "credentials_id": "old", + "managed_services_customer_managed_key_id": "def", + "storage_customer_managed_key_id": "def", + "is_no_public_ip_enabled": "true", + "deployment_name": "900150983cd24fb0", + "workspace_name": "labdata", + "network_id": "fgh", + "storage_configuration_id": "ghi", + "workspace_id": "1234", + }, State: map[string]any{ "account_id": "abc", "aws_region": "us-east-1", @@ -768,9 +842,8 @@ func TestResourceWorkspaceUpdate_Error(t *testing.T) { "storage_configuration_id": "ghi", "workspace_id": 1234, }, - Update: true, - RequiresNew: true, - ID: "abc/1234", + Update: true, + ID: "abc/1234", }.ExpectError(t, "Internal error happened") } @@ -822,8 +895,9 @@ func TestCreateFailsAndCleansUp(t *testing.T) { // Create a mock waiter mockWaiter := &provisioning.WaitGetWorkspaceRunning[provisioning.Workspace]{ WorkspaceId: 1234, + Response: mockFailedWorkspace, Poll: func(d time.Duration, f func(*provisioning.Workspace)) (*provisioning.Workspace, error) { - return mockFailedWorkspace, nil + return nil, errors.New("failed to reach RUNNING, got FAILED") }, } @@ -891,7 +965,7 @@ func TestCreateFailsAndCleansUp(t *testing.T) { }, Create: true, }.Apply(t) - assert.ErrorContains(t, err, "Workspace failed to create: Always fails, network error message: error: CREDENTIALS;error_msg: Message;") + assert.ErrorContains(t, err, "workspace status message: Always fails, network error message: error: credentials;error_msg: Message;") } func TestWorkspace_verifyWorkspaceReachable(t *testing.T) { @@ -955,7 +1029,7 @@ func TestEnsureTokenExists(t *testing.T) { err := ensureTokenExists(context.Background(), mockClient.WorkspaceClient, token) assert.NoError(t, err) assert.Equal(t, token.TokenID, "new-id") - assert.Equal(t, token.TokenValue, "new-value") + assert.Equal(t, token.TokenValue, SensitiveString("new-value")) } func TestEnsureTokenExists_NoRecreate(t *testing.T) { @@ -1003,7 +1077,7 @@ func TestExplainWorkspaceFailureCornerCase(t *testing.T) { NetworkId: "abc", WorkspaceStatusMessage: "🔥", } - assert.EqualError(t, explainWorkspaceFailure(context.Background(), mockClient.AccountClient, ws), "workspace status message: 🔥, network error message: cannot read network: 🐜") + assert.EqualError(t, explainWorkspaceFailure(context.Background(), mockClient.AccountClient, ws), "workspace status message: 🔥; network error message: cannot read network: 🐜") }) } @@ -1036,19 +1110,24 @@ func TestResourceWorkspaceUpdatePrivateAccessSettings(t *testing.T) { MockAccountClientFunc: func(a *mocks.MockAccountClient) { // Expect the Update call a.GetMockWorkspacesAPI().EXPECT().Update(mock.Anything, provisioning.UpdateWorkspaceRequest{ - WorkspaceId: 1234, - CredentialsId: "bcd", - NetworkId: "fgh", - StorageCustomerManagedKeyId: "def", - PrivateAccessSettingsId: "pas", + WorkspaceId: 1234, + AwsRegion: "us-east-1", + ManagedServicesCustomerManagedKeyId: "def", + StorageConfigurationId: "ghi", + CredentialsId: "bcd", + NetworkId: "fgh", + StorageCustomerManagedKeyId: "def", + PrivateAccessSettingsId: "pas", }).Return(mockWaiter, nil) // Expect the Get call to retrieve the updated workspace a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ WorkspaceId: 1234, }).Return(mockWorkspace, nil) + a.GetMockWorkspacesAPI().EXPECT().WaitGetWorkspaceRunning(mock.Anything, int64(1234), 20*time.Minute, mock.Anything).Return(mockWorkspace, nil) }, - Resource: ResourceMwsWorkspaces(), + MockWorkspaceClientsFunc: basicMockWorkspaceClients(t, setDefaultConfigHost), + Resource: ResourceMwsWorkspaces(), InstanceState: map[string]string{ "account_id": "abc", "aws_region": "us-east-1", @@ -1146,22 +1225,25 @@ func TestResourceWorkspaceCreateGcpManagedVPC(t *testing.T) { MockAccountClientFunc: func(a *mocks.MockAccountClient) { // Expect the Create call a.GetMockWorkspacesAPI().EXPECT().Create(mock.Anything, provisioning.CreateWorkspaceRequest{ - Cloud: "gcp", CloudResourceContainer: &provisioning.CloudResourceContainer{ Gcp: &provisioning.CustomerFacingGcpCloudResourceContainer{ ProjectId: "def", }, }, - Location: "bcd", - WorkspaceName: "labdata", + IsNoPublicIpEnabled: true, + DeploymentName: "900150983cd24fb0", + Location: "bcd", + WorkspaceName: "labdata", }).Return(mockWaiter, nil) // Expect the Get call to retrieve the workspace a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ WorkspaceId: 1234, }).Return(mockWorkspace, nil) + a.GetMockWorkspacesAPI().EXPECT().WaitGetWorkspaceRunning(mock.Anything, int64(1234), 20*time.Minute, mock.Anything).Return(mockWorkspace, nil) }, - Resource: ResourceMwsWorkspaces(), + MockWorkspaceClientsFunc: basicMockWorkspaceClients(t, setDefaultConfigHost, mockScimMe), + Resource: ResourceMwsWorkspaces(), HCL: ` account_id = "abc" workspace_name = "labdata" diff --git a/qa/testing.go b/qa/testing.go index 6f4d3a42be..7cb00f81a8 100644 --- a/qa/testing.go +++ b/qa/testing.go @@ -80,6 +80,8 @@ type ResourceFixture struct { MockWorkspaceClientFunc func(*mocks.MockWorkspaceClient) + MockWorkspaceClientsFunc func(map[int64]*mocks.MockWorkspaceClient) + MockAccountClientFunc func(*mocks.MockAccountClient) // The resource the unit test is testing. @@ -203,10 +205,10 @@ func (f ResourceFixture) setDatabricksEnvironmentForTest(client *common.Databric } func (f ResourceFixture) validateMocks() error { - isMockConfigured := f.MockAccountClientFunc != nil || f.MockWorkspaceClientFunc != nil + isMockConfigured := f.MockAccountClientFunc != nil || f.MockWorkspaceClientFunc != nil || f.MockWorkspaceClientsFunc != nil isFixtureConfigured := f.Fixtures != nil if isFixtureConfigured && isMockConfigured { - return fmt.Errorf("either (MockWorkspaceClientFunc, MockAccountClientFunc) or Fixtures may be set, not both") + return fmt.Errorf("either (MockWorkspaceClientFunc, MockWorkspaceClientsFunc, MockAccountClientFunc) or Fixtures may be set, not both") } return nil } @@ -230,10 +232,14 @@ func (f ResourceFixture) setupClient(t *testing.T) (*common.DatabricksClient, se return client, ss, err } mw := mocks.NewMockWorkspaceClient(t) + mws := map[int64]*mocks.MockWorkspaceClient{} ma := mocks.NewMockAccountClient(t) if f.MockWorkspaceClientFunc != nil { f.MockWorkspaceClientFunc(mw) } + if f.MockWorkspaceClientsFunc != nil { + f.MockWorkspaceClientsFunc(mws) + } if f.MockAccountClientFunc != nil { f.MockAccountClientFunc(ma) } @@ -244,6 +250,9 @@ func (f ResourceFixture) setupClient(t *testing.T) (*common.DatabricksClient, se } c.SetWorkspaceClient(mw.WorkspaceClient) c.SetAccountClient(ma.AccountClient) + for workspaceId, client := range mws { + c.SetWorkspaceClientForWorkspace(workspaceId, client.WorkspaceClient) + } c.Config.Credentials = testCredentialsProvider{token: token} return c, server{ Close: func() {}, From dd41a243bc8aa3d5ab8967f7f05816970f260241 Mon Sep 17 00:00:00 2001 From: Miles Yucht Date: Thu, 10 Apr 2025 21:34:32 +0200 Subject: [PATCH 05/16] fmt --- common/client.go | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/common/client.go b/common/client.go index b519529421..5fd4987d86 100644 --- a/common/client.go +++ b/common/client.go @@ -53,7 +53,7 @@ type DatabricksClient struct { *client.DatabricksClient // callback used to create API1.2 call wrapper, which simplifies unit testing - commandFactory func(context.Context, *DatabricksClient) CommandExecutor + commandFactory func(context.Context, *DatabricksClient) CommandExecutor // cachedWorkspaceClient is a cached workspace client authenticated to the workspace // configured for the provider cachedWorkspaceClient *databricks.WorkspaceClient @@ -63,10 +63,10 @@ type DatabricksClient struct { cachedWorkspaceClients map[int64]*databricks.WorkspaceClient // cachedAccountClient is a cached account client authenticated to the account // configured for the provider - cachedAccountClient *databricks.AccountClient + cachedAccountClient *databricks.AccountClient // mu synchronizes access to all cached clients. - mu sync.Mutex + mu sync.Mutex } // GetWorkspaceClient returns the Databricks WorkspaceClient or a diagnostics if that fails. From 7e6bdead0d0dd4f6535776d9157372c6f95fed5d Mon Sep 17 00:00:00 2001 From: Miles Yucht Date: Thu, 10 Apr 2025 22:12:04 +0200 Subject: [PATCH 06/16] work --- NEXT_CHANGELOG.md | 3 ++- mws/data_mws_workspaces_test.go | 4 ++-- mws/resource_mws_ncc_binding_test.go | 6 ------ 3 files changed, 4 insertions(+), 9 deletions(-) diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index af0619f01d..05e5638f7a 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -21,4 +21,5 @@ ### Internal Changes -* Add `TestMwsAccGcpWorkspaces` and `TestMwsAccGcpByovpcWorkspaces` to flaky test ([#4624](https://github.com/databricks/terraform-provider-databricks/pull/4624)). \ No newline at end of file +* Add `TestMwsAccGcpWorkspaces` and `TestMwsAccGcpByovpcWorkspaces` to flaky test ([#4624](https://github.com/databricks/terraform-provider-databricks/pull/4624)). +* Refactor `databricks_mws_workspaces` resource and datasource, as well as `databricks_mws_ncc_binding`, to use the Go SDK ([#4633](https://github.com/databricks/terraform-provider-databricks/pull/4633)). \ No newline at end of file diff --git a/mws/data_mws_workspaces_test.go b/mws/data_mws_workspaces_test.go index a08f6d157d..9f4a65be99 100755 --- a/mws/data_mws_workspaces_test.go +++ b/mws/data_mws_workspaces_test.go @@ -1,12 +1,12 @@ package mws import ( + "errors" "testing" "github.com/databricks/databricks-sdk-go/experimental/mocks" "github.com/databricks/databricks-sdk-go/service/provisioning" "github.com/databricks/terraform-provider-databricks/qa" - "github.com/stretchr/testify/assert" "github.com/stretchr/testify/mock" ) @@ -40,7 +40,7 @@ func TestDataSourceMwsWorkspaces(t *testing.T) { func TestDataSourceMwsWorkspaces_Error(t *testing.T) { qa.ResourceFixture{ MockAccountClientFunc: func(a *mocks.MockAccountClient) { - a.GetMockWorkspacesAPI().EXPECT().List(mock.Anything).Return(nil, assert.AnError) + a.GetMockWorkspacesAPI().EXPECT().List(mock.Anything).Return(nil, errors.New("i'm a teapot")) }, AccountID: "abc", Resource: DataSourceMwsWorkspaces(), diff --git a/mws/resource_mws_ncc_binding_test.go b/mws/resource_mws_ncc_binding_test.go index 9eb3c3d5d6..82bcaea30e 100644 --- a/mws/resource_mws_ncc_binding_test.go +++ b/mws/resource_mws_ncc_binding_test.go @@ -29,9 +29,6 @@ func TestResourceNccBindingCreate(t *testing.T) { WorkspaceId: 123456789, NetworkConnectivityConfigId: "ncc_id", }).Return(mockWaiter, nil) - a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ - WorkspaceId: 123456789, - }).Return(mockWorkspace, nil) }, Resource: ResourceMwsNccBinding(), AccountID: "abc", @@ -50,9 +47,6 @@ func TestResourceNccBindingUpdate(t *testing.T) { WorkspaceId: 123456789, NetworkConnectivityConfigId: "new_ncc_id", }).Return(mockWaiter, nil) - a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ - WorkspaceId: 123456789, - }).Return(mockWorkspace, nil) }, Resource: ResourceMwsNccBinding(), AccountID: "abc", From 40265e408763736d15865aaceeee52ca3da31aa6 Mon Sep 17 00:00:00 2001 From: Miles Yucht Date: Fri, 11 Apr 2025 07:37:01 +0200 Subject: [PATCH 07/16] fix test --- internal/acceptance/expect_not_destroyed.go | 29 +++++++++++++++ mws/mws_workspaces_test.go | 39 ++++++++------------- mws/resource_mws_workspaces.go | 1 + 3 files changed, 44 insertions(+), 25 deletions(-) create mode 100644 internal/acceptance/expect_not_destroyed.go diff --git a/internal/acceptance/expect_not_destroyed.go b/internal/acceptance/expect_not_destroyed.go new file mode 100644 index 0000000000..430437e035 --- /dev/null +++ b/internal/acceptance/expect_not_destroyed.go @@ -0,0 +1,29 @@ +package acceptance + +import ( + "context" + "fmt" + + "github.com/hashicorp/terraform-plugin-testing/plancheck" +) + +type expectNotDestroyed struct { + addr string +} + +func ExpectNotDestroyed(addr string) expectNotDestroyed { + return expectNotDestroyed{addr: addr} +} + +func (e expectNotDestroyed) CheckPlan(ctx context.Context, req plancheck.CheckPlanRequest, resp *plancheck.CheckPlanResponse) { + for _, resource := range req.Plan.ResourceChanges { + if resource.Address != e.addr { + continue + } + actions := resource.Change.Actions + if actions.DestroyBeforeCreate() || actions.CreateBeforeDestroy() || actions.Delete() { + resp.Error = fmt.Errorf("resource %s is marked for destruction", e.addr) + return + } + } +} diff --git a/mws/mws_workspaces_test.go b/mws/mws_workspaces_test.go index d01367e411..6e49d84c29 100644 --- a/mws/mws_workspaces_test.go +++ b/mws/mws_workspaces_test.go @@ -78,27 +78,6 @@ func TestMwsAccWorkspaces(t *testing.T) { }) } -type expectNotDestroyed struct { - addr string -} - -func ExpectNotDestroyed(addr string) expectNotDestroyed { - return expectNotDestroyed{addr: addr} -} - -func (e expectNotDestroyed) CheckPlan(ctx context.Context, req plancheck.CheckPlanRequest, resp *plancheck.CheckPlanResponse) { - for _, resource := range req.Plan.ResourceChanges { - if resource.Address != e.addr { - continue - } - actions := resource.Change.Actions - if actions.DestroyBeforeCreate() || actions.CreateBeforeDestroy() || actions.Delete() { - resp.Error = fmt.Errorf("resource %s is marked for destruction", e.addr) - return - } - } -} - func TestMwsAccWorkspaces_TokenUpdate(t *testing.T) { tokenUpdateTemplate := func(token, customTags string) string { tokenBlock := `` @@ -182,7 +161,7 @@ func TestMwsAccWorkspaces_TokenUpdate(t *testing.T) { // Updating the comment causes the old token to be deleted and a new one to be created Template: tokenUpdateTemplate(`comment = "test bar"`, ""), ConfigPlanChecks: resource.ConfigPlanChecks{ - PreApply: []plancheck.PlanCheck{ExpectNotDestroyed("databricks_mws_workspaces.this")}, + PreApply: []plancheck.PlanCheck{acceptance.ExpectNotDestroyed("databricks_mws_workspaces.this")}, }, Check: resource.ComposeAggregateTestCheckFunc( checkTokenExists, @@ -201,7 +180,7 @@ func TestMwsAccWorkspaces_TokenUpdate(t *testing.T) { // Modifying the tags doesn't change the token but does modify the workspace. Template: tokenUpdateTemplate(`comment = "test bar"`, `"Key" = "Value"`), ConfigPlanChecks: resource.ConfigPlanChecks{ - PreApply: []plancheck.PlanCheck{ExpectNotDestroyed("databricks_mws_workspaces.this")}, + PreApply: []plancheck.PlanCheck{acceptance.ExpectNotDestroyed("databricks_mws_workspaces.this")}, }, Check: func(s *terraform.State) error { state, ok := s.RootModule().Resources["databricks_mws_workspaces.this"] @@ -217,7 +196,7 @@ func TestMwsAccWorkspaces_TokenUpdate(t *testing.T) { // It is also possible to modify the token comment and tags at the same time. Template: tokenUpdateTemplate(`comment = "test quux"`, `"Key" = "Value2"`), ConfigPlanChecks: resource.ConfigPlanChecks{ - PreApply: []plancheck.PlanCheck{ExpectNotDestroyed("databricks_mws_workspaces.this")}, + PreApply: []plancheck.PlanCheck{acceptance.ExpectNotDestroyed("databricks_mws_workspaces.this")}, }, Check: func(s *terraform.State) error { state, ok := s.RootModule().Resources["databricks_mws_workspaces.this"] @@ -327,6 +306,7 @@ func TestMwsAccGcpPscWorkspaces(t *testing.T) { } func TestMwsAccAwsChangeToServicePrincipal(t *testing.T) { + acceptance.LoadDebugEnvIfRunsFromIDE(t, "account") if !acceptance.IsAws(t) { acceptance.Skipf(t)("TestMwsAccAwsChangeToServicePrincipal should only run on AWS") } @@ -442,17 +422,26 @@ func TestMwsAccAwsChangeToServicePrincipal(t *testing.T) { // Tolerate existing token Template: workspaceTemplate(`token { comment = "Test {var.STICKY_RANDOM}" }`) + servicePrincipal, ProtoV6ProviderFactories: providerFactory, + ConfigPlanChecks: resource.ConfigPlanChecks{ + PreApply: []plancheck.PlanCheck{acceptance.ExpectNotDestroyed("databricks_mws_workspaces.this")}, + }, }, acceptance.Step{ // Allow the token to be removed Template: workspaceTemplate(``) + servicePrincipal, ProtoV6ProviderFactories: providerFactory, + ConfigPlanChecks: resource.ConfigPlanChecks{ + PreApply: []plancheck.PlanCheck{acceptance.ExpectNotDestroyed("databricks_mws_workspaces.this")}, + }, }, acceptance.Step{ // Fail when adding the token back Template: workspaceTemplate(`token { comment = "Test {var.STICKY_RANDOM}" }`) + servicePrincipal, ProtoV6ProviderFactories: providerFactory, ExpectError: regexp.MustCompile(`cannot create token: the principal used by Databricks \(client ID .*\) is not authorized to create a token in this workspace`), + ConfigPlanChecks: resource.ConfigPlanChecks{ + PreApply: []plancheck.PlanCheck{acceptance.ExpectNotDestroyed("databricks_mws_workspaces.this")}, + }, }, acceptance.Step{ // Use the original provider for a final step to clean up the newly created service principal - Template: workspaceTemplate(``) + servicePrincipal, + Template: workspaceTemplate(`token { comment = "Test {var.STICKY_RANDOM}" }`) + servicePrincipal, }) } diff --git a/mws/resource_mws_workspaces.go b/mws/resource_mws_workspaces.go index 2731f8e7d5..3b596d373d 100644 --- a/mws/resource_mws_workspaces.go +++ b/mws/resource_mws_workspaces.go @@ -170,6 +170,7 @@ func ResourceMwsWorkspaces() common.Resource { "private_access_settings_id": {}, "managed_services_customer_managed_key_id": {}, "custom_tags": {}, + "token": {}, } workspaceSchema := common.StructToSchema(Workspace{}, From 50bb2790d1081eb7f46718e29cd4bbae1b02a036 Mon Sep 17 00:00:00 2001 From: Miles Yucht Date: Fri, 11 Apr 2025 09:56:45 +0200 Subject: [PATCH 08/16] work --- common/client.go | 1 - common/client_test.go | 98 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 98 insertions(+), 1 deletion(-) diff --git a/common/client.go b/common/client.go index 5fd4987d86..36ea4506c3 100644 --- a/common/client.go +++ b/common/client.go @@ -123,7 +123,6 @@ func (c *DatabricksClient) WorkspaceClientForWorkspace(ctx context.Context, work return nil, err } w.CurrentUser = newCachedMe(w.CurrentUser) - w.Config.AzureResourceID = workspace.AzureResourceId() c.cachedWorkspaceClients[workspace.WorkspaceId] = w return w, nil } diff --git a/common/client_test.go b/common/client_test.go index a43b654535..250a77514a 100644 --- a/common/client_test.go +++ b/common/client_test.go @@ -2,16 +2,21 @@ package common import ( "context" + "fmt" "log" "net/http" "path/filepath" "strings" "testing" + "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/client" "github.com/databricks/databricks-sdk-go/config" + "github.com/databricks/databricks-sdk-go/experimental/mocks" "github.com/databricks/databricks-sdk-go/service/iam" + "github.com/databricks/databricks-sdk-go/service/provisioning" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" "github.com/stretchr/testify/require" ) @@ -335,3 +340,96 @@ func TestCachedMe_Me_MakesSingleRequest(t *testing.T) { cm.Me(context.Background()) assert.Equal(t, 1, mock.count) } + +func TestWorkspaceClientForWorkspace_WorkspaceDoesNotExist(t *testing.T) { + mockAcc := mocks.NewMockAccountClient(t) + mockWorkspacesAPI := mockAcc.GetMockWorkspacesAPI() + + // Setup the mock to return an error for non-existent workspace + mockWorkspacesAPI.EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ + WorkspaceId: 12345, + }).Return(nil, fmt.Errorf("workspace not found")) + + // Create a DatabricksClient with the mock account client + dc := &DatabricksClient{ + DatabricksClient: &client.DatabricksClient{ + Config: &config.Config{}, + }, + } + dc.SetAccountClient(mockAcc.AccountClient) + + // Call the method with a non-existent workspace ID + _, err := dc.WorkspaceClientForWorkspace(context.Background(), 12345) + + // Verify the error + assert.Error(t, err) + assert.Contains(t, err.Error(), "workspace not found") +} + +func TestWorkspaceClientForWorkspace_WorkspaceExistsNotInCache(t *testing.T) { + mockAcc := mocks.NewMockAccountClient(t) + mockAcc.AccountClient.Config = &config.Config{ + Token: "dapi123", // Instantiating WorkspaceClient attempts authentication, this allows Configure() to complete quickly. + } + mockWorkspacesAPI := mockAcc.GetMockWorkspacesAPI() + + // Create a mock workspace + mockWorkspace := &provisioning.Workspace{ + WorkspaceId: 12345, + WorkspaceName: "test-workspace", + DeploymentName: "test-deployment", + } + + // Setup the mock to return the workspace + mockWorkspacesAPI.EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ + WorkspaceId: 12345, + }).Return(mockWorkspace, nil) + + // Create a DatabricksClient with the mock account client + dc := &DatabricksClient{ + DatabricksClient: &client.DatabricksClient{ + Config: &config.Config{}, + }, + } + dc.SetAccountClient(mockAcc.AccountClient) + + // Call the method with the workspace ID + workspaceClient, err := dc.WorkspaceClientForWorkspace(context.Background(), 12345) + + // Verify no error and client is returned + assert.NoError(t, err) + assert.NotNil(t, workspaceClient) + + // Verify the workspace client is configured correctly + assert.Equal(t, fmt.Sprintf("https://%s.cloud.databricks.com", mockWorkspace.DeploymentName), workspaceClient.Config.Host) + + // Verify the client is cached + dc.mu.Lock() + cachedClient, exists := dc.cachedWorkspaceClients[12345] + dc.mu.Unlock() + + assert.True(t, exists) + assert.Equal(t, workspaceClient, cachedClient) +} + +func TestWorkspaceClientForWorkspace_WorkspaceExistsInCache(t *testing.T) { + // Create a mock workspace client + mockWorkspaceClient := &databricks.WorkspaceClient{} + + // Create a DatabricksClient with the mock workspace client in cache + dc := &DatabricksClient{ + DatabricksClient: &client.DatabricksClient{ + Config: &config.Config{}, + }, + } + + // Set the workspace client in cache + dc.SetWorkspaceClientForWorkspace(12345, mockWorkspaceClient) + + // Call the method with the workspace ID + workspaceClient, err := dc.WorkspaceClientForWorkspace(context.Background(), 12345) + + // Verify no error and the cached client is returned + assert.NoError(t, err) + assert.Equal(t, mockWorkspaceClient, workspaceClient) +} From 78c0eb5954d5e93f8d814ce412c6a1cb04e7dd33 Mon Sep 17 00:00:00 2001 From: Miles Yucht Date: Fri, 11 Apr 2025 11:13:26 +0200 Subject: [PATCH 09/16] Fix schema diffs --- mws/resource_mws_workspaces.go | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/mws/resource_mws_workspaces.go b/mws/resource_mws_workspaces.go index 3b596d373d..01feb5cef0 100644 --- a/mws/resource_mws_workspaces.go +++ b/mws/resource_mws_workspaces.go @@ -154,6 +154,7 @@ func removeToken(ctx context.Context, w *databricks.WorkspaceClient, tokenID str // ResourceMwsWorkspaces manages E2 workspaces func ResourceMwsWorkspaces() common.Resource { var computedFields = map[string]struct{}{ + "cloud": {}, "workspace_id": {}, "workspace_url": {}, "workspace_status": {}, @@ -184,6 +185,7 @@ func ResourceMwsWorkspaces() common.Resource { } } s["account_id"].Sensitive = true + common.CustomizeSchemaPath(s, "account_id").SetRequired() s["deployment_name"].DiffSuppressFunc = func(k, old, new string, d *schema.ResourceData) bool { if old == "" && new != "" { return false @@ -215,6 +217,12 @@ func ResourceMwsWorkspaces() common.Resource { common.CustomizeSchemaPath(s, "gke_config").SetDeprecated(getGkeDeprecationMessage("gke_config", docOptions)) common.CustomizeSchemaPath(s, "gcp_managed_network_config", "gke_cluster_pod_ip_range").SetDeprecated(getGkeDeprecationMessage("gcp_managed_network_config.gke_cluster_pod_ip_range", docOptions)) common.CustomizeSchemaPath(s, "gcp_managed_network_config", "gke_cluster_service_ip_range").SetDeprecated(getGkeDeprecationMessage("gcp_managed_network_config.gke_cluster_service_ip_range", docOptions)) + common.CustomizeSchemaPath(s, "gcp_managed_network_config", "subnet_cidr").SetRequired() + common.CustomizeSchemaPath(s, "workspace_name").SetRequired() + common.CustomizeSchemaPath(s, "cloud_resource_container", "gcp", "project_id").SetRequired() + for _, field := range []string{"authoritative_user_email", "authoritative_user_full_name", "customer_name"} { + common.CustomizeSchemaPath(s, "external_customer_info", field).SetRequired() + } return s }) p := common.NewPairSeparatedID("account_id", "workspace_id", "/").Schema( From e7f2dfdec29b944dd38e0284da71d34155e3b4c0 Mon Sep 17 00:00:00 2001 From: Miles Yucht Date: Fri, 11 Apr 2025 11:15:23 +0200 Subject: [PATCH 10/16] work --- mws/resource_mws_workspaces.go | 1 + 1 file changed, 1 insertion(+) diff --git a/mws/resource_mws_workspaces.go b/mws/resource_mws_workspaces.go index 01feb5cef0..83284f300a 100644 --- a/mws/resource_mws_workspaces.go +++ b/mws/resource_mws_workspaces.go @@ -219,6 +219,7 @@ func ResourceMwsWorkspaces() common.Resource { common.CustomizeSchemaPath(s, "gcp_managed_network_config", "gke_cluster_service_ip_range").SetDeprecated(getGkeDeprecationMessage("gcp_managed_network_config.gke_cluster_service_ip_range", docOptions)) common.CustomizeSchemaPath(s, "gcp_managed_network_config", "subnet_cidr").SetRequired() common.CustomizeSchemaPath(s, "workspace_name").SetRequired() + common.CustomizeSchemaPath(s, "cloud_resource_container", "gcp").SetMinItems(1) common.CustomizeSchemaPath(s, "cloud_resource_container", "gcp", "project_id").SetRequired() for _, field := range []string{"authoritative_user_email", "authoritative_user_full_name", "customer_name"} { common.CustomizeSchemaPath(s, "external_customer_info", field).SetRequired() From 169b077ee1da46f3e7bf57379dd24a007e4e21c1 Mon Sep 17 00:00:00 2001 From: Miles Yucht Date: Fri, 11 Apr 2025 11:38:10 +0200 Subject: [PATCH 11/16] fix test --- mws/resource_mws_workspaces_test.go | 1 + 1 file changed, 1 insertion(+) diff --git a/mws/resource_mws_workspaces_test.go b/mws/resource_mws_workspaces_test.go index 0982866d76..a453afc618 100644 --- a/mws/resource_mws_workspaces_test.go +++ b/mws/resource_mws_workspaces_test.go @@ -290,6 +290,7 @@ func TestResourceWorkspaceCreateGcpCmk(t *testing.T) { WorkspaceId: 1234, WorkspaceStatus: provisioning.WorkspaceStatusRunning, WorkspaceName: "labdata", + AccountId: "abc", DeploymentName: "900150983cd24fb0", Cloud: "gcp", Location: "bcd", From a399087a37fa84e531e2d5ce1438dae83723ecd9 Mon Sep 17 00:00:00 2001 From: Miles Yucht Date: Fri, 11 Apr 2025 11:48:01 +0200 Subject: [PATCH 12/16] fix --- mws/resource_mws_workspaces.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mws/resource_mws_workspaces.go b/mws/resource_mws_workspaces.go index 83284f300a..c997740b9e 100644 --- a/mws/resource_mws_workspaces.go +++ b/mws/resource_mws_workspaces.go @@ -219,7 +219,7 @@ func ResourceMwsWorkspaces() common.Resource { common.CustomizeSchemaPath(s, "gcp_managed_network_config", "gke_cluster_service_ip_range").SetDeprecated(getGkeDeprecationMessage("gcp_managed_network_config.gke_cluster_service_ip_range", docOptions)) common.CustomizeSchemaPath(s, "gcp_managed_network_config", "subnet_cidr").SetRequired() common.CustomizeSchemaPath(s, "workspace_name").SetRequired() - common.CustomizeSchemaPath(s, "cloud_resource_container", "gcp").SetMinItems(1) + common.CustomizeSchemaPath(s, "cloud_resource_container", "gcp").SetRequired() common.CustomizeSchemaPath(s, "cloud_resource_container", "gcp", "project_id").SetRequired() for _, field := range []string{"authoritative_user_email", "authoritative_user_full_name", "customer_name"} { common.CustomizeSchemaPath(s, "external_customer_info", field).SetRequired() From cebb9f39e82396dfeed5273f2421bac4b2a08731 Mon Sep 17 00:00:00 2001 From: Miles Yucht Date: Fri, 11 Apr 2025 14:00:40 +0200 Subject: [PATCH 13/16] empty commit From 74685dcd2c639a4741a4893422dcacaaf3ef9d51 Mon Sep 17 00:00:00 2001 From: Miles Yucht Date: Mon, 14 Apr 2025 09:54:49 +0000 Subject: [PATCH 14/16] some fixes --- docs/resources/mws_workspaces.md | 18 ++++++++++------ mws/resource_mws_workspaces.go | 37 +++++++++++++++++++------------- 2 files changed, 33 insertions(+), 22 deletions(-) diff --git a/docs/resources/mws_workspaces.md b/docs/resources/mws_workspaces.md index a252302da1..df833388d9 100644 --- a/docs/resources/mws_workspaces.md +++ b/docs/resources/mws_workspaces.md @@ -323,6 +323,7 @@ The following arguments are available: * `deployment_name` - (Optional) part of URL as in `https://-.cloud.databricks.com`. Deployment name cannot be used until a deployment name prefix is defined. Please contact your Databricks representative. Once a new deployment prefix is added/updated, it only will affect the new workspaces created. * `workspace_name` - name of the workspace, will appear on UI. * `network_id` - (Optional) `network_id` from [networks](mws_networks.md). +* `credentials_id` - (AWS only) ID of the workspace's credential configuration object. * `aws_region` - (AWS only) region of VPC. * `storage_configuration_id` - (AWS only)`storage_configuration_id` from [storage configuration](mws_storage_configurations.md). * `managed_services_customer_managed_key_id` - (Optional) `customer_managed_key_id` from [customer managed keys](mws_customer_managed_keys.md) with `use_cases` set to `MANAGED_SERVICES`. This is used to encrypt the workspace's notebook and secret data in the control plane. @@ -331,7 +332,9 @@ The following arguments are available: * `cloud_resource_container` - (GCP only) A block that specifies GCP workspace configurations, consisting of following blocks: * `gcp` - A block that consists of the following field: * `project_id` - The Google Cloud project ID, which the workspace uses to instantiate cloud resources for your workspace. -* `gke_config` - (GCP only) A block that specifies GKE configuration for the Databricks workspace: +* `gcp_managed_network_config` - (GCP only) A block that describes the network configuration for workspaces with Databricks-managed networks. + * `subnet_cidr` - The IP range from which to allocate GKE cluster nodes. No bigger than `/9` and no smaller than `/29`. +* `gke_config` - (GCP only, deprecated) A block that specifies GKE configuration for the Databricks workspace: * `connectivity_type`: Specifies the network connectivity types for the GKE nodes and the GKE master network. Possible values are: `PRIVATE_NODE_PUBLIC_MASTER`, `PUBLIC_NODE_PUBLIC_MASTER`. * `master_ip_range`: The IP range from which to allocate GKE cluster master resources. This field will be ignored if GKE private cluster is not enabled. It must be exactly as big as `/28`. * `private_access_settings_id` - (Optional) Canonical unique identifier of [databricks_mws_private_access_settings](mws_private_access_settings.md) in Databricks Account. @@ -351,24 +354,25 @@ You can specify a `token` block in the body of the workspace resource, so that T On AWS, the following arguments could be modified after the workspace is running: -* `network_id` - Modifying [networks on running workspaces](mws_networks.md#modifying-networks-on-running-workspaces-aws-only) would require three separate `terraform apply` steps. * `credentials_id` -* `storage_customer_managed_key_id` -* `private_access_settings_id` * `custom_tags` +* `managed_services_customer_managed_key_id` +* `network_id` - Modifying [networks on running workspaces](mws_networks.md#modifying-networks-on-running-workspaces-aws-only) would require three separate `terraform apply` steps. +* `private_access_settings_id` +* `storage_customer_managed_key_id` ## Attribute Reference In addition to all arguments above, the following attributes are exported: * `id` - (String) Canonical unique identifier for the workspace, of the format `/` +* `creation_time` - (Integer) time when workspace was created +* `custom_tags` - (Map) Custom Tags (if present) added to workspace +* `gcp_workspace_sa` - (String, GCP only) identifier of a service account created for the workspace in form of `db-@prod-gcp-.iam.gserviceaccount.com` * `workspace_id` - (String) workspace id * `workspace_status_message` - (String) updates on workspace status * `workspace_status` - (String) workspace status -* `creation_time` - (Integer) time when workspace was created * `workspace_url` - (String) URL of the workspace -* `custom_tags` - (Map) Custom Tags (if present) added to workspace -* `gcp_workspace_sa` - (String, GCP only) identifier of a service account created for the workspace in form of `db-@prod-gcp-.iam.gserviceaccount.com` ## Timeouts diff --git a/mws/resource_mws_workspaces.go b/mws/resource_mws_workspaces.go index c997740b9e..a42b62cc7f 100644 --- a/mws/resource_mws_workspaces.go +++ b/mws/resource_mws_workspaces.go @@ -28,8 +28,12 @@ import ( // this may help with local DNS cache issues. const DefaultProvisionTimeout = 20 * time.Minute -// Workspace is the object that contains all the information for deploying a workspace -type Workspace struct { +// workspaceReachableRequestTimeout is the amount of seconds to wait when polling the +// newly created workspace's SCIM Me endpoint to check if the workspace is reachable. +const workspaceReachableRequestTimeout = 10 * time.Second + +// workspace is the object that contains all the information for deploying a workspace +type workspace struct { provisioning.Workspace CustomerManagedKeyID string `json:"customer_managed_key_id,omitempty"` // just for compatibility, will be removed GcpWorkspaceSa string `json:"gcp_workspace_sa" tf:"computed"` @@ -49,7 +53,7 @@ func verifyWorkspaceReachable(ctx context.Context, w *databricks.WorkspaceClient } return false })).Wait(ctx, func(ctx context.Context) error { - ctx, cancel := context.WithTimeout(ctx, 10*time.Second) + ctx, cancel := context.WithTimeout(ctx, workspaceReachableRequestTimeout) defer cancel() _, err := w.CurrentUser.Me(ctx) return err @@ -154,27 +158,27 @@ func removeToken(ctx context.Context, w *databricks.WorkspaceClient, tokenID str // ResourceMwsWorkspaces manages E2 workspaces func ResourceMwsWorkspaces() common.Resource { var computedFields = map[string]struct{}{ - "cloud": {}, - "workspace_id": {}, - "workspace_url": {}, - "workspace_status": {}, - "workspace_status_message": {}, + "azure_workspace_info": {}, "creation_time": {}, "gke_config": {}, "pricing_tier": {}, + "workspace_id": {}, + "workspace_status": {}, + "workspace_status_message": {}, + "workspace_url": {}, } var workspaceRunningUpdatesAllowed = map[string]struct{}{ "credentials_id": {}, + "custom_tags": {}, + "managed_services_customer_managed_key_id": {}, "network_id": {}, - "storage_customer_managed_key_id": {}, "private_access_settings_id": {}, - "managed_services_customer_managed_key_id": {}, - "custom_tags": {}, + "storage_customer_managed_key_id": {}, "token": {}, } - workspaceSchema := common.StructToSchema(Workspace{}, + workspaceSchema := common.StructToSchema(workspace{}, func(s map[string]*schema.Schema) map[string]*schema.Schema { for name, fieldSchema := range s { if _, ok := computedFields[name]; ok { @@ -258,7 +262,7 @@ func ResourceMwsWorkspaces() common.Resource { } var createWorkspaceRequest provisioning.CreateWorkspaceRequest common.DataToStructPointer(d, workspaceSchema, &createWorkspaceRequest) - var workspaceConfig Workspace + var workspaceConfig workspace common.DataToStructPointer(d, workspaceSchema, &workspaceConfig) // Validate required fields by cloud @@ -275,6 +279,9 @@ func ResourceMwsWorkspaces() common.Resource { log.Print("[INFO] Using existing customer_managed_key_id as value for new managed_services_customer_managed_key_id") createWorkspaceRequest.ManagedServicesCustomerManagedKeyId = customerManagedKeyId.(string) } + // is_no_public_ip_enabled always needs to be serialized. It is true by default, but if it is disabled + // by a customer, it will be omitted by the default request serializer. + createWorkspaceRequest.ForceSendFields = append(createWorkspaceRequest.ForceSendFields, "IsNoPublicIpEnabled") // Create the workspace. If creation fails, clean it up and return. wait, err := a.Workspaces.Create(ctx, createWorkspaceRequest) @@ -332,7 +339,7 @@ func ResourceMwsWorkspaces() common.Resource { if err != nil { return err } - var workspaceConfig Workspace + var workspaceConfig workspace common.DataToStructPointer(d, workspaceSchema, &workspaceConfig) workspace, err := a.Workspaces.Get(ctx, provisioning.GetWorkspaceRequest{WorkspaceId: common.MustInt64(workspaceID)}) if err != nil { @@ -370,7 +377,7 @@ func ResourceMwsWorkspaces() common.Resource { } var updateWorkspaceRequest provisioning.UpdateWorkspaceRequest common.DataToStructPointer(d, workspaceSchema, &updateWorkspaceRequest) - var workspaceConfig Workspace + var workspaceConfig workspace common.DataToStructPointer(d, workspaceSchema, &workspaceConfig) // WorkspaceId in UpdateWorkspaceRequest is a path parameter, thus tagged with `json:"-"`. // This causes it not to be set in DataToStructPointer. Instead, the workspace ID can be From 77621c643079a2dc42e8d3772e580ec4e9f70734 Mon Sep 17 00:00:00 2001 From: Miles Yucht Date: Mon, 14 Apr 2025 13:31:17 +0200 Subject: [PATCH 15/16] work --- mws/resource_mws_workspaces.go | 12 ++++++------ mws/resource_mws_workspaces_test.go | 12 ++++++++++-- 2 files changed, 16 insertions(+), 8 deletions(-) diff --git a/mws/resource_mws_workspaces.go b/mws/resource_mws_workspaces.go index a42b62cc7f..340dd1c95b 100644 --- a/mws/resource_mws_workspaces.go +++ b/mws/resource_mws_workspaces.go @@ -169,13 +169,13 @@ func ResourceMwsWorkspaces() common.Resource { } var workspaceRunningUpdatesAllowed = map[string]struct{}{ - "credentials_id": {}, - "custom_tags": {}, + "credentials_id": {}, + "custom_tags": {}, "managed_services_customer_managed_key_id": {}, - "network_id": {}, - "private_access_settings_id": {}, - "storage_customer_managed_key_id": {}, - "token": {}, + "network_id": {}, + "private_access_settings_id": {}, + "storage_customer_managed_key_id": {}, + "token": {}, } workspaceSchema := common.StructToSchema(workspace{}, diff --git a/mws/resource_mws_workspaces_test.go b/mws/resource_mws_workspaces_test.go index a453afc618..348af5da52 100644 --- a/mws/resource_mws_workspaces_test.go +++ b/mws/resource_mws_workspaces_test.go @@ -91,6 +91,7 @@ func TestResourceWorkspaceCreate(t *testing.T) { CustomTags: map[string]string{ "SoldToCode": "1234", }, + ForceSendFields: []string{"IsNoPublicIpEnabled"}, }).Return(mockWaiter, nil) a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ WorkspaceId: 1234, @@ -157,7 +158,8 @@ func TestResourceWorkspaceCreateGcp(t *testing.T) { GcpManagedNetworkConfig: &provisioning.GcpManagedNetworkConfig{ SubnetCidr: "a", }, - WorkspaceName: "labdata", + WorkspaceName: "labdata", + ForceSendFields: []string{"IsNoPublicIpEnabled"}, }).Return(mockWaiter, nil) a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ WorkspaceId: 1234, @@ -253,7 +255,8 @@ func TestResourceWorkspaceCreateGcpPsc(t *testing.T) { GcpManagedNetworkConfig: &provisioning.GcpManagedNetworkConfig{ SubnetCidr: "a", }, - WorkspaceName: "labdata", + WorkspaceName: "labdata", + ForceSendFields: []string{"IsNoPublicIpEnabled"}, }).Return(mockWaiter, nil) a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ WorkspaceId: 1234, @@ -329,6 +332,7 @@ func TestResourceWorkspaceCreateGcpCmk(t *testing.T) { WorkspaceName: "labdata", ManagedServicesCustomerManagedKeyId: "managed_services_cmk", StorageCustomerManagedKeyId: "storage_cmk", + ForceSendFields: []string{"IsNoPublicIpEnabled"}, }).Return(mockWaiter, nil) a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ WorkspaceId: 1234, @@ -396,6 +400,7 @@ func TestResourceWorkspaceCreateWithIsNoPublicIPEnabledFalse(t *testing.T) { NetworkId: "fgh", ManagedServicesCustomerManagedKeyId: "def", StorageCustomerManagedKeyId: "def", + ForceSendFields: []string{"IsNoPublicIpEnabled"}, }).Return(mockWaiter, nil) a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ WorkspaceId: 1234, @@ -456,6 +461,7 @@ func TestResourceWorkspaceCreateLegacyConfig(t *testing.T) { StorageConfigurationId: "ghi", NetworkId: "fgh", ManagedServicesCustomerManagedKeyId: "def", + ForceSendFields: []string{"IsNoPublicIpEnabled"}, }).Return(mockWaiter, nil) a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ WorkspaceId: 1234, @@ -926,6 +932,7 @@ func TestCreateFailsAndCleansUp(t *testing.T) { NetworkId: "fgh", ManagedServicesCustomerManagedKeyId: "def", StorageCustomerManagedKeyId: "def", + ForceSendFields: []string{"IsNoPublicIpEnabled"}, }).Return(mockWaiter, nil) // Expect the Get call to retrieve the failed workspace @@ -1235,6 +1242,7 @@ func TestResourceWorkspaceCreateGcpManagedVPC(t *testing.T) { DeploymentName: "900150983cd24fb0", Location: "bcd", WorkspaceName: "labdata", + ForceSendFields: []string{"IsNoPublicIpEnabled"}, }).Return(mockWaiter, nil) // Expect the Get call to retrieve the workspace From e0d9a7444815c5753f78c04397af3daaf000aea9 Mon Sep 17 00:00:00 2001 From: Miles Yucht Date: Mon, 14 Apr 2025 14:34:46 +0200 Subject: [PATCH 16/16] add tests --- mws/resource_mws_workspaces.go | 4 + mws/resource_mws_workspaces_test.go | 402 +++++++++++++++++++++++----- 2 files changed, 339 insertions(+), 67 deletions(-) diff --git a/mws/resource_mws_workspaces.go b/mws/resource_mws_workspaces.go index 340dd1c95b..e9ec0e78ab 100644 --- a/mws/resource_mws_workspaces.go +++ b/mws/resource_mws_workspaces.go @@ -149,6 +149,10 @@ func removeToken(ctx context.Context, w *databricks.WorkspaceClient, tokenID str tflog.Debug(ctx, fmt.Sprintf("unable to delete token with ID %s from workspace using the provided service principal, continuing", tokenID)) return nil } + if apierr.IsMissing(err) { + tflog.Debug(ctx, fmt.Sprintf("token with ID %s not found, skipping deletion", tokenID)) + return nil + } if err != nil { return fmt.Errorf("cannot remove token: %w", err) } diff --git a/mws/resource_mws_workspaces_test.go b/mws/resource_mws_workspaces_test.go index 348af5da52..017dae349f 100644 --- a/mws/resource_mws_workspaces_test.go +++ b/mws/resource_mws_workspaces_test.go @@ -100,20 +100,20 @@ func TestResourceWorkspaceCreate(t *testing.T) { }, MockWorkspaceClientsFunc: basicMockWorkspaceClients(t, setDefaultConfigHost, mockScimMe), Resource: ResourceMwsWorkspaces(), - State: map[string]any{ - "account_id": "abc", - "aws_region": "us-east-1", - "credentials_id": "bcd", - "managed_services_customer_managed_key_id": "def", - "storage_customer_managed_key_id": "def", - "deployment_name": "900150983cd24fb0", - "workspace_name": "labdata", - "network_id": "fgh", - "storage_configuration_id": "ghi", - "custom_tags": map[string]any{ - "SoldToCode": "1234", - }, - }, + HCL: ` + account_id = "abc" + aws_region = "us-east-1" + credentials_id = "bcd" + managed_services_customer_managed_key_id = "def" + storage_customer_managed_key_id = "def" + deployment_name = "900150983cd24fb0" + workspace_name = "labdata" + network_id = "fgh" + storage_configuration_id = "ghi" + custom_tags = { + SoldToCode = "1234" + } + `, Create: true, }.Apply(t) assert.NoError(t, err) @@ -409,18 +409,18 @@ func TestResourceWorkspaceCreateWithIsNoPublicIPEnabledFalse(t *testing.T) { }, MockWorkspaceClientsFunc: basicMockWorkspaceClients(t, setDefaultConfigHost, mockScimMe), Resource: ResourceMwsWorkspaces(), - State: map[string]any{ - "account_id": "abc", - "aws_region": "us-east-1", - "credentials_id": "bcd", - "managed_services_customer_managed_key_id": "def", - "storage_customer_managed_key_id": "def", - "deployment_name": "900150983cd24fb0", - "workspace_name": "labdata", - "is_no_public_ip_enabled": false, - "network_id": "fgh", - "storage_configuration_id": "ghi", - }, + HCL: ` + account_id = "abc" + aws_region = "us-east-1" + credentials_id = "bcd" + managed_services_customer_managed_key_id = "def" + storage_customer_managed_key_id = "def" + deployment_name = "900150983cd24fb0" + workspace_name = "labdata" + is_no_public_ip_enabled = false + network_id = "fgh" + storage_configuration_id = "ghi" + `, Create: true, }.Apply(t) assert.NoError(t, err) @@ -470,16 +470,16 @@ func TestResourceWorkspaceCreateLegacyConfig(t *testing.T) { }, MockWorkspaceClientsFunc: basicMockWorkspaceClients(t, setDefaultConfigHost, mockScimMe), Resource: ResourceMwsWorkspaces(), - State: map[string]any{ - "account_id": "abc", - "aws_region": "us-east-1", - "credentials_id": "bcd", - "customer_managed_key_id": "def", - "deployment_name": "900150983cd24fb0", - "workspace_name": "labdata", - "network_id": "fgh", - "storage_configuration_id": "ghi", - }, + HCL: ` + account_id = "abc" + aws_region = "us-east-1" + credentials_id = "bcd" + customer_managed_key_id = "def" + deployment_name = "900150983cd24fb0" + workspace_name = "labdata" + network_id = "fgh" + storage_configuration_id = "ghi" + `, Create: true, }.Apply(t) assert.NoError(t, err) @@ -716,20 +716,19 @@ func TestResourceWorkspaceUpdate_NotAllowed(t *testing.T) { "storage_configuration_id": "ghi", "workspace_id": "1234", }, - State: map[string]any{ - "account_id": "THIS_IS_CHANGING", - - "aws_region": "us-east-1", - "credentials_id": "bcd", - "managed_services_customer_managed_key_id": "def", - "storage_customer_managed_key_id": "def", - "deployment_name": "900150983cd24fb0", - "workspace_name": "labdata", - "is_no_public_ip_enabled": true, - "network_id": "fgh", - "storage_configuration_id": "ghi", - "workspace_id": 1234, - }, + HCL: ` + account_id = "THIS_IS_CHANGING" + aws_region = "us-east-1" + credentials_id = "bcd" + managed_services_customer_managed_key_id = "def" + storage_customer_managed_key_id = "def" + deployment_name = "900150983cd24fb0" + workspace_name = "labdata" + is_no_public_ip_enabled = true + network_id = "fgh" + storage_configuration_id = "ghi" + workspace_id = 1234 + `, Update: true, ID: "abc/1234", }.ExpectError(t, "changes require new: account_id") @@ -941,9 +940,7 @@ func TestCreateFailsAndCleansUp(t *testing.T) { }).Return(mockFailedWorkspace, nil) // Expect the Get call to retrieve the network with errors - a.GetMockNetworksAPI().EXPECT().Get(mock.Anything, provisioning.GetNetworkRequest{ - NetworkId: "fgh", - }).Return(mockNetwork, nil) + a.GetMockNetworksAPI().EXPECT().Get(mock.Anything, provisioning.GetNetworkRequest{NetworkId: "fgh"}).Return(mockNetwork, nil) // Expect the Delete call to clean up the failed workspace a.GetMockWorkspacesAPI().EXPECT().Delete(mock.Anything, provisioning.DeleteWorkspaceRequest{ @@ -1149,22 +1146,22 @@ func TestResourceWorkspaceUpdatePrivateAccessSettings(t *testing.T) { "storage_configuration_id": "ghi", "workspace_id": "1234", }, - State: map[string]any{ - "account_id": "abc", - "aws_region": "us-east-1", - "credentials_id": "bcd", - "managed_services_customer_managed_key_id": "def", - "storage_customer_managed_key_id": "def", - "deployment_name": "900150983cd24fb0", - "workspace_name": "labdata", - "is_no_public_ip_enabled": true, - "network_id": "fgh", - "storage_configuration_id": "ghi", - "private_access_settings_id": "pas", - "workspace_id": 1234, - }, Update: true, ID: "abc/1234", + HCL: ` + account_id = "abc" + aws_region = "us-east-1" + credentials_id = "bcd" + managed_services_customer_managed_key_id = "def" + storage_customer_managed_key_id = "def" + deployment_name = "900150983cd24fb0" + workspace_name = "labdata" + private_access_settings_id = "pas" + is_no_public_ip_enabled = true + network_id = "fgh" + storage_configuration_id = "ghi" + workspace_id = 1234 + `, }.Apply(t) assert.NoError(t, err) assert.Equal(t, "abc/1234", d.Id(), "Id should be the same as in reading") @@ -1283,3 +1280,274 @@ func TestSensitiveDataInLogs(t *testing.T) { assert.NotContains(t, fmt.Sprintf("%#v", tk), "sensitive") assert.NotContains(t, fmt.Sprintf("%+v", tk), "sensitive") } + +func TestResourceWorkspaceAddToken(t *testing.T) { + // Define a mock workspace that can be reused + mockWorkspace := &provisioning.Workspace{ + WorkspaceId: 1234, + WorkspaceStatus: provisioning.WorkspaceStatusRunning, + WorkspaceName: "labdata", + DeploymentName: "900150983cd24fb0", + AwsRegion: "us-east-1", + CredentialsId: "bcd", + StorageConfigurationId: "ghi", + NetworkId: "fgh", + ManagedServicesCustomerManagedKeyId: "def", + StorageCustomerManagedKeyId: "def", + AccountId: "abc", + } + + addToken := func(m *mocks.MockWorkspaceClient) { + // Create a mock workspace client with token API expectations + mockTokensAPI := m.GetMockTokensAPI() + + // Expect the list token call to return no existing tokens + mockTokensAPI.EXPECT(). + List(mock.Anything). + Return(&listing.SliceIterator[settings.PublicTokenInfo]{}) + + // Expect the create token call for the new token + mockTokensAPI.EXPECT(). + Create(mock.Anything, settings.CreateTokenRequest{ + LifetimeSeconds: 3600, + Comment: "New token comment", + }). + Return(&settings.CreateTokenResponse{ + TokenValue: "new-token-value", + TokenInfo: &settings.PublicTokenInfo{ + TokenId: "new-token-id", + }, + }, nil) + } + + d, err := qa.ResourceFixture{ + MockAccountClientFunc: func(a *mocks.MockAccountClient) { + // Expect the Get call to retrieve the workspace + a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ + WorkspaceId: 1234, + }).Return(mockWorkspace, nil) + a.GetMockWorkspacesAPI().EXPECT().WaitGetWorkspaceRunning(mock.Anything, int64(1234), 20*time.Minute, mock.Anything).Return(mockWorkspace, nil) + }, + MockWorkspaceClientsFunc: basicMockWorkspaceClients(t, setDefaultConfigHost, addToken), + Resource: ResourceMwsWorkspaces(), + InstanceState: map[string]string{ + "account_id": "abc", + "aws_region": "us-east-1", + "credentials_id": "bcd", + "managed_services_customer_managed_key_id": "def", + "storage_customer_managed_key_id": "def", + "deployment_name": "900150983cd24fb0", + "workspace_name": "labdata", + "is_no_public_ip_enabled": "true", + "network_id": "fgh", + "storage_configuration_id": "ghi", + "workspace_id": "1234", + }, + HCL: ` + account_id = "abc" + aws_region = "us-east-1" + credentials_id = "bcd" + managed_services_customer_managed_key_id = "def" + storage_customer_managed_key_id = "def" + deployment_name = "900150983cd24fb0" + workspace_name = "labdata" + is_no_public_ip_enabled = true + network_id = "fgh" + storage_configuration_id = "ghi" + workspace_id = 1234 + token { + comment = "New token comment" + lifetime_seconds = 3600 + } + `, + Update: true, + ID: "abc/1234", + }.Apply(t) + assert.NoError(t, err) + assert.Equal(t, "abc/1234", d.Id(), "Id should be the same as in reading") + + // Verify that the token was added to the state + token := d.Get("token").([]any)[0].(map[string]any) + assert.Equal(t, "new-token-id", token["token_id"]) + assert.Equal(t, "new-token-value", token["token_value"]) +} + +func TestResourceWorkspaceUpdateToken(t *testing.T) { + // a helper function to set up token API mocks + updateToken := func(c *mocks.MockWorkspaceClient) { + mockTokensAPI := c.GetMockTokensAPI() + + // Expect the list token call + mockTokensAPI.EXPECT(). + List(mock.Anything). + Return(&listing.SliceIterator[settings.PublicTokenInfo]{ + { + TokenId: "old-token-id", + }, + }) + + // Expect the revoke token call for the old token + mockTokensAPI.EXPECT(). + Delete(mock.Anything, settings.RevokeTokenRequest{TokenId: "old-token-id"}). + Return(nil) + + // Expect the create token call for the new token + mockTokensAPI.EXPECT(). + Create(mock.Anything, settings.CreateTokenRequest{ + LifetimeSeconds: 3600, + Comment: "New token comment", + }). + Return(&settings.CreateTokenResponse{ + TokenValue: "new-token-value", + TokenInfo: &settings.PublicTokenInfo{ + TokenId: "new-token-id", + }, + }, nil) + } + // Define a mock workspace that can be reused + mockWorkspace := &provisioning.Workspace{ + WorkspaceId: 1234, + WorkspaceStatus: provisioning.WorkspaceStatusRunning, + WorkspaceName: "labdata", + DeploymentName: "900150983cd24fb0", + AwsRegion: "us-east-1", + CredentialsId: "bcd", + StorageConfigurationId: "ghi", + NetworkId: "fgh", + ManagedServicesCustomerManagedKeyId: "def", + StorageCustomerManagedKeyId: "def", + AccountId: "abc", + } + + d, err := qa.ResourceFixture{ + MockAccountClientFunc: func(a *mocks.MockAccountClient) { + // Expect the Get call to retrieve the workspace + a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ + WorkspaceId: 1234, + }).Return(mockWorkspace, nil) + a.GetMockWorkspacesAPI().EXPECT().WaitGetWorkspaceRunning(mock.Anything, int64(1234), 20*time.Minute, mock.Anything).Return(mockWorkspace, nil) + }, + MockWorkspaceClientsFunc: basicMockWorkspaceClients(t, setDefaultConfigHost, updateToken), + Resource: ResourceMwsWorkspaces(), + InstanceState: map[string]string{ + "account_id": "abc", + "aws_region": "us-east-1", + "credentials_id": "bcd", + "managed_services_customer_managed_key_id": "def", + "storage_customer_managed_key_id": "def", + "deployment_name": "900150983cd24fb0", + "workspace_name": "labdata", + "is_no_public_ip_enabled": "true", + "network_id": "fgh", + "storage_configuration_id": "ghi", + "workspace_id": "1234", + "token.#": "1", + "token.0.comment": "Old token comment", + "token.0.lifetime_seconds": "3600", + "token.0.token_id": "old-token-id", + "token.0.token_value": "old-token-value", + }, + HCL: ` + account_id = "abc" + aws_region = "us-east-1" + credentials_id = "bcd" + managed_services_customer_managed_key_id = "def" + storage_customer_managed_key_id = "def" + deployment_name = "900150983cd24fb0" + workspace_name = "labdata" + is_no_public_ip_enabled = true + network_id = "fgh" + storage_configuration_id = "ghi" + workspace_id = 1234 + token { + comment = "New token comment" + lifetime_seconds = 3600 + } + `, + Update: true, + ID: "abc/1234", + }.Apply(t) + assert.NoError(t, err) + assert.Equal(t, "abc/1234", d.Id(), "Id should be the same as in reading") + + // Verify that the token was updated in the state + token := d.Get("token").([]any)[0].(map[string]any) + assert.Equal(t, "new-token-id", token["token_id"]) + assert.Equal(t, "new-token-value", token["token_value"]) +} + +func TestResourceWorkspaceDeleteToken(t *testing.T) { + // Define a mock workspace that can be reused + mockWorkspace := &provisioning.Workspace{ + WorkspaceId: 1234, + WorkspaceStatus: provisioning.WorkspaceStatusRunning, + WorkspaceName: "labdata", + DeploymentName: "900150983cd24fb0", + AwsRegion: "us-east-1", + CredentialsId: "bcd", + StorageConfigurationId: "ghi", + NetworkId: "fgh", + ManagedServicesCustomerManagedKeyId: "def", + StorageCustomerManagedKeyId: "def", + AccountId: "abc", + } + + revokeToken := func(m *mocks.MockWorkspaceClient) { + // Create a mock workspace client with token API expectations + mockTokensAPI := m.GetMockTokensAPI() + + // Expect the revoke token call for the old token + mockTokensAPI.EXPECT(). + Delete(mock.Anything, settings.RevokeTokenRequest{TokenId: "old-token-id"}). + Return(nil) + } + + d, err := qa.ResourceFixture{ + MockAccountClientFunc: func(a *mocks.MockAccountClient) { + // Expect the Get call to retrieve the workspace + a.GetMockWorkspacesAPI().EXPECT().Get(mock.Anything, provisioning.GetWorkspaceRequest{ + WorkspaceId: 1234, + }).Return(mockWorkspace, nil) + a.GetMockWorkspacesAPI().EXPECT().WaitGetWorkspaceRunning(mock.Anything, int64(1234), 20*time.Minute, mock.Anything).Return(mockWorkspace, nil) + }, + MockWorkspaceClientsFunc: basicMockWorkspaceClients(t, setDefaultConfigHost, revokeToken), + Resource: ResourceMwsWorkspaces(), + InstanceState: map[string]string{ + "account_id": "abc", + "aws_region": "us-east-1", + "credentials_id": "bcd", + "managed_services_customer_managed_key_id": "def", + "storage_customer_managed_key_id": "def", + "deployment_name": "900150983cd24fb0", + "workspace_name": "labdata", + "is_no_public_ip_enabled": "true", + "network_id": "fgh", + "storage_configuration_id": "ghi", + "workspace_id": "1234", + "token.#": "1", + "token.0.comment": "Old token comment", + "token.0.lifetime_seconds": "3600", + "token.0.token_id": "old-token-id", + "token.0.token_value": "old-token-value", + }, + HCL: ` + account_id = "abc" + aws_region = "us-east-1" + credentials_id = "bcd" + managed_services_customer_managed_key_id = "def" + storage_customer_managed_key_id = "def" + deployment_name = "900150983cd24fb0" + workspace_name = "labdata" + is_no_public_ip_enabled = true + network_id = "fgh" + storage_configuration_id = "ghi" + workspace_id = 1234 + `, + Update: true, + ID: "abc/1234", + }.Apply(t) + assert.NoError(t, err) + + // Verify that the token was removed from the state + assert.Len(t, d.Get("token"), 0) +}