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)
 	}