From 8b02f662f733bebfd8be5b5089847954adc50c5d Mon Sep 17 00:00:00 2001
From: Mason Cao <32713681+masonc08@users.noreply.github.com>
Date: Mon, 24 Aug 2026 18:57:26 +0000
Subject: [PATCH] test: cover the relayed-proxy 401 refresh-failure hint
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Adds test_failed_refresh_surfaces_reauth_hint to TestRetryOn401 — forces a 401 plus a
refresh that raises and asserts the `databricks auth login` hint prints and the response
is still relayed. Guards the fix commit on this branch.

Co-authored-by: Isaac <no-reply@databricks.com>
---
 tests/test_gateway_proxy.py | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/tests/test_gateway_proxy.py b/tests/test_gateway_proxy.py
index f103aca2..163067da 100644
--- a/tests/test_gateway_proxy.py
+++ b/tests/test_gateway_proxy.py
@@ -337,6 +337,22 @@ def refresh(self):
         self._t = "Bearer-tok2"
 
 
+class _RefreshFailsCache:
+    """A cache whose forced refresh raises, mimicking a dead Databricks OAuth
+    session that can't be re-minted non-interactively."""
+
+    def __init__(self):
+        self.refreshed = 0
+
+    @property
+    def token(self):
+        return "tok1"
+
+    def refresh(self):
+        self.refreshed += 1
+        raise RuntimeError("mint failed")
+
+
 def _handle_handler(client, cache, wfile) -> gateway_proxy._ProxyHandler:
     h = object.__new__(gateway_proxy._ProxyHandler)
     h.client = client
@@ -396,6 +412,19 @@ def test_success_first_try_does_not_refresh(self):
         assert client.sent_tokens == ["Bearer tok1"]
         assert b"hi" in bytes(out.data)
 
+    def test_failed_refresh_surfaces_reauth_hint(self, capsys):
+        # When the forced refresh itself fails (the Databricks OAuth session is
+        # dead, not just the access token), the user must be told to run
+        # `databricks auth login` rather than left with a bare 401 that reads as
+        # an Anthropic `/login` prompt. We still retry + relay whatever comes.
+        client = _FakeClient([_FakeResp(401, b"a"), _FakeResp(401, b"b")])
+        cache = _RefreshFailsCache()
+        out = _Collect()
+        _handle_handler(client, cache, out)._handle()
+        assert cache.refreshed == 1  # a refresh was attempted
+        assert "databricks auth login" in capsys.readouterr().err
+        assert b"401" in bytes(out.data)  # the response is still relayed
+
 
 class TestStartProxyPortFallback:
     def test_falls_back_to_free_port_when_cached_port_busy(self, monkeypatch):