Stabilize Gemini auth recovery e2e

Author: AarushiShah-db

tests/test_e2e.py

@@ -10,9 +10,12 @@
 
 from __future__ import annotations
 
+import json
 import os
 import shutil
 import subprocess
+from urllib import error as urllib_error
+from urllib import request as urllib_request
 
 import pytest
 
@@ -57,6 +60,39 @@ def _run_agent(
     )
 
 
+def _run_gemini_gateway_smoke(workspace: str, model: str, token: str) -> str:
+    """Call the Gemini gateway directly with a text-only prompt.
+
+    This keeps auth recovery coverage focused on the recovered Databricks token
+    instead of Gemini CLI's separate tool-calling request shape.
+    """
+    url = f"{build_tool_base_url('gemini', workspace)}/v1beta/models/{model}:generateContent"
+    payload = {
+        "contents": [
+            {"role": "user", "parts": [{"text": "say hi in 5 words or less"}]},
+        ],
+    }
+    req = urllib_request.Request(
+        url,
+        data=json.dumps(payload).encode("utf-8"),
+        headers={
+            "Authorization": f"Bearer {token}",
+            "Content-Type": "application/json",
+            "Accept": "application/json",
+        },
+        method="POST",
+    )
+    try:
+        with urllib_request.urlopen(req, timeout=30) as response:
+            body = response.read().decode("utf-8")
+    except urllib_error.HTTPError as exc:
+        body = exc.read().decode("utf-8", errors="replace") if exc.fp else ""
+        raise AssertionError(f"Gemini gateway smoke failed: HTTP {exc.code}: {body[:500]}") from exc
+
+    data = json.loads(body)
+    return data.get("candidates", [{}])[0].get("content", {}).get("parts", [{}])[0].get("text", "")
+
+
 # ---------------------------------------------------------------------------
 # Databricks auth / token
 # ---------------------------------------------------------------------------
@@ -878,11 +914,4 @@ def test_recovers_when_initial_token_empty(
             "get_databricks_token may not be retrying after auth login."
         )
 
-        env = gemini.build_runtime_env(e2e_workspace, model, recovered_token)
-        cmd = gemini.validate_cmd("gemini")
-        result = _run_agent(cmd, env=env, timeout=90)
-        combined = (result.stdout + result.stderr).strip()
-        assert result.returncode == 0 and combined, (
-            f"Gemini failed after auth recovery: rc={result.returncode} "
-            f"stdout={result.stdout[:300]!r} stderr={result.stderr[:300]!r}"
-        )
+        assert _run_gemini_gateway_smoke(e2e_workspace, model, recovered_token).strip()