test(wizard): cover _validate_and_report branches; inject completion for DI

litellm/setup_wizard.py

@@ -487,7 +487,11 @@ class SetupWizard:
             return _KeyUnverified(reason=f"{type(exc).__name__}: {exc}")
 
     @staticmethod
-    def _validate_and_report(provider: Dict, api_key: str) -> str:
+    def _validate_and_report(
+        provider: Dict,
+        api_key: str,
+        completion: Callable[..., object] = litellm.completion,
+    ) -> str:
         """
         Validate credentials with a live completion and print the result.
         Offers a re-entry loop on failure. Returns the final (possibly re-entered) key.
@@ -501,7 +505,7 @@ class SetupWizard:
                 f"  {grey('Testing connection to ' + provider['name'] + '...')}",
                 flush=True,
             )
-            result = SetupWizard._classify_key(test_model, api_key, litellm.completion)
+            result = SetupWizard._classify_key(test_model, api_key, completion)
             match result:
                 case _KeyValid():
                     print(f"  {green(_CHECK)} {bold(provider['name'])} connected successfully")

tests/test_litellm/test_setup_wizard.py

@@ -249,6 +249,86 @@ def test_classify_key_unverified_on_generic_error():
     assert "not found" in result.reason
 
 
+def _auth_error():
+    return AuthenticationError(
+        message="invalid key",
+        llm_provider="gemini",
+        model="gemini/gemini-3.5-flash",
+    )
+
+
+_GEMINI = {
+    "name": "Google Gemini",
+    "env_key": "GEMINI_API_KEY",
+    "key_hint": "AIza...",
+    "test_model": "gemini/gemini-3.5-flash",
+}
+
+
+def test_validate_and_report_skips_when_no_test_model():
+    """Providers without a test_model (Azure/Bedrock/Ollama) return the key untouched."""
+    calls = {"n": 0}
+
+    def fake_completion(**kwargs):
+        calls["n"] += 1
+
+    key = SetupWizard._validate_and_report({"name": "Azure", "test_model": None}, "az-key", fake_completion)
+    assert key == "az-key"
+    assert calls["n"] == 0  # no validation attempted
+
+
+def test_validate_and_report_valid_returns_key(capsys):
+    key = SetupWizard._validate_and_report(_GEMINI, "good-key", lambda **_: object())
+    assert key == "good-key"
+    assert "connected successfully" in capsys.readouterr().out
+
+
+def test_validate_and_report_unverified_surfaces_reason(capsys, monkeypatch):
+    """A non-auth failure must print the real reason, not 'invalid API key'."""
+    monkeypatch.setattr("builtins.input", lambda *_: "n")  # decline re-entry
+
+    def fake_completion(**_):
+        raise RateLimitError(
+            message="Resource has been exhausted (check quota)",
+            llm_provider="gemini",
+            model="gemini/gemini-3.5-flash",
+        )
+
+    key = SetupWizard._validate_and_report(_GEMINI, "valid-but-throttled", fake_completion)
+    out = capsys.readouterr().out
+    assert key == "valid-but-throttled"
+    assert "could not verify key" in out
+    assert "check quota" in out
+    assert "invalid API key" not in out
+
+
+def test_validate_and_report_invalid_key_prompt(capsys, monkeypatch):
+    monkeypatch.setattr("builtins.input", lambda *_: "n")  # decline re-entry
+
+    key = SetupWizard._validate_and_report(_GEMINI, "bad-key", lambda **_: (_ for _ in ()).throw(_auth_error()))
+    out = capsys.readouterr().out
+    assert key == "bad-key"
+    assert "invalid API key" in out
+
+
+def test_validate_and_report_reentry_accepts_new_key(monkeypatch):
+    """On re-entry the newly typed key is validated and returned once it works."""
+    inputs = iter(["y", "good-key"])  # yes re-enter, then the replacement key
+    monkeypatch.setattr("builtins.input", lambda *_: next(inputs))
+
+    seen = {"keys": []}
+
+    def fake_completion(**kwargs):
+        seen["keys"].append(kwargs["api_key"])
+        if kwargs["api_key"] == "bad-key":
+            raise _auth_error()
+        return object()
+
+    key = SetupWizard._validate_and_report(_GEMINI, "bad-key", fake_completion)
+    assert key == "good-key"
+    assert seen["keys"] == ["bad-key", "good-key"]
+
+
 def test_run_setup_wizard_enables_debug_when_requested(monkeypatch):
     import litellm
     import litellm.setup_wizard as wiz