Use MockWorkspaceClient from SDK instead of WithImpl mocking (#1134)

## Changes
Use MockWorkspaceClient from SDK instead of WithImpl mocking

bundle/bundle.go

@@ -140,6 +140,13 @@ func (b *Bundle) WorkspaceClient() *databricks.WorkspaceClient {
 	return b.client
 }
 
+// SetWorkpaceClient sets the workspace client for this bundle.
+// This is used to inject a mock client for testing.
+func (b *Bundle) SetWorkpaceClient(w *databricks.WorkspaceClient) {
+	b.clientOnce.Do(func() {})
+	b.client = w
+}
+
 // CacheDir returns directory to use for temporary files for this bundle.
 // Scoped to the bundle's target.
 func (b *Bundle) CacheDir(ctx context.Context, paths ...string) (string, error) {

bundle/config/mutator/resolve_resource_references_test.go

@@ -2,122 +2,19 @@ package mutator
 
 import (
 	"context"
+	"fmt"
 	"testing"
 
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/bundle/config"
 	"github.com/databricks/cli/bundle/config/variable"
-	"github.com/databricks/databricks-sdk-go/service/compute"
+	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
+
+	"github.com/databricks/databricks-sdk-go/experimental/mocks"
+	"github.com/databricks/databricks-sdk-go/service/compute"
 )
 
-type MockClusterService struct{}
-
-// ChangeOwner implements compute.ClustersService.
-func (MockClusterService) ChangeOwner(ctx context.Context, request compute.ChangeClusterOwner) error {
-	panic("unimplemented")
-}
-
-// Create implements compute.ClustersService.
-func (MockClusterService) Create(ctx context.Context, request compute.CreateCluster) (*compute.CreateClusterResponse, error) {
-	panic("unimplemented")
-}
-
-// Delete implements compute.ClustersService.
-func (MockClusterService) Delete(ctx context.Context, request compute.DeleteCluster) error {
-	panic("unimplemented")
-}
-
-// Edit implements compute.ClustersService.
-func (MockClusterService) Edit(ctx context.Context, request compute.EditCluster) error {
-	panic("unimplemented")
-}
-
-// Events implements compute.ClustersService.
-func (MockClusterService) Events(ctx context.Context, request compute.GetEvents) (*compute.GetEventsResponse, error) {
-	panic("unimplemented")
-}
-
-// Get implements compute.ClustersService.
-func (MockClusterService) Get(ctx context.Context, request compute.GetClusterRequest) (*compute.ClusterDetails, error) {
-	panic("unimplemented")
-}
-
-// GetPermissionLevels implements compute.ClustersService.
-func (MockClusterService) GetPermissionLevels(ctx context.Context, request compute.GetClusterPermissionLevelsRequest) (*compute.GetClusterPermissionLevelsResponse, error) {
-	panic("unimplemented")
-}
-
-// GetPermissions implements compute.ClustersService.
-func (MockClusterService) GetPermissions(ctx context.Context, request compute.GetClusterPermissionsRequest) (*compute.ClusterPermissions, error) {
-	panic("unimplemented")
-}
-
-// List implements compute.ClustersService.
-func (MockClusterService) List(ctx context.Context, request compute.ListClustersRequest) (*compute.ListClustersResponse, error) {
-	return &compute.ListClustersResponse{
-		Clusters: []compute.ClusterDetails{
-			{ClusterId: "1234-5678-abcd", ClusterName: "Some Custom Cluster"},
-			{ClusterId: "9876-5432-xywz", ClusterName: "Some Other Name"},
-		},
-	}, nil
-}
-
-// ListNodeTypes implements compute.ClustersService.
-func (MockClusterService) ListNodeTypes(ctx context.Context) (*compute.ListNodeTypesResponse, error) {
-	panic("unimplemented")
-}
-
-// ListZones implements compute.ClustersService.
-func (MockClusterService) ListZones(ctx context.Context) (*compute.ListAvailableZonesResponse, error) {
-	panic("unimplemented")
-}
-
-// PermanentDelete implements compute.ClustersService.
-func (MockClusterService) PermanentDelete(ctx context.Context, request compute.PermanentDeleteCluster) error {
-	panic("unimplemented")
-}
-
-// Pin implements compute.ClustersService.
-func (MockClusterService) Pin(ctx context.Context, request compute.PinCluster) error {
-	panic("unimplemented")
-}
-
-// Resize implements compute.ClustersService.
-func (MockClusterService) Resize(ctx context.Context, request compute.ResizeCluster) error {
-	panic("unimplemented")
-}
-
-// Restart implements compute.ClustersService.
-func (MockClusterService) Restart(ctx context.Context, request compute.RestartCluster) error {
-	panic("unimplemented")
-}
-
-// SetPermissions implements compute.ClustersService.
-func (MockClusterService) SetPermissions(ctx context.Context, request compute.ClusterPermissionsRequest) (*compute.ClusterPermissions, error) {
-	panic("unimplemented")
-}
-
-// SparkVersions implements compute.ClustersService.
-func (MockClusterService) SparkVersions(ctx context.Context) (*compute.GetSparkVersionsResponse, error) {
-	panic("unimplemented")
-}
-
-// Start implements compute.ClustersService.
-func (MockClusterService) Start(ctx context.Context, request compute.StartCluster) error {
-	panic("unimplemented")
-}
-
-// Unpin implements compute.ClustersService.
-func (MockClusterService) Unpin(ctx context.Context, request compute.UnpinCluster) error {
-	panic("unimplemented")
-}
-
-// UpdatePermissions implements compute.ClustersService.
-func (MockClusterService) UpdatePermissions(ctx context.Context, request compute.ClusterPermissionsRequest) (*compute.ClusterPermissions, error) {
-	panic("unimplemented")
-}
-
 func TestResolveClusterReference(t *testing.T) {
 	clusterRef1 := "Some Custom Cluster"
 	clusterRef2 := "Some Other Name"
@@ -142,7 +39,15 @@ func TestResolveClusterReference(t *testing.T) {
 		},
 	}
 
-	b.WorkspaceClient().Clusters.WithImpl(MockClusterService{})
+	m := mocks.NewMockWorkspaceClient(t)
+	b.SetWorkpaceClient(m.WorkspaceClient)
+	clusterApi := m.GetMockClustersAPI()
+	clusterApi.EXPECT().GetByClusterName(mock.Anything, clusterRef1).Return(&compute.ClusterDetails{
+		ClusterId: "1234-5678-abcd",
+	}, nil)
+	clusterApi.EXPECT().GetByClusterName(mock.Anything, clusterRef2).Return(&compute.ClusterDetails{
+		ClusterId: "9876-5432-xywz",
+	}, nil)
 
 	err := bundle.Apply(context.Background(), b, ResolveResourceReferences())
 	require.NoError(t, err)
@@ -168,7 +73,10 @@ func TestResolveNonExistentClusterReference(t *testing.T) {
 		},
 	}
 
-	b.WorkspaceClient().Clusters.WithImpl(MockClusterService{})
+	m := mocks.NewMockWorkspaceClient(t)
+	b.SetWorkpaceClient(m.WorkspaceClient)
+	clusterApi := m.GetMockClustersAPI()
+	clusterApi.EXPECT().GetByClusterName(mock.Anything, clusterRef).Return(nil, fmt.Errorf("ClusterDetails named '%s' does not exist", clusterRef))
 
 	err := bundle.Apply(context.Background(), b, ResolveResourceReferences())
 	require.ErrorContains(t, err, "failed to resolve cluster: Random, err: ClusterDetails named 'Random' does not exist")
@@ -188,7 +96,9 @@ func TestNoLookupIfVariableIsSet(t *testing.T) {
 		},
 	}
 
-	b.WorkspaceClient().Clusters.WithImpl(MockClusterService{})
+	m := mocks.NewMockWorkspaceClient(t)
+	b.SetWorkpaceClient(m.WorkspaceClient)
+
 	b.Config.Variables["my-cluster-id"].Set("random value")
 
 	err := bundle.Apply(context.Background(), b, ResolveResourceReferences())

bundle/permissions/workspace_root_test.go

@@ -7,85 +7,16 @@ import (
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/bundle/config"
 	"github.com/databricks/cli/bundle/config/resources"
+	"github.com/databricks/databricks-sdk-go/experimental/mocks"
 	"github.com/databricks/databricks-sdk-go/service/jobs"
 	"github.com/databricks/databricks-sdk-go/service/ml"
 	"github.com/databricks/databricks-sdk-go/service/pipelines"
 	"github.com/databricks/databricks-sdk-go/service/serving"
 	"github.com/databricks/databricks-sdk-go/service/workspace"
+	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 )
 
-type MockWorkspaceClient struct {
-	t *testing.T
-}
-
-// Delete implements workspace.WorkspaceService.
-func (MockWorkspaceClient) Delete(ctx context.Context, request workspace.Delete) error {
-	panic("unimplemented")
-}
-
-// Export implements workspace.WorkspaceService.
-func (MockWorkspaceClient) Export(ctx context.Context, request workspace.ExportRequest) (*workspace.ExportResponse, error) {
-	panic("unimplemented")
-}
-
-// GetPermissionLevels implements workspace.WorkspaceService.
-func (MockWorkspaceClient) GetPermissionLevels(ctx context.Context, request workspace.GetWorkspaceObjectPermissionLevelsRequest) (*workspace.GetWorkspaceObjectPermissionLevelsResponse, error) {
-	panic("unimplemented")
-}
-
-// GetPermissions implements workspace.WorkspaceService.
-func (MockWorkspaceClient) GetPermissions(ctx context.Context, request workspace.GetWorkspaceObjectPermissionsRequest) (*workspace.WorkspaceObjectPermissions, error) {
-	panic("unimplemented")
-}
-
-// GetStatus implements workspace.WorkspaceService.
-func (MockWorkspaceClient) GetStatus(ctx context.Context, request workspace.GetStatusRequest) (*workspace.ObjectInfo, error) {
-	return &workspace.ObjectInfo{
-		ObjectId: 1234, ObjectType: "directories", Path: "/Users/foo@bar.com",
-	}, nil
-}
-
-// Import implements workspace.WorkspaceService.
-func (MockWorkspaceClient) Import(ctx context.Context, request workspace.Import) error {
-	panic("unimplemented")
-}
-
-// List implements workspace.WorkspaceService.
-func (MockWorkspaceClient) List(ctx context.Context, request workspace.ListWorkspaceRequest) (*workspace.ListResponse, error) {
-	panic("unimplemented")
-}
-
-// Mkdirs implements workspace.WorkspaceService.
-func (MockWorkspaceClient) Mkdirs(ctx context.Context, request workspace.Mkdirs) error {
-	panic("unimplemented")
-}
-
-// SetPermissions implements workspace.WorkspaceService.
-func (MockWorkspaceClient) SetPermissions(ctx context.Context, request workspace.WorkspaceObjectPermissionsRequest) (*workspace.WorkspaceObjectPermissions, error) {
-	panic("unimplemented")
-}
-
-// UpdatePermissions implements workspace.WorkspaceService.
-func (m MockWorkspaceClient) UpdatePermissions(ctx context.Context, request workspace.WorkspaceObjectPermissionsRequest) (*workspace.WorkspaceObjectPermissions, error) {
-	require.Equal(m.t, "1234", request.WorkspaceObjectId)
-	require.Equal(m.t, "directories", request.WorkspaceObjectType)
-	require.Contains(m.t, request.AccessControlList, workspace.WorkspaceObjectAccessControlRequest{
-		UserName:        "TestUser",
-		PermissionLevel: "CAN_MANAGE",
-	})
-	require.Contains(m.t, request.AccessControlList, workspace.WorkspaceObjectAccessControlRequest{
-		GroupName:       "TestGroup",
-		PermissionLevel: "CAN_READ",
-	})
-	require.Contains(m.t, request.AccessControlList, workspace.WorkspaceObjectAccessControlRequest{
-		ServicePrincipalName: "TestServicePrincipal",
-		PermissionLevel:      "CAN_RUN",
-	})
-
-	return nil, nil
-}
-
 func TestApplyWorkspaceRootPermissions(t *testing.T) {
 	b := &bundle.Bundle{
 		Config: config.Root{
@@ -122,7 +53,21 @@ func TestApplyWorkspaceRootPermissions(t *testing.T) {
 		},
 	}
 
-	b.WorkspaceClient().Workspace.WithImpl(MockWorkspaceClient{t})
+	m := mocks.NewMockWorkspaceClient(t)
+	b.SetWorkpaceClient(m.WorkspaceClient)
+	workspaceApi := m.GetMockWorkspaceAPI()
+	workspaceApi.EXPECT().GetStatusByPath(mock.Anything, "/Users/foo@bar.com").Return(&workspace.ObjectInfo{
+		ObjectId: 1234,
+	}, nil)
+	workspaceApi.EXPECT().UpdatePermissions(mock.Anything, workspace.WorkspaceObjectPermissionsRequest{
+		AccessControlList: []workspace.WorkspaceObjectAccessControlRequest{
+			{UserName: "TestUser", PermissionLevel: "CAN_MANAGE"},
+			{GroupName: "TestGroup", PermissionLevel: "CAN_READ"},
+			{ServicePrincipalName: "TestServicePrincipal", PermissionLevel: "CAN_RUN"},
+		},
+		WorkspaceObjectId:   "1234",
+		WorkspaceObjectType: "directories",
+	}).Return(nil, nil)
 
 	err := bundle.Apply(context.Background(), b, ApplyWorkspaceRootPermissions())
 	require.NoError(t, err)

bundle/python/warning_test.go

@@ -7,122 +7,13 @@ import (
 	"github.com/databricks/cli/bundle"
 	"github.com/databricks/cli/bundle/config"
 	"github.com/databricks/cli/bundle/config/resources"
+	"github.com/databricks/databricks-sdk-go/experimental/mocks"
 	"github.com/databricks/databricks-sdk-go/service/compute"
 	"github.com/databricks/databricks-sdk-go/service/jobs"
+	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 )
 
-type MockClusterService struct{}
-
-// ChangeOwner implements compute.ClustersService.
-func (MockClusterService) ChangeOwner(ctx context.Context, request compute.ChangeClusterOwner) error {
-	panic("unimplemented")
-}
-
-// Create implements compute.ClustersService.
-func (MockClusterService) Create(ctx context.Context, request compute.CreateCluster) (*compute.CreateClusterResponse, error) {
-	panic("unimplemented")
-}
-
-// Delete implements compute.ClustersService.
-func (MockClusterService) Delete(ctx context.Context, request compute.DeleteCluster) error {
-	panic("unimplemented")
-}
-
-// Edit implements compute.ClustersService.
-func (MockClusterService) Edit(ctx context.Context, request compute.EditCluster) error {
-	panic("unimplemented")
-}
-
-// Events implements compute.ClustersService.
-func (MockClusterService) Events(ctx context.Context, request compute.GetEvents) (*compute.GetEventsResponse, error) {
-	panic("unimplemented")
-}
-
-// Get implements compute.ClustersService.
-func (MockClusterService) Get(ctx context.Context, request compute.GetClusterRequest) (*compute.ClusterDetails, error) {
-	clusterDetails := map[string]*compute.ClusterDetails{
-		"test-key-1": {
-			SparkVersion: "12.2.x-scala2.12",
-		},
-		"test-key-2": {
-			SparkVersion: "13.2.x-scala2.12",
-		},
-	}
-
-	return clusterDetails[request.ClusterId], nil
-}
-
-// GetPermissionLevels implements compute.ClustersService.
-func (MockClusterService) GetPermissionLevels(ctx context.Context, request compute.GetClusterPermissionLevelsRequest) (*compute.GetClusterPermissionLevelsResponse, error) {
-	panic("unimplemented")
-}
-
-// GetPermissions implements compute.ClustersService.
-func (MockClusterService) GetPermissions(ctx context.Context, request compute.GetClusterPermissionsRequest) (*compute.ClusterPermissions, error) {
-	panic("unimplemented")
-}
-
-// List implements compute.ClustersService.
-func (MockClusterService) List(ctx context.Context, request compute.ListClustersRequest) (*compute.ListClustersResponse, error) {
-	panic("unimplemented")
-}
-
-// ListNodeTypes implements compute.ClustersService.
-func (MockClusterService) ListNodeTypes(ctx context.Context) (*compute.ListNodeTypesResponse, error) {
-	panic("unimplemented")
-}
-
-// ListZones implements compute.ClustersService.
-func (MockClusterService) ListZones(ctx context.Context) (*compute.ListAvailableZonesResponse, error) {
-	panic("unimplemented")
-}
-
-// PermanentDelete implements compute.ClustersService.
-func (MockClusterService) PermanentDelete(ctx context.Context, request compute.PermanentDeleteCluster) error {
-	panic("unimplemented")
-}
-
-// Pin implements compute.ClustersService.
-func (MockClusterService) Pin(ctx context.Context, request compute.PinCluster) error {
-	panic("unimplemented")
-}
-
-// Resize implements compute.ClustersService.
-func (MockClusterService) Resize(ctx context.Context, request compute.ResizeCluster) error {
-	panic("unimplemented")
-}
-
-// Restart implements compute.ClustersService.
-func (MockClusterService) Restart(ctx context.Context, request compute.RestartCluster) error {
-	panic("unimplemented")
-}
-
-// SetPermissions implements compute.ClustersService.
-func (MockClusterService) SetPermissions(ctx context.Context, request compute.ClusterPermissionsRequest) (*compute.ClusterPermissions, error) {
-	panic("unimplemented")
-}
-
-// SparkVersions implements compute.ClustersService.
-func (MockClusterService) SparkVersions(ctx context.Context) (*compute.GetSparkVersionsResponse, error) {
-	panic("unimplemented")
-}
-
-// Start implements compute.ClustersService.
-func (MockClusterService) Start(ctx context.Context, request compute.StartCluster) error {
-	panic("unimplemented")
-}
-
-// Unpin implements compute.ClustersService.
-func (MockClusterService) Unpin(ctx context.Context, request compute.UnpinCluster) error {
-	panic("unimplemented")
-}
-
-// UpdatePermissions implements compute.ClustersService.
-func (MockClusterService) UpdatePermissions(ctx context.Context, request compute.ClusterPermissionsRequest) (*compute.ClusterPermissions, error) {
-	panic("unimplemented")
-}
-
 func TestIncompatibleWheelTasksWithNewCluster(t *testing.T) {
 	b := &bundle.Bundle{
 		Config: config.Root{
@@ -246,7 +137,12 @@ func TestIncompatibleWheelTasksWithExistingClusterId(t *testing.T) {
 		},
 	}
 
-	b.WorkspaceClient().Clusters.WithImpl(MockClusterService{})
+	m := mocks.NewMockWorkspaceClient(t)
+	b.SetWorkpaceClient(m.WorkspaceClient)
+	clustersApi := m.GetMockClustersAPI()
+	clustersApi.EXPECT().GetByClusterId(mock.Anything, "test-key-1").Return(&compute.ClusterDetails{
+		SparkVersion: "12.2.x-scala2.12",
+	}, nil)
 
 	require.True(t, hasIncompatibleWheelTasks(context.Background(), b))
 }
@@ -335,7 +231,12 @@ func TestNoIncompatibleWheelTasks(t *testing.T) {
 		},
 	}
 
-	b.WorkspaceClient().Clusters.WithImpl(MockClusterService{})
+	m := mocks.NewMockWorkspaceClient(t)
+	b.SetWorkpaceClient(m.WorkspaceClient)
+	clustersApi := m.GetMockClustersAPI()
+	clustersApi.EXPECT().GetByClusterId(mock.Anything, "test-key-2").Return(&compute.ClusterDetails{
+		SparkVersion: "13.2.x-scala2.12",
+	}, nil)
 
 	require.False(t, hasIncompatibleWheelTasks(context.Background(), b))
 }

go.mod

@@ -55,6 +55,7 @@ require (
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
 	github.com/mattn/go-colorable v0.1.13 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
+	github.com/stretchr/objx v0.5.0 // indirect
 	github.com/zclconf/go-cty v1.14.1 // indirect
 	go.opencensus.io v0.24.0 // indirect
 	go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.46.1 // indirect

go.sum

@@ -148,6 +148,7 @@ github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
 github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
+github.com/stretchr/objx v0.5.0 h1:1zr/of2m5FGMsad5YfcqgdqdWrIhu+EBEJRhR1U7z/c=
 github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
 github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=