Add tests

Author: renaudhartert-db

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java

@@ -297,8 +297,7 @@ private Token exchange(String code, String state) {
 
   static class CallbackResponseHandler implements HttpHandler {
     private final Logger LOG = LoggerFactory.getLogger(getClass().getName());
-    // Protects params
-    private final Object lock = new Object();
+    protected final Object lock = new Object(); // protected for testing
     private volatile Map<String, String> params;
     private final Optional<Duration> timeout;
 
@@ -343,10 +342,7 @@ public void handleInner(HttpExchange exchange) throws IOException {
               });
 
       sendSuccess(exchange);
-      synchronized (lock) {
-        params = theseParams;
-        lock.notify();
-      }
+      setParams(theseParams);
     }
 
     private void sendError(
@@ -411,5 +407,12 @@ public Map<String, String> getParams() {
         return params;
       }
     }
+
+    void setParams(Map<String, String> params) {
+      synchronized (lock) {
+        this.params = params;
+        lock.notify();
+      }
+    }
   }
 }

databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ConsentTest.java

@@ -2,8 +2,13 @@
 
 import static org.junit.jupiter.api.Assertions.*;
 
+import com.databricks.sdk.core.DatabricksException;
 import java.time.Duration;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.Optional;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 import org.junit.jupiter.api.Test;
 
 public class ConsentTest {
@@ -22,7 +27,6 @@ public void testConsentWithBrowserAuthTimeout() {
             .withBrowserTimeout(Optional.of(Duration.ofSeconds(30)))
             .build();
 
-    // Verify that the timeout is properly set
     assertEquals(Optional.of(Duration.ofSeconds(30)), consent.getBrowserTimeout());
   }
 
@@ -39,7 +43,90 @@ public void testConsentWithoutBrowserAuthTimeout() {
             .withVerifier("test-verifier")
             .build();
 
-    // Verify that the timeout is empty when not set
     assertEquals(Optional.empty(), consent.getBrowserTimeout());
   }
+
+  @Test
+  public void testTimeoutLogicWithShortTimeout() throws InterruptedException {
+    // Test that timeout is enforced correctly.
+    Consent.CallbackResponseHandler handler =
+        new Consent.CallbackResponseHandler(Optional.of(Duration.ofMillis(100))); // 100ms timeout
+
+    long startTime = System.currentTimeMillis();
+
+    try {
+      handler.getParams();
+      fail("Expected timeout exception");
+    } catch (DatabricksException e) {
+      long elapsedTime = System.currentTimeMillis() - startTime;
+      assertTrue(
+          elapsedTime >= 100, "Timeout should have taken at least 100ms, but took " + elapsedTime);
+      assertTrue(e.getMessage().contains("timed out after 0 seconds"));
+    }
+  }
+
+  @Test
+  public void testTimeoutLogicWithNoTimeout() throws InterruptedException {
+    // Test that no timeout means indefinite wait.
+    Consent.CallbackResponseHandler handler = new Consent.CallbackResponseHandler(Optional.empty());
+
+    CountDownLatch latch = new CountDownLatch(1);
+
+    Thread setterThread =
+        new Thread(
+            () -> {
+              try {
+                Thread.sleep(50);
+                synchronized (handler.lock) {
+                  Map<String, String> params = new HashMap<>();
+                  params.put("code", "test-code");
+                  params.put("state", "test-state");
+                  handler.setParams(params);
+                }
+                latch.countDown();
+              } catch (InterruptedException e) {
+                Thread.currentThread().interrupt();
+              }
+            });
+
+    setterThread.start();
+
+    Map<String, String> result = handler.getParams();
+    assertNotNull(result);
+    assertEquals("test-code", result.get("code"));
+    assertEquals("test-state", result.get("state"));
+    assertTrue(latch.await(1, TimeUnit.SECONDS));
+  }
+
+  @Test
+  public void testTimeoutLogicWithParamsSetBeforeTimeout() throws InterruptedException {
+    // Test that if params are set before timeout, no exception is thrown.
+    Consent.CallbackResponseHandler handler =
+        new Consent.CallbackResponseHandler(Optional.of(Duration.ofSeconds(1)));
+
+    CountDownLatch latch = new CountDownLatch(1);
+
+    Thread setterThread =
+        new Thread(
+            () -> {
+              try {
+                Thread.sleep(50);
+                synchronized (handler.lock) {
+                  Map<String, String> params = new HashMap<>();
+                  params.put("code", "test-code");
+                  handler.setParams(params);
+                }
+                latch.countDown();
+              } catch (InterruptedException e) {
+                Thread.currentThread().interrupt();
+              }
+            });
+
+    setterThread.start();
+
+    Map<String, String> result = handler.getParams();
+    assertNotNull(result);
+    assertEquals("test-code", result.get("code"));
+    assertTrue(latch.await(1, TimeUnit.SECONDS));
+  }
 }