From ad88558a8f499dcdd01b43d583917be4cddd503c Mon Sep 17 00:00:00 2001
From: Omer Lachish <289488+rauchy@users.noreply.github.com>
Date: Fri, 18 Oct 2024 13:29:06 +0200
Subject: [PATCH] [Fix] Fix Ubuntu unit tests on Github Actions (#366)

## Changes
<!-- Summary of your changes that are easy to understand -->
About a week ago, unit tests running on `ubuntu-latest` started failing,
likely due to a new version of 22.04 being rolled out. As a result, the
actual Azure CLI was getting executed instead of the stubbed version
during tests. This PR fixes that by making sure the unit tests directly
target the stubbed CLI.

## Tests
<!-- How is this tested? -->

Co-authored-by: Omer Lachish <rauchy@users.noreply.github.com>
---
 .../sdk/core/AzureCliCredentialsProvider.java         |  5 ++++-
 .../com/databricks/sdk/DatabricksAuthManualTest.java  |  9 ++++++---
 .../com/databricks/sdk/core/utils/TestOSUtils.java    | 11 ++++++++++-
 3 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/AzureCliCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/AzureCliCredentialsProvider.java
index c176e17b6..b08dfb8b4 100644
--- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/AzureCliCredentialsProvider.java
+++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/AzureCliCredentialsProvider.java
@@ -19,10 +19,13 @@ public String authType() {
   }
 
   public CliTokenSource tokenSourceFor(DatabricksConfig config, String resource) {
+    String azPath =
+        Optional.ofNullable(config.getEnv()).map(env -> env.get("AZ_PATH")).orElse("az");
+
     List<String> cmd =
         new ArrayList<>(
             Arrays.asList(
-                "az", "account", "get-access-token", "--resource", resource, "--output", "json"));
+                azPath, "account", "get-access-token", "--resource", resource, "--output", "json"));
     Optional<String> subscription = getSubscription(config);
     if (subscription.isPresent()) {
       // This will fail if the user has access to the workspace, but not to the subscription
diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/DatabricksAuthManualTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/DatabricksAuthManualTest.java
index 95396b75f..c467d5f4b 100644
--- a/databricks-sdk-java/src/test/java/com/databricks/sdk/DatabricksAuthManualTest.java
+++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/DatabricksAuthManualTest.java
@@ -13,7 +13,8 @@ void azureCliWorkspaceHeaderPresent() {
     StaticEnv env =
         new StaticEnv()
             .with("HOME", TestOSUtils.resource("/testdata/azure"))
-            .with("PATH", "testdata:/bin");
+            .with("PATH", "/bin:testdata")
+            .with("AZ_PATH", TestOSUtils.resource("/testdata/az"));
     String azureWorkspaceResourceId =
         "/subscriptions/123/resourceGroups/abc/providers/Microsoft.Databricks/workspaces/abc123";
     DatabricksConfig config =
@@ -32,7 +33,8 @@ void azureCliUserWithManagementAccess() {
     StaticEnv env =
         new StaticEnv()
             .with("HOME", TestOSUtils.resource("/testdata/azure"))
-            .with("PATH", "testdata:/bin");
+            .with("PATH", "/bin:testdata")
+            .with("AZ_PATH", TestOSUtils.resource("/testdata/az"));
     String azureWorkspaceResourceId =
         "/subscriptions/123/resourceGroups/abc/providers/Microsoft.Databricks/workspaces/abc123";
     DatabricksConfig config =
@@ -50,7 +52,8 @@ void azureCliUserNoManagementAccess() {
     StaticEnv env =
         new StaticEnv()
             .with("HOME", TestOSUtils.resource("/testdata/azure"))
-            .with("PATH", "testdata:/bin")
+            .with("PATH", "/bin:testdata")
+            .with("AZ_PATH", TestOSUtils.resource("/testdata/az"))
             .with("FAIL_IF", "https://management.core.windows.net/");
     String azureWorkspaceResourceId =
         "/subscriptions/123/resourceGroups/abc/providers/Microsoft.Databricks/workspaces/abc123";
diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/utils/TestOSUtils.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/utils/TestOSUtils.java
index 591049558..e4c19c021 100644
--- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/utils/TestOSUtils.java
+++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/utils/TestOSUtils.java
@@ -28,6 +28,15 @@ public static String resource(String file) {
     if (resource == null) {
       fail("Asset not found: " + file);
     }
-    return resource.getFile();
+
+    String filePath = resource.getFile();
+    File f = new File(filePath);
+
+    // Make the file executable
+    if (!f.setExecutable(true)) {
+      fail("Failed to set the file as executable: " + file);
+    }
+
+    return filePath;
   }
 }