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

## Changes
Use MockWorkspaceClient from SDK instead of WithImpl mocking
This commit is contained in:
Andrew Nester 2024-01-19 14:12:58 +00:00 committed by GitHub
parent 7067782cf1
commit 1b6241746e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 61 additions and 296 deletions

View File

@ -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) {

View File

@ -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())

View File

@ -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)

View File

@ -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))
}

1
go.mod
View File

@ -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

1
go.sum generated
View File

@ -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=