From 04c6371a0a323a95e1195e853e1b0b6588378bb9 Mon Sep 17 00:00:00 2001
From: Satvik Ramaprasad <satvik.ramaprasad@databricks.com>
Date: Tue, 3 Sep 2024 10:29:07 +0530
Subject: [PATCH] Handle login.html incorrect validation for PL

---
 .../sdk/core/error/PrivateLinkInfo.java       |  3 +-
 .../sdk/core/error/PrivateLinkInfoTest.java   | 39 +++++++++++++++++++
 2 files changed, 41 insertions(+), 1 deletion(-)
 create mode 100644 databricks-sdk-java/src/test/java/com/databricks/sdk/core/error/PrivateLinkInfoTest.java

diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/error/PrivateLinkInfo.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/error/PrivateLinkInfo.java
index 9ef75aa7f..6b11453ae 100644
--- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/error/PrivateLinkInfo.java
+++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/error/PrivateLinkInfo.java
@@ -53,7 +53,8 @@ public String errorMessage() {
 
   public static boolean isPrivateLinkRedirect(Response resp) {
     return resp.getUrl().getPath().equals("/login.html")
-        && resp.getUrl().getQuery().contains("error=private-link-validation-error");
+        && (resp.getUrl().getQuery() != null
+            && resp.getUrl().getQuery().contains("error=private-link-validation-error"));
   }
 
   static PrivateLinkValidationError createPrivateLinkValidationError(Response resp) {
diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/error/PrivateLinkInfoTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/error/PrivateLinkInfoTest.java
new file mode 100644
index 000000000..0c6848a3a
--- /dev/null
+++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/error/PrivateLinkInfoTest.java
@@ -0,0 +1,39 @@
+package com.databricks.sdk.core.error;
+
+import com.databricks.sdk.core.error.platform.*;
+import com.databricks.sdk.core.http.Request;
+import com.databricks.sdk.core.http.Response;
+import org.junit.jupiter.api.Test;
+
+public class PrivateLinkInfoTest {
+  @Test
+  void testIsPrivateLinkRedirectWithLoginHtmlAndQueryParamSet() {
+    Response response =
+        new Response(
+            new Request("GET", "https://example.com/login.html")
+                .withQueryParam("error", "private-link-validation-error"),
+            307,
+            "Temporary Redirect",
+            null);
+    assert PrivateLinkInfo.isPrivateLinkRedirect(response) == true;
+  }
+
+  @Test
+  void testIsPrivateLinkRedirectWithLoginHtmlAndQueryParamNotSet() {
+    Response response =
+        new Response(
+            new Request("GET", "https://example.com/login.html"), 307, "Temporary Redirect", null);
+    assert PrivateLinkInfo.isPrivateLinkRedirect(response) == false;
+  }
+
+  @Test
+  void testIsPrivateLinkRedirectNotLoginPage() {
+    Response response =
+        new Response(
+            new Request("GET", "https://example.com/not-login.html"),
+            307,
+            "Temporary Redirect",
+            null);
+    assert PrivateLinkInfo.isPrivateLinkRedirect(response) == false;
+  }
+}