sync logic with go

Author: hectorcast-db

databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java

@@ -532,11 +532,7 @@ public DatabricksConfig setAzureTenantId(String azureTenantId) {
   }
 
   public AzureEnvironment getAzureEnvironment() {
-    String env = "PUBLIC";
-    if (azureEnvironment != null) {
-      env = azureEnvironment;
-    }
-    return AzureEnvironment.getEnvironment(env);
+    return getDatabricksEnvironment().getAzureEnvironment();
   }
 
   public DatabricksConfig setAzureEnvironment(String azureEnvironment) {
@@ -871,10 +867,11 @@ public DatabricksEnvironment getDatabricksEnvironment() {
       return this.databricksEnvironment;
     }
 
-    if (this.host == null && this.azureWorkspaceResourceId != null) {
+    if ((this.host == null || this.azureEnvironment != null)
+        && this.azureWorkspaceResourceId != null) {
       String azureEnv = "PUBLIC";
-      if (this.azureEnvironment != null) {
-        azureEnv = this.azureEnvironment;
+      if (this.azureEnvironment != null && !this.azureEnvironment.isEmpty()) {
+        azureEnv = this.azureEnvironment.toUpperCase();
       }
       for (DatabricksEnvironment env : DatabricksEnvironment.ALL_ENVIRONMENTS) {
         if (env.getCloud() != Cloud.AZURE) {

databricks-sdk-java/src/test/java/com/databricks/sdk/DatabricksAuthManualTest.java

@@ -50,6 +50,7 @@ void azureCliUserWithManagementAccess() {
         createConfigWithMockClient()
             .setAuthType("azure-cli")
             .setHost("https://x")
+            .setAzureEnvironment("PUBLIC")
             .setAzureWorkspaceResourceId(azureWorkspaceResourceId);
     resolveConfig(config, env);
     Map<String, String> headers = config.authenticate();