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] -

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