From 22623abe5570f93657105254fc4bd2c109683f79 Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Tue, 20 Feb 2024 17:00:14 +0100
Subject: [PATCH 1/3] Fix CLI nightlies to run on UC workspaces

---
 internal/helpers.go        | 7 ++++---
 internal/testutil/cloud.go | 3 +++
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/internal/helpers.go b/internal/helpers.go
index 6377ae07e7..eaaf80c1d7 100644
--- a/internal/helpers.go
+++ b/internal/helpers.go
@@ -484,8 +484,9 @@ func TemporaryRepo(t *testing.T, w *databricks.WorkspaceClient) string {
 func GetNodeTypeId(env string) string {
 	if env == "gcp" {
 		return "n1-standard-4"
-	} else if env == "aws" {
-		return "i3.xlarge"
+	} else if env == "azure" {
+		return "Standard_DS4_v2"
 	}
-	return "Standard_DS4_v2"
+	// We default to AWS because our "aws-prod-ucws" test environment has CLOUD_ENV set to "ucws"
+	return "i3.xlarge"
 }
diff --git a/internal/testutil/cloud.go b/internal/testutil/cloud.go
index 50bbf67f25..e547069f38 100644
--- a/internal/testutil/cloud.go
+++ b/internal/testutil/cloud.go
@@ -41,6 +41,9 @@ func GetCloud(t *testing.T) Cloud {
 		return Azure
 	case "gcp":
 		return GCP
+	// CLOUD_ENV is set to "ucws" in the "aws-prod-ucws" test environment
+	case "ucws":
+		return AWS
 	default:
 		t.Fatalf("Unknown cloud environment: %s", env)
 	}

From 7cd61e9080281738d2da155d87fda6267cb5630b Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Wed, 21 Feb 2024 11:43:07 +0100
Subject: [PATCH 2/3] -

---
 internal/helpers.go | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/internal/helpers.go b/internal/helpers.go
index eaaf80c1d7..e9627b0c49 100644
--- a/internal/helpers.go
+++ b/internal/helpers.go
@@ -484,9 +484,10 @@ func TemporaryRepo(t *testing.T, w *databricks.WorkspaceClient) string {
 func GetNodeTypeId(env string) string {
 	if env == "gcp" {
 		return "n1-standard-4"
-	} else if env == "azure" {
-		return "Standard_DS4_v2"
 	}
-	// We default to AWS because our "aws-prod-ucws" test environment has CLOUD_ENV set to "ucws"
-	return "i3.xlarge"
+	// aws-prod-ucws has CLOUD_ENV set to "ucws"
+	else if env == "aws" || env == "ucws" {
+		return "i3.xlarge"
+	}
+	return "Standard_DS4_v2"
 }

From b9a0b88e622b5d38ab636fb5bc97d0a1b50d337d Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Wed, 21 Feb 2024 11:49:21 +0100
Subject: [PATCH 3/3] lint

---
 internal/helpers.go | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/internal/helpers.go b/internal/helpers.go
index e9627b0c49..bacc970b8c 100644
--- a/internal/helpers.go
+++ b/internal/helpers.go
@@ -484,9 +484,8 @@ func TemporaryRepo(t *testing.T, w *databricks.WorkspaceClient) string {
 func GetNodeTypeId(env string) string {
 	if env == "gcp" {
 		return "n1-standard-4"
-	}
-	// aws-prod-ucws has CLOUD_ENV set to "ucws"
-	else if env == "aws" || env == "ucws" {
+	} else if env == "aws" || env == "ucws" {
+		// aws-prod-ucws has CLOUD_ENV set to "ucws"
 		return "i3.xlarge"
 	}
 	return "Standard_DS4_v2"