Move cloud-specific node type helper

pietern · pietern · commit 020f14646c5b · 2024-12-12T17:10:16.000+01:00
diff --git a/internal/bundle/basic_test.go b/internal/bundle/basic_test.go
@@ -5,17 +5,16 @@ import (
 	"path/filepath"
 	"testing"
 
-	"github.com/databricks/cli/internal"
 	"github.com/databricks/cli/internal/acc"
-	"github.com/databricks/cli/libs/env"
+	"github.com/databricks/cli/internal/testutil"
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/require"
 )
 
 func TestAccBasicBundleDeployWithFailOnActiveRuns(t *testing.T) {
 	ctx, _ := acc.WorkspaceTest(t)
 
-	nodeTypeId := internal.GetNodeTypeId(env.Get(ctx, "CLOUD_ENV"))
+	nodeTypeId := testutil.GetCloud(t).NodeTypeID()
 	uniqueId := uuid.New().String()
 	root, err := initTestTemplate(t, ctx, "basic", map[string]any{
 		"unique_id":     uniqueId,
diff --git a/internal/bundle/bind_resource_test.go b/internal/bundle/bind_resource_test.go
@@ -6,7 +6,6 @@ import (
 	"path/filepath"
 	"testing"
 
-	"github.com/databricks/cli/internal"
 	"github.com/databricks/cli/internal/acc"
 	"github.com/databricks/cli/internal/testcli"
 	"github.com/databricks/cli/internal/testutil"
@@ -24,7 +23,7 @@ func TestAccBindJobToExistingJob(t *testing.T) {
 	ctx, wt := acc.WorkspaceTest(t)
 	gt := &generateJobTest{T: wt, w: wt.W}
 
-	nodeTypeId := internal.GetNodeTypeId(env)
+	nodeTypeId := testutil.GetCloud(t).NodeTypeID()
 	uniqueId := uuid.New().String()
 	bundleRoot, err := initTestTemplate(t, ctx, "basic", map[string]any{
 		"unique_id":     uniqueId,
@@ -89,7 +88,7 @@ func TestAccAbortBind(t *testing.T) {
 	ctx, wt := acc.WorkspaceTest(t)
 	gt := &generateJobTest{T: wt, w: wt.W}
 
-	nodeTypeId := internal.GetNodeTypeId(env)
+	nodeTypeId := testutil.GetCloud(t).NodeTypeID()
 	uniqueId := uuid.New().String()
 	bundleRoot, err := initTestTemplate(t, ctx, "basic", map[string]any{
 		"unique_id":     uniqueId,
diff --git a/internal/bundle/clusters_test.go b/internal/bundle/clusters_test.go
@@ -4,10 +4,8 @@ import (
 	"fmt"
 	"testing"
 
-	"github.com/databricks/cli/internal"
 	"github.com/databricks/cli/internal/acc"
 	"github.com/databricks/cli/internal/testutil"
-	"github.com/databricks/cli/libs/env"
 	"github.com/databricks/databricks-sdk-go/service/compute"
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/require"
@@ -20,7 +18,7 @@ func TestAccDeployBundleWithCluster(t *testing.T) {
 		t.Skip("Skipping test for AWS cloud because it is not permitted to create clusters")
 	}
 
-	nodeTypeId := internal.GetNodeTypeId(env.Get(ctx, "CLOUD_ENV"))
+	nodeTypeId := testutil.GetCloud(t).NodeTypeID()
 	uniqueId := uuid.New().String()
 	root, err := initTestTemplate(t, ctx, "clusters", map[string]any{
 		"unique_id":     uniqueId,
diff --git a/internal/bundle/deploy_test.go b/internal/bundle/deploy_test.go
@@ -11,10 +11,9 @@ import (
 	"testing"
 
 	"github.com/databricks/cli/cmd/root"
-	"github.com/databricks/cli/internal"
 	"github.com/databricks/cli/internal/acc"
 	"github.com/databricks/cli/internal/testcli"
-	"github.com/databricks/cli/libs/env"
+	"github.com/databricks/cli/internal/testutil"
 	"github.com/databricks/databricks-sdk-go"
 	"github.com/databricks/databricks-sdk-go/apierr"
 	"github.com/databricks/databricks-sdk-go/service/catalog"
@@ -133,7 +132,7 @@ func TestAccBundlePipelineDeleteWithoutAutoApprove(t *testing.T) {
 	ctx, wt := acc.WorkspaceTest(t)
 	w := wt.W
 
-	nodeTypeId := internal.GetNodeTypeId(env.Get(ctx, "CLOUD_ENV"))
+	nodeTypeId := testutil.GetCloud(t).NodeTypeID()
 	uniqueId := uuid.New().String()
 	bundleRoot, err := initTestTemplate(t, ctx, "deploy_then_remove_resources", map[string]any{
 		"unique_id":     uniqueId,
@@ -219,7 +218,7 @@ properties such as the 'catalog' or 'storage' are changed:
 func TestAccDeployBasicBundleLogs(t *testing.T) {
 	ctx, wt := acc.WorkspaceTest(t)
 
-	nodeTypeId := internal.GetNodeTypeId(env.Get(ctx, "CLOUD_ENV"))
+	nodeTypeId := testutil.GetCloud(t).NodeTypeID()
 	uniqueId := uuid.New().String()
 	root, err := initTestTemplate(t, ctx, "basic", map[string]any{
 		"unique_id":     uniqueId,
diff --git a/internal/bundle/deploy_then_remove_resources_test.go b/internal/bundle/deploy_then_remove_resources_test.go
@@ -5,9 +5,8 @@ import (
 	"path/filepath"
 	"testing"
 
-	"github.com/databricks/cli/internal"
 	"github.com/databricks/cli/internal/acc"
-	"github.com/databricks/cli/libs/env"
+	"github.com/databricks/cli/internal/testutil"
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -17,7 +16,7 @@ func TestAccBundleDeployThenRemoveResources(t *testing.T) {
 	ctx, wt := acc.WorkspaceTest(t)
 	w := wt.W
 
-	nodeTypeId := internal.GetNodeTypeId(env.Get(ctx, "CLOUD_ENV"))
+	nodeTypeId := testutil.GetCloud(t).NodeTypeID()
 	uniqueId := uuid.New().String()
 	bundleRoot, err := initTestTemplate(t, ctx, "deploy_then_remove_resources", map[string]any{
 		"unique_id":     uniqueId,
diff --git a/internal/bundle/deploy_to_shared_test.go b/internal/bundle/deploy_to_shared_test.go
@@ -4,17 +4,16 @@ import (
 	"fmt"
 	"testing"
 
-	"github.com/databricks/cli/internal"
 	"github.com/databricks/cli/internal/acc"
-	"github.com/databricks/cli/libs/env"
+	"github.com/databricks/cli/internal/testutil"
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/require"
 )
 
 func TestAccDeployBasicToSharedWorkspacePath(t *testing.T) {
 	ctx, wt := acc.WorkspaceTest(t)
 
-	nodeTypeId := internal.GetNodeTypeId(env.Get(ctx, "CLOUD_ENV"))
+	nodeTypeId := testutil.GetCloud(t).NodeTypeID()
 	uniqueId := uuid.New().String()
 
 	currentUser, err := wt.W.CurrentUser.Me(ctx)
diff --git a/internal/bundle/deployment_state_test.go b/internal/bundle/deployment_state_test.go
@@ -7,7 +7,6 @@ import (
 	"testing"
 
 	"github.com/databricks/cli/bundle/deploy"
-	"github.com/databricks/cli/internal"
 	"github.com/databricks/cli/internal/acc"
 	"github.com/databricks/cli/internal/testutil"
 	"github.com/google/uuid"
@@ -21,7 +20,7 @@ func TestAccFilesAreSyncedCorrectlyWhenNoSnapshot(t *testing.T) {
 	ctx, wt := acc.WorkspaceTest(t)
 	w := wt.W
 
-	nodeTypeId := internal.GetNodeTypeId(env)
+	nodeTypeId := testutil.GetCloud(t).NodeTypeID()
 	uniqueId := uuid.New().String()
 	bundleRoot, err := initTestTemplate(t, ctx, "basic", map[string]any{
 		"unique_id":     uniqueId,
diff --git a/internal/bundle/destroy_test.go b/internal/bundle/destroy_test.go
@@ -6,9 +6,8 @@ import (
 	"path/filepath"
 	"testing"
 
-	"github.com/databricks/cli/internal"
 	"github.com/databricks/cli/internal/acc"
-	"github.com/databricks/cli/libs/env"
+	"github.com/databricks/cli/internal/testutil"
 	"github.com/databricks/databricks-sdk-go/apierr"
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
@@ -19,7 +18,7 @@ func TestAccBundleDestroy(t *testing.T) {
 	ctx, wt := acc.WorkspaceTest(t)
 	w := wt.W
 
-	nodeTypeId := internal.GetNodeTypeId(env.Get(ctx, "CLOUD_ENV"))
+	nodeTypeId := testutil.GetCloud(t).NodeTypeID()
 	uniqueId := uuid.New().String()
 	bundleRoot, err := initTestTemplate(t, ctx, "deploy_then_remove_resources", map[string]any{
 		"unique_id":     uniqueId,
diff --git a/internal/bundle/generate_job_test.go b/internal/bundle/generate_job_test.go
@@ -77,16 +77,6 @@ func (gt *generateJobTest) createTestJob(ctx context.Context) int64 {
 	t := gt.T
 	w := gt.w
 
-	var nodeTypeId string
-	switch testutil.GetCloud(t) {
-	case testutil.AWS:
-		nodeTypeId = "i3.xlarge"
-	case testutil.Azure:
-		nodeTypeId = "Standard_DS4_v2"
-	case testutil.GCP:
-		nodeTypeId = "n1-standard-4"
-	}
-
 	tmpdir := acc.TemporaryWorkspaceDir(t, "generate-job-")
 	f, err := filer.NewWorkspaceFilesClient(w, tmpdir)
 	require.NoError(t, err)
@@ -102,7 +92,7 @@ func (gt *generateJobTest) createTestJob(ctx context.Context) int64 {
 				NewCluster: &compute.ClusterSpec{
 					SparkVersion: "13.3.x-scala2.12",
 					NumWorkers:   1,
-					NodeTypeId:   nodeTypeId,
+					NodeTypeId:   testutil.GetCloud(t).NodeTypeID(),
 					SparkConf: map[string]string{
 						"spark.databricks.enableWsfs":                         "true",
 						"spark.databricks.hive.metastore.glueCatalog.enabled": "true",
diff --git a/internal/bundle/generate_pipeline_test.go b/internal/bundle/generate_pipeline_test.go
@@ -9,7 +9,6 @@ import (
 	"strings"
 	"testing"
 
-	"github.com/databricks/cli/internal"
 	"github.com/databricks/cli/internal/acc"
 	"github.com/databricks/cli/internal/testcli"
 	"github.com/databricks/cli/internal/testutil"
@@ -96,8 +95,7 @@ func (gt *generatePipelineTest) createTestPipeline(ctx context.Context) (string,
 	err = f.Write(ctx, "test.py", strings.NewReader("print('Hello!')"))
 	require.NoError(t, err)
 
-	env := testutil.GetEnvOrSkipTest(t, "CLOUD_ENV")
-	nodeTypeId := internal.GetNodeTypeId(env)
+	nodeTypeId := testutil.GetCloud(t).NodeTypeID()
 
 	name := testutil.RandomName("generated-pipeline-")
 	resp, err := w.Pipelines.Create(ctx, pipelines.CreatePipeline{
diff --git a/internal/bundle/job_metadata_test.go b/internal/bundle/job_metadata_test.go
@@ -10,9 +10,8 @@ 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/internal/testutil"
 	"github.com/databricks/cli/libs/filer"
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
@@ -23,7 +22,7 @@ func TestAccJobsMetadataFile(t *testing.T) {
 	ctx, wt := acc.WorkspaceTest(t)
 	w := wt.W
 
-	nodeTypeId := internal.GetNodeTypeId(env.Get(ctx, "CLOUD_ENV"))
+	nodeTypeId := testutil.GetCloud(t).NodeTypeID()
 	uniqueId := uuid.New().String()
 	bundleRoot, err := initTestTemplate(t, ctx, "job_metadata", map[string]any{
 		"unique_id":     uniqueId,
diff --git a/internal/bundle/local_state_staleness_test.go b/internal/bundle/local_state_staleness_test.go
@@ -4,9 +4,8 @@ import (
 	"context"
 	"testing"
 
-	"github.com/databricks/cli/internal"
 	"github.com/databricks/cli/internal/acc"
-	"github.com/databricks/cli/libs/env"
+	"github.com/databricks/cli/internal/testutil"
 	"github.com/databricks/databricks-sdk-go/listing"
 	"github.com/databricks/databricks-sdk-go/service/jobs"
 	"github.com/google/uuid"
@@ -25,7 +24,7 @@ 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.Get(ctx, "CLOUD_ENV"))
+	nodeTypeId := testutil.GetCloud(t).NodeTypeID()
 	uniqueId := uuid.New().String()
 	initialize := func() string {
 		root, err := initTestTemplate(t, ctx, "basic", map[string]any{
diff --git a/internal/bundle/python_wheel_test.go b/internal/bundle/python_wheel_test.go
@@ -3,7 +3,6 @@ package bundle
 import (
 	"testing"
 
-	"github.com/databricks/cli/internal"
 	"github.com/databricks/cli/internal/acc"
 	"github.com/databricks/cli/internal/testutil"
 	"github.com/databricks/cli/libs/env"
@@ -14,7 +13,7 @@ import (
 func runPythonWheelTest(t *testing.T, templateName, sparkVersion string, pythonWheelWrapper bool) {
 	ctx, _ := acc.WorkspaceTest(t)
 
-	nodeTypeId := internal.GetNodeTypeId(env.Get(ctx, "CLOUD_ENV"))
+	nodeTypeId := testutil.GetCloud(t).NodeTypeID()
 	instancePoolId := env.Get(ctx, "TEST_INSTANCE_POOL_ID")
 	bundleRoot, err := initTestTemplate(t, ctx, templateName, map[string]any{
 		"node_type_id":         nodeTypeId,
diff --git a/internal/bundle/spark_jar_test.go b/internal/bundle/spark_jar_test.go
@@ -4,7 +4,6 @@ import (
 	"context"
 	"testing"
 
-	"github.com/databricks/cli/internal"
 	"github.com/databricks/cli/internal/acc"
 	"github.com/databricks/cli/internal/testutil"
 	"github.com/databricks/cli/libs/env"
@@ -13,8 +12,7 @@ import (
 )
 
 func runSparkJarTestCommon(t *testing.T, ctx context.Context, sparkVersion, artifactPath string) {
-	cloudEnv := testutil.GetEnvOrSkipTest(t, "CLOUD_ENV")
-	nodeTypeId := internal.GetNodeTypeId(cloudEnv)
+	nodeTypeId := testutil.GetCloud(t).NodeTypeID()
 	tmpDir := t.TempDir()
 	instancePoolId := env.Get(ctx, "TEST_INSTANCE_POOL_ID")
 	bundleRoot, err := initTestTemplateWithBundleRoot(t, ctx, "spark_jar_task", map[string]any{
diff --git a/internal/helpers.go b/internal/helpers.go
@@ -15,16 +15,6 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func GetNodeTypeId(env string) string {
-	if env == "gcp" {
-		return "n1-standard-4"
-	} else if env == "aws" || env == "ucws" {
-		// aws-prod-ucws has CLOUD_ENV set to "ucws"
-		return "i3.xlarge"
-	}
-	return "Standard_DS4_v2"
-}
-
 func setupLocalFiler(t testutil.TestingT) (filer.Filer, string) {
 	t.Log(testutil.GetEnvOrSkipTest(t, "CLOUD_ENV"))
 
diff --git a/internal/python/python_tasks_test.go b/internal/python/python_tasks_test.go
@@ -14,7 +14,6 @@ import (
 	"time"
 
 	"github.com/databricks/cli/bundle/run/output"
-	"github.com/databricks/cli/internal"
 	"github.com/databricks/cli/internal/acc"
 	"github.com/databricks/cli/internal/testutil"
 	"github.com/databricks/cli/libs/filer"
@@ -129,7 +128,7 @@ func runPythonTasks(t *testing.T, tw *testFiles, opts testOpts) {
 
 	w := tw.w
 
-	nodeTypeId := internal.GetNodeTypeId(env)
+	nodeTypeId := testutil.GetCloud(t).NodeTypeID()
 	tasks := make([]jobs.SubmitTask, 0)
 	if opts.includeNotebookTasks {
 		tasks = append(tasks, GenerateNotebookTasks(tw.pyNotebookPath, sparkVersions, nodeTypeId)...)
diff --git a/internal/tags_test.go b/internal/tags_test.go
@@ -1,41 +1,27 @@
 package internal
 
 import (
-	"context"
 	"strings"
 	"testing"
 
+	"github.com/databricks/cli/internal/acc"
 	"github.com/databricks/cli/internal/testutil"
-	"github.com/databricks/databricks-sdk-go"
 	"github.com/databricks/databricks-sdk-go/service/compute"
 	"github.com/databricks/databricks-sdk-go/service/jobs"
 	"github.com/stretchr/testify/require"
 )
 
 func testTags(t *testing.T, tags map[string]string) error {
-	var nodeTypeId string
-	switch testutil.GetCloud(t) {
-	case testutil.AWS:
-		nodeTypeId = "i3.xlarge"
-	case testutil.Azure:
-		nodeTypeId = "Standard_DS4_v2"
-	case testutil.GCP:
-		nodeTypeId = "n1-standard-4"
-	}
-
-	w, err := databricks.NewWorkspaceClient()
-	require.NoError(t, err)
-
-	ctx := context.Background()
-	resp, err := w.Jobs.Create(ctx, jobs.CreateJob{
+	ctx, wt := acc.WorkspaceTest(t)
+	resp, err := wt.W.Jobs.Create(ctx, jobs.CreateJob{
 		Name: testutil.RandomName("test-tags-"),
 		Tasks: []jobs.Task{
 			{
 				TaskKey: "test",
 				NewCluster: &compute.ClusterSpec{
 					SparkVersion: "13.3.x-scala2.12",
 					NumWorkers:   1,
-					NodeTypeId:   nodeTypeId,
+					NodeTypeId:   testutil.GetCloud(t).NodeTypeID(),
 				},
 				SparkPythonTask: &jobs.SparkPythonTask{
 					PythonFile: "/doesnt_exist.py",
@@ -47,7 +33,7 @@ func testTags(t *testing.T, tags map[string]string) error {
 
 	if resp != nil {
 		t.Cleanup(func() {
-			_ = w.Jobs.DeleteByJobId(ctx, resp.JobId)
+			_ = wt.W.Jobs.DeleteByJobId(ctx, resp.JobId)
 			// Cannot enable errchecking there, tests fail with:
 			//   Error: Received unexpected error:
 			//   Job 0 does not exist.
diff --git a/internal/testutil/cloud.go b/internal/testutil/cloud.go
