diff --git a/internal/bundle/deploy_then_remove_resources_test.go b/internal/bundle/deploy_then_remove_resources_test.go index 73860593c0..72baf798cd 100644 --- a/internal/bundle/deploy_then_remove_resources_test.go +++ b/internal/bundle/deploy_then_remove_resources_test.go @@ -1,38 +1,33 @@ package bundle import ( - "context" "os" "path/filepath" "testing" - "github.com/databricks/cli/internal" - "github.com/databricks/databricks-sdk-go" + "github.com/databricks/cli/internal/acc" "github.com/google/uuid" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) func TestAccBundleDeployThenRemoveResources(t *testing.T) { - env := internal.GetEnvOrSkipTest(t, "CLOUD_ENV") - t.Log(env) + ctx, wt := acc.WorkspaceTest(t) + w := wt.W uniqueId := uuid.New().String() - bundleRoot, err := initTestTemplate(t, "deploy_then_remove_resources", map[string]any{ + bundleRoot, err := initTestTemplate(t, ctx, "deploy_then_remove_resources", map[string]any{ "unique_id": uniqueId, }) require.NoError(t, err) // deploy pipeline - err = deployBundle(t, bundleRoot) - require.NoError(t, err) - - w, err := databricks.NewWorkspaceClient() + err = deployBundle(t, ctx, bundleRoot) require.NoError(t, err) // assert pipeline is created pipelineName := "test-bundle-pipeline-" + uniqueId - pipeline, err := w.Pipelines.GetByName(context.Background(), pipelineName) + pipeline, err := w.Pipelines.GetByName(ctx, pipelineName) require.NoError(t, err) assert.Equal(t, pipeline.Name, pipelineName) @@ -41,15 +36,15 @@ func TestAccBundleDeployThenRemoveResources(t *testing.T) { require.NoError(t, err) // deploy again - err = deployBundle(t, bundleRoot) + err = deployBundle(t, ctx, bundleRoot) require.NoError(t, err) // assert pipeline is deleted - _, err = w.Pipelines.GetByName(context.Background(), pipelineName) + _, err = w.Pipelines.GetByName(ctx, pipelineName) assert.ErrorContains(t, err, "does not exist") t.Cleanup(func() { - err = destroyBundle(t, bundleRoot) + err = destroyBundle(t, ctx, bundleRoot) require.NoError(t, err) }) } diff --git a/internal/bundle/empty_bundle_test.go b/internal/bundle/empty_bundle_test.go index 9b39368f45..36883ae001 100644 --- a/internal/bundle/empty_bundle_test.go +++ b/internal/bundle/empty_bundle_test.go @@ -6,14 +6,13 @@ import ( "path/filepath" "testing" - "github.com/databricks/cli/internal" + "github.com/databricks/cli/internal/acc" "github.com/google/uuid" "github.com/stretchr/testify/require" ) func TestAccEmptyBundleDeploy(t *testing.T) { - env := internal.GetEnvOrSkipTest(t, "CLOUD_ENV") - t.Log(env) + ctx, _ := acc.WorkspaceTest(t) // create empty bundle tmpDir := t.TempDir() @@ -27,11 +26,11 @@ func TestAccEmptyBundleDeploy(t *testing.T) { f.Close() // deploy empty bundle - err = deployBundle(t, tmpDir) + err = deployBundle(t, ctx, tmpDir) require.NoError(t, err) t.Cleanup(func() { - err = destroyBundle(t, tmpDir) + err = destroyBundle(t, ctx, tmpDir) require.NoError(t, err) }) } diff --git a/internal/bundle/generate_job_test.go b/internal/bundle/generate_job_test.go index e9445abc55..e6f1578092 100644 --- a/internal/bundle/generate_job_test.go +++ b/internal/bundle/generate_job_test.go @@ -10,6 +10,7 @@ import ( "testing" "github.com/databricks/cli/internal" + "github.com/databricks/cli/internal/acc" "github.com/databricks/cli/internal/testutil" "github.com/databricks/cli/libs/filer" "github.com/databricks/databricks-sdk-go" @@ -20,23 +21,22 @@ import ( ) func TestAccGenerateFromExistingJobAndDeploy(t *testing.T) { - env := internal.GetEnvOrSkipTest(t, "CLOUD_ENV") - t.Log(env) + ctx, wt := acc.WorkspaceTest(t) + gt := &generateJobTest{T: t, w: wt.W} uniqueId := uuid.New().String() - bundleRoot, err := initTestTemplate(t, "with_includes", map[string]any{ + bundleRoot, err := initTestTemplate(t, ctx, "with_includes", map[string]any{ "unique_id": uniqueId, }) require.NoError(t, err) - jobId := createTestJob(t) + jobId := gt.createTestJob(ctx) t.Cleanup(func() { - destroyJob(t, jobId) - require.NoError(t, err) + gt.destroyJob(ctx, jobId) }) t.Setenv("BUNDLE_ROOT", bundleRoot) - c := internal.NewCobraTestRunner(t, "bundle", "generate", "job", + c := internal.NewCobraTestRunnerWithContext(t, ctx, "bundle", "generate", "job", "--existing-job-id", fmt.Sprint(jobId), "--config-dir", filepath.Join(bundleRoot, "resources"), "--source-dir", filepath.Join(bundleRoot, "src")) @@ -61,15 +61,22 @@ func TestAccGenerateFromExistingJobAndDeploy(t *testing.T) { require.Contains(t, generatedYaml, "spark_version: 13.3.x-scala2.12") require.Contains(t, generatedYaml, "num_workers: 1") - err = deployBundle(t, bundleRoot) + err = deployBundle(t, ctx, bundleRoot) require.NoError(t, err) - err = destroyBundle(t, bundleRoot) + err = destroyBundle(t, ctx, bundleRoot) require.NoError(t, err) +} +type generateJobTest struct { + T *testing.T + w *databricks.WorkspaceClient } -func createTestJob(t *testing.T) int64 { +func (gt *generateJobTest) createTestJob(ctx context.Context) int64 { + t := gt.T + w := gt.w + var nodeTypeId string switch testutil.GetCloud(t) { case testutil.AWS: @@ -80,10 +87,6 @@ func createTestJob(t *testing.T) int64 { nodeTypeId = "n1-standard-4" } - w, err := databricks.NewWorkspaceClient() - require.NoError(t, err) - - ctx := context.Background() tmpdir := internal.TemporaryWorkspaceDir(t, w) f, err := filer.NewWorkspaceFilesClient(w, tmpdir) require.NoError(t, err) @@ -112,13 +115,9 @@ func createTestJob(t *testing.T) int64 { return resp.JobId } -func destroyJob(t *testing.T, jobId int64) { - w, err := databricks.NewWorkspaceClient() - require.NoError(t, err) - - ctx := context.Background() - err = w.Jobs.Delete(ctx, jobs.DeleteJob{ +func (gt *generateJobTest) destroyJob(ctx context.Context, jobId int64) { + err := gt.w.Jobs.Delete(ctx, jobs.DeleteJob{ JobId: jobId, }) - require.NoError(t, err) + require.NoError(gt.T, err) } diff --git a/internal/bundle/generate_pipeline_test.go b/internal/bundle/generate_pipeline_test.go index 7b2323e6a9..0005e29fa1 100644 --- a/internal/bundle/generate_pipeline_test.go +++ b/internal/bundle/generate_pipeline_test.go @@ -10,6 +10,7 @@ import ( "testing" "github.com/databricks/cli/internal" + "github.com/databricks/cli/internal/acc" "github.com/databricks/cli/libs/filer" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/service/pipelines" @@ -18,23 +19,22 @@ import ( ) func TestAccGenerateFromExistingPipelineAndDeploy(t *testing.T) { - env := internal.GetEnvOrSkipTest(t, "CLOUD_ENV") - t.Log(env) + ctx, wt := acc.WorkspaceTest(t) + gt := &generatePipelineTest{T: t, w: wt.W} uniqueId := uuid.New().String() - bundleRoot, err := initTestTemplate(t, "with_includes", map[string]any{ + bundleRoot, err := initTestTemplate(t, ctx, "with_includes", map[string]any{ "unique_id": uniqueId, }) require.NoError(t, err) - pipelineId := createTestPipeline(t) + pipelineId := gt.createTestPipeline(ctx) t.Cleanup(func() { - destroyPipeline(t, pipelineId) - require.NoError(t, err) + gt.destroyPipeline(ctx, pipelineId) }) t.Setenv("BUNDLE_ROOT", bundleRoot) - c := internal.NewCobraTestRunner(t, "bundle", "generate", "pipeline", + c := internal.NewCobraTestRunnerWithContext(t, ctx, "bundle", "generate", "pipeline", "--existing-pipeline-id", fmt.Sprint(pipelineId), "--config-dir", filepath.Join(bundleRoot, "resources"), "--source-dir", filepath.Join(bundleRoot, "src")) @@ -61,18 +61,22 @@ func TestAccGenerateFromExistingPipelineAndDeploy(t *testing.T) { require.Contains(t, generatedYaml, "- file:") require.Contains(t, generatedYaml, fmt.Sprintf("path: %s", filepath.Join("..", "src", "test.py"))) - err = deployBundle(t, bundleRoot) + err = deployBundle(t, ctx, bundleRoot) require.NoError(t, err) - err = destroyBundle(t, bundleRoot) + err = destroyBundle(t, ctx, bundleRoot) require.NoError(t, err) } -func createTestPipeline(t *testing.T) string { - w, err := databricks.NewWorkspaceClient() - require.NoError(t, err) +type generatePipelineTest struct { + T *testing.T + w *databricks.WorkspaceClient +} + +func (gt *generatePipelineTest) createTestPipeline(ctx context.Context) string { + t := gt.T + w := gt.w - ctx := context.Background() tmpdir := internal.TemporaryWorkspaceDir(t, w) f, err := filer.NewWorkspaceFilesClient(w, tmpdir) require.NoError(t, err) @@ -103,13 +107,9 @@ func createTestPipeline(t *testing.T) string { return resp.PipelineId } -func destroyPipeline(t *testing.T, pipelineId string) { - w, err := databricks.NewWorkspaceClient() - require.NoError(t, err) - - ctx := context.Background() - err = w.Pipelines.Delete(ctx, pipelines.DeletePipelineRequest{ +func (gt *generatePipelineTest) destroyPipeline(ctx context.Context, pipelineId string) { + err := gt.w.Pipelines.Delete(ctx, pipelines.DeletePipelineRequest{ PipelineId: pipelineId, }) - require.NoError(t, err) + require.NoError(gt.T, err) } diff --git a/internal/bundle/helpers.go b/internal/bundle/helpers.go index 681edc2d6c..2c2b2dac94 100644 --- a/internal/bundle/helpers.go +++ b/internal/bundle/helpers.go @@ -15,7 +15,7 @@ import ( "github.com/databricks/cli/libs/template" ) -func initTestTemplate(t *testing.T, templateName string, config map[string]any) (string, error) { +func initTestTemplate(t *testing.T, ctx context.Context, templateName string, config map[string]any) (string, error) { templateRoot := filepath.Join("bundles", templateName) bundleRoot := t.TempDir() @@ -24,7 +24,7 @@ func initTestTemplate(t *testing.T, templateName string, config map[string]any) return "", err } - ctx := root.SetWorkspaceClient(context.Background(), nil) + ctx = root.SetWorkspaceClient(ctx, nil) cmd := cmdio.NewIO(flags.OutputJSON, strings.NewReader(""), os.Stdout, os.Stderr, "bundles") ctx = cmdio.InContext(ctx, cmd) @@ -46,15 +46,14 @@ func writeConfigFile(t *testing.T, config map[string]any) (string, error) { return filepath, err } -func deployBundle(t *testing.T, path string) error { +func deployBundle(t *testing.T, ctx context.Context, path string) error { t.Setenv("BUNDLE_ROOT", path) - c := internal.NewCobraTestRunner(t, "bundle", "deploy", "--force-lock") + c := internal.NewCobraTestRunnerWithContext(t, ctx, "bundle", "deploy", "--force-lock") _, _, err := c.Run() return err } -func runResource(t *testing.T, path string, key string) (string, error) { - ctx := context.Background() +func runResource(t *testing.T, ctx context.Context, path string, key string) (string, error) { ctx = cmdio.NewContext(ctx, cmdio.Default()) c := internal.NewCobraTestRunnerWithContext(t, ctx, "bundle", "run", key) @@ -62,8 +61,7 @@ func runResource(t *testing.T, path string, key string) (string, error) { return stdout.String(), err } -func runResourceWithParams(t *testing.T, path string, key string, params ...string) (string, error) { - ctx := context.Background() +func runResourceWithParams(t *testing.T, ctx context.Context, path string, key string, params ...string) (string, error) { ctx = cmdio.NewContext(ctx, cmdio.Default()) args := make([]string, 0) @@ -74,9 +72,9 @@ func runResourceWithParams(t *testing.T, path string, key string, params ...stri return stdout.String(), err } -func destroyBundle(t *testing.T, path string) error { +func destroyBundle(t *testing.T, ctx context.Context, path string) error { t.Setenv("BUNDLE_ROOT", path) - c := internal.NewCobraTestRunner(t, "bundle", "destroy", "--auto-approve") + c := internal.NewCobraTestRunnerWithContext(t, ctx, "bundle", "destroy", "--auto-approve") _, _, err := c.Run() return err } diff --git a/internal/bundle/job_metadata_test.go b/internal/bundle/job_metadata_test.go index 3e2bb7f032..0d8a431e4f 100644 --- a/internal/bundle/job_metadata_test.go +++ b/internal/bundle/job_metadata_test.go @@ -12,23 +12,21 @@ import ( "github.com/databricks/cli/bundle/config" "github.com/databricks/cli/bundle/metadata" "github.com/databricks/cli/internal" + "github.com/databricks/cli/internal/acc" + "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/filer" - "github.com/databricks/databricks-sdk-go" "github.com/google/uuid" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) func TestAccJobsMetadataFile(t *testing.T) { - env := internal.GetEnvOrSkipTest(t, "CLOUD_ENV") - t.Log(env) + ctx, wt := acc.WorkspaceTest(t) + w := wt.W - w, err := databricks.NewWorkspaceClient() - require.NoError(t, err) - - nodeTypeId := internal.GetNodeTypeId(env) + nodeTypeId := internal.GetNodeTypeId(env.Get(ctx, "CLOUD_ENV")) uniqueId := uuid.New().String() - bundleRoot, err := initTestTemplate(t, "job_metadata", map[string]any{ + bundleRoot, err := initTestTemplate(t, ctx, "job_metadata", map[string]any{ "unique_id": uniqueId, "node_type_id": nodeTypeId, "spark_version": "13.2.x-snapshot-scala2.12", @@ -36,12 +34,12 @@ func TestAccJobsMetadataFile(t *testing.T) { require.NoError(t, err) // deploy bundle - err = deployBundle(t, bundleRoot) + err = deployBundle(t, ctx, bundleRoot) require.NoError(t, err) // Cleanup the deployed bundle t.Cleanup(func() { - err = destroyBundle(t, bundleRoot) + err = destroyBundle(t, ctx, bundleRoot) require.NoError(t, err) }) diff --git a/internal/bundle/local_state_staleness_test.go b/internal/bundle/local_state_staleness_test.go index 06cfe0e0d9..872ac8a8e8 100644 --- a/internal/bundle/local_state_staleness_test.go +++ b/internal/bundle/local_state_staleness_test.go @@ -5,7 +5,8 @@ import ( "testing" "github.com/databricks/cli/internal" - "github.com/databricks/databricks-sdk-go" + "github.com/databricks/cli/internal/acc" + "github.com/databricks/cli/libs/env" "github.com/databricks/databricks-sdk-go/listing" "github.com/databricks/databricks-sdk-go/service/jobs" "github.com/google/uuid" @@ -14,11 +15,8 @@ import ( ) func TestAccLocalStateStaleness(t *testing.T) { - env := internal.GetEnvOrSkipTest(t, "CLOUD_ENV") - t.Log(env) - - w, err := databricks.NewWorkspaceClient() - require.NoError(t, err) + ctx, wt := acc.WorkspaceTest(t) + w := wt.W // The approach for this test is as follows: // 1) First deploy of bundle instance A @@ -27,10 +25,10 @@ func TestAccLocalStateStaleness(t *testing.T) { // Because of deploy (2), the locally cached state of bundle instance A should be stale. // Then for deploy (3), it must use the remote state over the stale local state. - nodeTypeId := internal.GetNodeTypeId(env) + nodeTypeId := internal.GetNodeTypeId(env.Get(ctx, "CLOUD_ENV")) uniqueId := uuid.New().String() initialize := func() string { - root, err := initTestTemplate(t, "basic", map[string]any{ + root, err := initTestTemplate(t, ctx, "basic", map[string]any{ "unique_id": uniqueId, "node_type_id": nodeTypeId, "spark_version": "13.2.x-snapshot-scala2.12", @@ -38,26 +36,28 @@ func TestAccLocalStateStaleness(t *testing.T) { require.NoError(t, err) t.Cleanup(func() { - err = destroyBundle(t, root) + err = destroyBundle(t, ctx, root) require.NoError(t, err) }) return root } + var err error + bundleA := initialize() bundleB := initialize() // 1) Deploy bundle A - err = deployBundle(t, bundleA) + err = deployBundle(t, ctx, bundleA) require.NoError(t, err) // 2) Deploy bundle B - err = deployBundle(t, bundleB) + err = deployBundle(t, ctx, bundleB) require.NoError(t, err) // 3) Deploy bundle A again - err = deployBundle(t, bundleA) + err = deployBundle(t, ctx, bundleA) require.NoError(t, err) // Assert that there is only a single job in the workspace corresponding to this bundle. diff --git a/internal/bundle/python_wheel_test.go b/internal/bundle/python_wheel_test.go index c94ed93a3d..fc14fd17bd 100644 --- a/internal/bundle/python_wheel_test.go +++ b/internal/bundle/python_wheel_test.go @@ -4,24 +4,17 @@ import ( "testing" "github.com/databricks/cli/internal" + "github.com/databricks/cli/internal/acc" + "github.com/databricks/cli/libs/env" "github.com/google/uuid" "github.com/stretchr/testify/require" ) func runPythonWheelTest(t *testing.T, sparkVersion string, pythonWheelWrapper bool) { - env := internal.GetEnvOrSkipTest(t, "CLOUD_ENV") - t.Log(env) - - var nodeTypeId string - if env == "gcp" { - nodeTypeId = "n1-standard-4" - } else if env == "aws" { - nodeTypeId = "i3.xlarge" - } else { - nodeTypeId = "Standard_DS4_v2" - } - - bundleRoot, err := initTestTemplate(t, "python_wheel_task", map[string]any{ + ctx, _ := acc.WorkspaceTest(t) + + nodeTypeId := internal.GetNodeTypeId(env.Get(ctx, "CLOUD_ENV")) + bundleRoot, err := initTestTemplate(t, ctx, "python_wheel_task", map[string]any{ "node_type_id": nodeTypeId, "unique_id": uuid.New().String(), "spark_version": sparkVersion, @@ -29,20 +22,20 @@ func runPythonWheelTest(t *testing.T, sparkVersion string, pythonWheelWrapper bo }) require.NoError(t, err) - err = deployBundle(t, bundleRoot) + err = deployBundle(t, ctx, bundleRoot) require.NoError(t, err) t.Cleanup(func() { - destroyBundle(t, bundleRoot) + destroyBundle(t, ctx, bundleRoot) }) - out, err := runResource(t, bundleRoot, "some_other_job") + out, err := runResource(t, ctx, bundleRoot, "some_other_job") require.NoError(t, err) require.Contains(t, out, "Hello from my func") require.Contains(t, out, "Got arguments:") require.Contains(t, out, "['my_test_code', 'one', 'two']") - out, err = runResourceWithParams(t, bundleRoot, "some_other_job", "--python-params=param1,param2") + out, err = runResourceWithParams(t, ctx, bundleRoot, "some_other_job", "--python-params=param1,param2") require.NoError(t, err) require.Contains(t, out, "Hello from my func") require.Contains(t, out, "Got arguments:")