package deploy import ( "bytes" "context" "encoding/json" "io" "os" "path/filepath" "testing" "github.com/databricks/cli/bundle" "github.com/databricks/cli/bundle/config" "github.com/databricks/cli/bundle/deploy/files" mockfiler "github.com/databricks/cli/internal/mocks/libs/filer" "github.com/databricks/cli/internal/testutil" "github.com/databricks/cli/libs/filer" "github.com/databricks/cli/libs/sync" "github.com/databricks/databricks-sdk-go/service/iam" "github.com/stretchr/testify/mock" "github.com/stretchr/testify/require" ) type snapshortStateExpectations struct { localToRemoteNames map[string]string remoteToLocalNames map[string]string } type statePullExpectations struct { seq int filesInDevelopmentState []File snapshotState *snapshortStateExpectations } type statePullOpts struct { files []File seq int localFiles []string localNotebooks []string expects statePullExpectations withExistingSnapshot bool localState *DeploymentState } func testStatePull(t *testing.T, opts statePullOpts) { s := &statePull{func(b *bundle.Bundle) (filer.Filer, error) { f := mockfiler.NewMockFiler(t) deploymentStateData, err := json.Marshal(DeploymentState{ Version: DeploymentStateVersion, Seq: int64(opts.seq), Files: opts.files, }) require.NoError(t, err) f.EXPECT().Read(mock.Anything, DeploymentStateFileName).Return(io.NopCloser(bytes.NewReader(deploymentStateData)), nil) return f, nil }} b := &bundle.Bundle{ RootPath: t.TempDir(), Config: config.Root{ Bundle: config.Bundle{ Target: "default", }, Workspace: config.Workspace{ StatePath: "/state", CurrentUser: &config.User{ User: &iam.User{ UserName: "test-user", }, }, }, }, } ctx := context.Background() for _, file := range opts.localFiles { testutil.Touch(t, filepath.Join(b.RootPath, "bar"), file) } for _, file := range opts.localNotebooks { testutil.TouchNotebook(t, filepath.Join(b.RootPath, "bar"), file) } if opts.withExistingSnapshot { opts, err := files.GetSyncOptions(ctx, b) require.NoError(t, err) snapshotPath, err := sync.SnapshotPath(opts) require.NoError(t, err) err = os.WriteFile(snapshotPath, []byte("snapshot"), 0644) require.NoError(t, err) } if opts.localState != nil { statePath, err := getPathToStateFile(ctx, b) require.NoError(t, err) data, err := json.Marshal(opts.localState) require.NoError(t, err) err = os.WriteFile(statePath, data, 0644) require.NoError(t, err) } diags := bundle.Apply(ctx, b, s) require.NoError(t, diags.Error()) // Check that deployment state was written statePath, err := getPathToStateFile(ctx, b) require.NoError(t, err) data, err := os.ReadFile(statePath) require.NoError(t, err) var state DeploymentState err = json.Unmarshal(data, &state) require.NoError(t, err) require.Equal(t, int64(opts.expects.seq), state.Seq) require.Len(t, state.Files, len(opts.expects.filesInDevelopmentState)) for i, file := range opts.expects.filesInDevelopmentState { require.Equal(t, file.LocalPath, state.Files[i].LocalPath) } if opts.expects.snapshotState != nil { syncOpts, err := files.GetSyncOptions(ctx, b) require.NoError(t, err) snapshotPath, err := sync.SnapshotPath(syncOpts) require.NoError(t, err) _, err = os.Stat(snapshotPath) require.NoError(t, err) data, err = os.ReadFile(snapshotPath) require.NoError(t, err) var snapshot sync.Snapshot err = json.Unmarshal(data, &snapshot) require.NoError(t, err) snapshotState := snapshot.SnapshotState require.Len(t, snapshotState.LocalToRemoteNames, len(opts.expects.snapshotState.localToRemoteNames)) for local, remote := range opts.expects.snapshotState.localToRemoteNames { require.Equal(t, remote, snapshotState.LocalToRemoteNames[local]) } require.Len(t, snapshotState.RemoteToLocalNames, len(opts.expects.snapshotState.remoteToLocalNames)) for remote, local := range opts.expects.snapshotState.remoteToLocalNames { require.Equal(t, local, snapshotState.RemoteToLocalNames[remote]) } } } var stateFiles []File = []File{ { LocalPath: "bar/t1.py", IsNotebook: false, }, { LocalPath: "bar/t2.py", IsNotebook: false, }, { LocalPath: "bar/notebook.py", IsNotebook: true, }, } func TestStatePull(t *testing.T) { testStatePull(t, statePullOpts{ seq: 1, files: stateFiles, localFiles: []string{"t1.py", "t2.py"}, localNotebooks: []string{"notebook.py"}, expects: statePullExpectations{ seq: 1, filesInDevelopmentState: []File{ { LocalPath: "bar/t1.py", }, { LocalPath: "bar/t2.py", }, { LocalPath: "bar/notebook.py", }, }, snapshotState: &snapshortStateExpectations{ localToRemoteNames: map[string]string{ "bar/t1.py": "bar/t1.py", "bar/t2.py": "bar/t2.py", "bar/notebook.py": "bar/notebook", }, remoteToLocalNames: map[string]string{ "bar/t1.py": "bar/t1.py", "bar/t2.py": "bar/t2.py", "bar/notebook": "bar/notebook.py", }, }, }, }) } func TestStatePullSnapshotExists(t *testing.T) { testStatePull(t, statePullOpts{ withExistingSnapshot: true, seq: 1, files: stateFiles, localFiles: []string{"t1.py", "t2.py"}, expects: statePullExpectations{ seq: 1, filesInDevelopmentState: []File{ { LocalPath: "bar/t1.py", }, { LocalPath: "bar/t2.py", }, { LocalPath: "bar/notebook.py", }, }, snapshotState: &snapshortStateExpectations{ localToRemoteNames: map[string]string{ "bar/t1.py": "bar/t1.py", "bar/t2.py": "bar/t2.py", "bar/notebook.py": "bar/notebook", }, remoteToLocalNames: map[string]string{ "bar/t1.py": "bar/t1.py", "bar/t2.py": "bar/t2.py", "bar/notebook": "bar/notebook.py", }, }, }, }) } func TestStatePullNoState(t *testing.T) { s := &statePull{func(b *bundle.Bundle) (filer.Filer, error) { f := mockfiler.NewMockFiler(t) f.EXPECT().Read(mock.Anything, DeploymentStateFileName).Return(nil, os.ErrNotExist) return f, nil }} b := &bundle.Bundle{ RootPath: t.TempDir(), Config: config.Root{ Bundle: config.Bundle{ Target: "default", }, Workspace: config.Workspace{ StatePath: "/state", }, }, } ctx := context.Background() diags := bundle.Apply(ctx, b, s) require.NoError(t, diags.Error()) // Check that deployment state was not written statePath, err := getPathToStateFile(ctx, b) require.NoError(t, err) _, err = os.Stat(statePath) require.True(t, os.IsNotExist(err)) } func TestStatePullOlderState(t *testing.T) { testStatePull(t, statePullOpts{ seq: 1, files: stateFiles, localFiles: []string{"t1.py", "t2.py"}, localNotebooks: []string{"notebook.py"}, localState: &DeploymentState{ Version: DeploymentStateVersion, Seq: 2, Files: []File{ { LocalPath: "bar/t1.py", }, }, }, expects: statePullExpectations{ seq: 2, filesInDevelopmentState: []File{ { LocalPath: "bar/t1.py", }, }, }, }) } func TestStatePullNewerState(t *testing.T) { testStatePull(t, statePullOpts{ seq: 1, files: stateFiles, localFiles: []string{"t1.py", "t2.py"}, localNotebooks: []string{"notebook.py"}, localState: &DeploymentState{ Version: DeploymentStateVersion, Seq: 0, Files: []File{ { LocalPath: "bar/t1.py", }, }, }, expects: statePullExpectations{ seq: 1, filesInDevelopmentState: []File{ { LocalPath: "bar/t1.py", }, { LocalPath: "bar/t2.py", }, { LocalPath: "bar/notebook.py", }, }, snapshotState: &snapshortStateExpectations{ localToRemoteNames: map[string]string{ "bar/t1.py": "bar/t1.py", "bar/t2.py": "bar/t2.py", "bar/notebook.py": "bar/notebook", }, remoteToLocalNames: map[string]string{ "bar/t1.py": "bar/t1.py", "bar/t2.py": "bar/t2.py", "bar/notebook": "bar/notebook.py", }, }, }, }) } func TestStatePullAndFileIsRemovedLocally(t *testing.T) { testStatePull(t, statePullOpts{ seq: 1, files: stateFiles, localFiles: []string{"t2.py"}, // t1.py is removed locally localNotebooks: []string{"notebook.py"}, expects: statePullExpectations{ seq: 1, filesInDevelopmentState: []File{ { LocalPath: "bar/t1.py", }, { LocalPath: "bar/t2.py", }, { LocalPath: "bar/notebook.py", }, }, snapshotState: &snapshortStateExpectations{ localToRemoteNames: map[string]string{ "bar/t1.py": "bar/t1.py", "bar/t2.py": "bar/t2.py", "bar/notebook.py": "bar/notebook", }, remoteToLocalNames: map[string]string{ "bar/t1.py": "bar/t1.py", "bar/t2.py": "bar/t2.py", "bar/notebook": "bar/notebook.py", }, }, }, }) } func TestStatePullAndNotebookIsRemovedLocally(t *testing.T) { testStatePull(t, statePullOpts{ seq: 1, files: stateFiles, localFiles: []string{"t1.py", "t2.py"}, localNotebooks: []string{}, // notebook.py is removed locally expects: statePullExpectations{ seq: 1, filesInDevelopmentState: []File{ { LocalPath: "bar/t1.py", }, { LocalPath: "bar/t2.py", }, { LocalPath: "bar/notebook.py", }, }, snapshotState: &snapshortStateExpectations{ localToRemoteNames: map[string]string{ "bar/t1.py": "bar/t1.py", "bar/t2.py": "bar/t2.py", "bar/notebook.py": "bar/notebook", }, remoteToLocalNames: map[string]string{ "bar/t1.py": "bar/t1.py", "bar/t2.py": "bar/t2.py", "bar/notebook": "bar/notebook.py", }, }, }, }) } func TestStatePullNewerDeploymentStateVersion(t *testing.T) { s := &statePull{func(b *bundle.Bundle) (filer.Filer, error) { f := mockfiler.NewMockFiler(t) deploymentStateData, err := json.Marshal(DeploymentState{ Version: DeploymentStateVersion + 1, Seq: 1, CliVersion: "1.2.3", Files: []File{ { LocalPath: "bar/t1.py", }, { LocalPath: "bar/t2.py", }, }, }) require.NoError(t, err) f.EXPECT().Read(mock.Anything, DeploymentStateFileName).Return(io.NopCloser(bytes.NewReader(deploymentStateData)), nil) return f, nil }} b := &bundle.Bundle{ RootPath: t.TempDir(), Config: config.Root{ Bundle: config.Bundle{ Target: "default", }, Workspace: config.Workspace{ StatePath: "/state", }, }, } ctx := context.Background() diags := bundle.Apply(ctx, b, s) require.True(t, diags.HasError()) require.ErrorContains(t, diags.Error(), "remote deployment state is incompatible with the current version of the CLI, please upgrade to at least 1.2.3") }