From 4ac1c1655b1736c00143c528cbbdba2baa190286 Mon Sep 17 00:00:00 2001
From: shreyas-goenka <88374338+shreyas-goenka@users.noreply.github.com>
Date: Wed, 21 Feb 2024 18:36:03 +0530
Subject: [PATCH] Fix CLI nightlies on our UC workspaces (#1225)

This PR fixes some test helper functions so that they work properly on
our test environments for AWS and Azure UC workspaces.
---
 internal/helpers.go        | 3 ++-
 internal/testutil/cloud.go | 3 +++
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/internal/helpers.go b/internal/helpers.go
index ca5aa25e..49dc9f4c 100644
--- a/internal/helpers.go
+++ b/internal/helpers.go
@@ -523,7 +523,8 @@ 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" {
+	} else if env == "aws" || env == "ucws" {
+		// aws-prod-ucws has CLOUD_ENV set to "ucws"
 		return "i3.xlarge"
 	}
 	return "Standard_DS4_v2"
diff --git a/internal/testutil/cloud.go b/internal/testutil/cloud.go
index 50bbf67f..e547069f 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)
 	}