Pyra/tests/unit/test_setup_wizard.py

"""Tests for setup wizard personalization and model-discovery helpers."""

from unittest.mock import MagicMock

import pytest


def test_use_case_plugin_mapping_all_categories_have_entries():
    from pyra.setup.wizard import _USE_CASE_PLUGINS
    assert all(len(v) > 0 for v in _USE_CASE_PLUGINS.values())


def test_use_case_plugin_mapping_has_expected_categories():
    from pyra.setup.wizard import _USE_CASE_PLUGINS
    assert "Email" in _USE_CASE_PLUGINS
    assert "Development & servers" in _USE_CASE_PLUGINS
    assert "Research & web" in _USE_CASE_PLUGINS


def test_use_case_email_contains_email_plugin():
    from pyra.setup.wizard import _USE_CASE_PLUGINS
    assert "email" in _USE_CASE_PLUGINS["Email"]


def test_use_case_dev_contains_ssh_and_docker():
    from pyra.setup.wizard import _USE_CASE_PLUGINS
    assert "ssh_tool" in _USE_CASE_PLUGINS["Development & servers"]
    assert "docker_tool" in _USE_CASE_PLUGINS["Development & servers"]


def test_use_case_file_management_contains_cloud_stores():
    from pyra.setup.wizard import _USE_CASE_PLUGINS
    plugins = _USE_CASE_PLUGINS["File management"]
    assert "gdrive" in plugins
    assert "onedrive" in plugins
    assert "dropbox_tool" in plugins


def test_suggest_plugins_empty_use_cases_returns_early(monkeypatch):
    calls = []
    import pyra.setup.wizard as wiz
    monkeypatch.setattr(wiz.console, "print", lambda *a, **kw: calls.append(a))
    wiz._suggest_plugins([])
    assert calls == []


def test_suggest_plugins_unknown_use_case_returns_early(monkeypatch):
    calls = []
    import pyra.setup.wizard as wiz
    monkeypatch.setattr(wiz.console, "print", lambda *a, **kw: calls.append(a))
    wiz._suggest_plugins(["Not a real category"])
    assert calls == []


def test_suggest_plugins_valid_use_case_calls_print(monkeypatch):
    calls = []
    import pyra.setup.wizard as wiz
    monkeypatch.setattr(wiz.console, "print", lambda *a, **kw: calls.append(str(a)))
    wiz._suggest_plugins(["Email"])
    assert len(calls) > 0


def test_suggest_plugins_panel_text_contains_plugin_name(monkeypatch):
    from rich.panel import Panel
    panels = []

    import pyra.setup.wizard as wiz

    def capture_print(*args, **kwargs):
        for a in args:
            if isinstance(a, Panel):
                panels.append(a.renderable)

    monkeypatch.setattr(wiz.console, "print", capture_print)
    wiz._suggest_plugins(["Email"])
    assert any("email" in str(p) for p in panels)


def test_suggest_plugins_multiple_categories(monkeypatch):
    from rich.panel import Panel
    panels = []

    import pyra.setup.wizard as wiz

    def capture_print(*args, **kwargs):
        for a in args:
            if isinstance(a, Panel):
                panels.append(a.renderable)

    monkeypatch.setattr(wiz.console, "print", capture_print)
    wiz._suggest_plugins(["Email", "Development & servers"])
    combined = " ".join(str(p) for p in panels)
    assert "email" in combined
    assert "ssh_tool" in combined


# ── _fetch_local_models ────────────────────────────────────────────────────────

def test_fetch_local_models_lmstudio_returns_model_ids(monkeypatch):
    import pyra.setup.wizard as wiz
    mock_resp = MagicMock()
    mock_resp.json.return_value = {"data": [{"id": "gemma-4b"}, {"id": "llama3"}]}
    mock_resp.raise_for_status = lambda: None
    monkeypatch.setattr(wiz.httpx, "get", lambda *a, **kw: mock_resp)
    from pyra.setup.providers import get_provider
    assert wiz._fetch_local_models(get_provider("lmstudio")) == ["gemma-4b", "llama3"]


def test_fetch_local_models_ollama_returns_model_names(monkeypatch):
    import pyra.setup.wizard as wiz
    mock_resp = MagicMock()
    mock_resp.json.return_value = {"models": [{"name": "llama3:latest"}, {"name": "mistral"}]}
    mock_resp.raise_for_status = lambda: None
    monkeypatch.setattr(wiz.httpx, "get", lambda *a, **kw: mock_resp)
    from pyra.setup.providers import get_provider
    assert wiz._fetch_local_models(get_provider("ollama")) == ["llama3:latest", "mistral"]


def test_fetch_local_models_returns_empty_on_connection_error(monkeypatch):
    import pyra.setup.wizard as wiz
    monkeypatch.setattr(wiz.httpx, "get", MagicMock(side_effect=Exception("conn refused")))
    from pyra.setup.providers import get_provider
    assert wiz._fetch_local_models(get_provider("lmstudio")) == []


def test_fetch_local_models_returns_empty_when_no_base_url():
    import pyra.setup.wizard as wiz
    from pyra.setup.providers import Provider
    provider = Provider(
        id="test", display_name="Test", requires_key=False,
        default_model="x", litellm_prefix="openai/", group="Local",
    )
    assert wiz._fetch_local_models(provider) == []


# ── _fetch_lmstudio_available_models ──────────────────────────────────────────

def test_fetch_lmstudio_available_models_returns_ids(monkeypatch):
    import pyra.setup.wizard as wiz
    mock_resp = MagicMock()
    mock_resp.json.return_value = {"data": [{"id": "model-a"}, {"id": "model-b"}]}
    mock_resp.raise_for_status = lambda: None
    monkeypatch.setattr(wiz.httpx, "get", lambda *a, **kw: mock_resp)
    assert wiz._fetch_lmstudio_available_models() == ["model-a", "model-b"]


def test_fetch_lmstudio_available_models_returns_empty_on_error(monkeypatch):
    import pyra.setup.wizard as wiz
    monkeypatch.setattr(wiz.httpx, "get", MagicMock(side_effect=Exception("not found")))
    assert wiz._fetch_lmstudio_available_models() == []


def test_fetch_lmstudio_available_models_empty_data(monkeypatch):
    import pyra.setup.wizard as wiz
    mock_resp = MagicMock()
    mock_resp.json.return_value = {"data": []}
    mock_resp.raise_for_status = lambda: None
    monkeypatch.setattr(wiz.httpx, "get", lambda *a, **kw: mock_resp)
    assert wiz._fetch_lmstudio_available_models() == []


# ── _load_lmstudio_model ──────────────────────────────────────────────────────

def test_load_lmstudio_model_returns_true_on_success(monkeypatch):
    import pyra.setup.wizard as wiz
    mock_resp = MagicMock()
    mock_resp.is_success = True
    monkeypatch.setattr(wiz.httpx, "post", lambda *a, **kw: mock_resp)
    assert wiz._load_lmstudio_model("gemma-4b") is True


def test_load_lmstudio_model_returns_false_on_api_failure(monkeypatch):
    import pyra.setup.wizard as wiz
    mock_resp = MagicMock()
    mock_resp.is_success = False
    monkeypatch.setattr(wiz.httpx, "post", lambda *a, **kw: mock_resp)
    assert wiz._load_lmstudio_model("gemma-4b") is False


def test_load_lmstudio_model_returns_false_on_exception(monkeypatch):
    import pyra.setup.wizard as wiz
    monkeypatch.setattr(wiz.httpx, "post", MagicMock(side_effect=Exception("timeout")))
    assert wiz._load_lmstudio_model("gemma-4b") is False


# ── _classify_error ───────────────────────────────────────────────────────────

def _fake_llm_exc(name: str) -> Exception:
    """Create a fake litellm exception with the given class name."""
    cls = type(name, (Exception,), {"__module__": "litellm.exceptions"})
    return cls("test error")


def test_classify_auth_error():
    from pyra.setup.wizard import _classify_error
    label, hint = _classify_error(_fake_llm_exc("AuthenticationError"))
    assert "key" in label.lower()
    assert "provider" in hint.lower() or "dashboard" in hint.lower()


def test_classify_not_found_error():
    from pyra.setup.wizard import _classify_error
    label, hint = _classify_error(_fake_llm_exc("NotFoundError"))
    assert "model" in label.lower()
    assert "model" in hint.lower()


def test_classify_rate_limit_error():
    from pyra.setup.wizard import _classify_error
    label, _ = _classify_error(_fake_llm_exc("RateLimitError"))
    assert "rate" in label.lower()


def test_classify_service_unavailable():
    from pyra.setup.wizard import _classify_error
    label, _ = _classify_error(_fake_llm_exc("ServiceUnavailableError"))
    assert "unavailable" in label.lower()


def test_classify_api_connection_error():
    from pyra.setup.wizard import _classify_error
    label, hint = _classify_error(_fake_llm_exc("APIConnectionError"))
    assert "reach" in label.lower() or "network" in label.lower() or "connect" in label.lower()
    assert "internet" in hint.lower() or "network" in hint.lower()


def test_classify_timeout_error():
    from pyra.setup.wizard import _classify_error
    label, _ = _classify_error(_fake_llm_exc("TimeoutError"))
    assert "time" in label.lower()


def test_classify_bad_request_error():
    from pyra.setup.wizard import _classify_error
    label, hint = _classify_error(_fake_llm_exc("BadRequestError"))
    assert "request" in label.lower() or "bad" in label.lower()
    assert "model" in hint.lower()


def test_classify_httpx_connect_error():
    import httpx
    from pyra.setup.wizard import _classify_error
    label, hint = _classify_error(httpx.ConnectError("refused"))
    assert "reach" in label.lower() or "reachable" in label.lower()
    assert "running" in hint.lower() or "listening" in hint.lower()


def test_classify_httpx_timeout():
    import httpx
    from pyra.setup.wizard import _classify_error
    label, _ = _classify_error(httpx.TimeoutException("timeout"))
    assert "time" in label.lower()


def test_classify_generic_error():
    from pyra.setup.wizard import _classify_error
    label, hint = _classify_error(ValueError("something went wrong"))
    assert len(label) > 0
    assert "something went wrong" in hint


def test_classify_error_always_returns_two_strings():
    from pyra.setup.wizard import _classify_error
    for exc in (RuntimeError("boom"), OSError("disk"), KeyError("k")):
        label, hint = _classify_error(exc)
        assert isinstance(label, str) and len(label) > 0
        assert isinstance(hint, str) and len(hint) > 0


# ── _check_local_server retry behaviour ──────────────────────────────────────

def _silent_print(*a, **kw):
    pass


def test_check_local_server_success(monkeypatch):
    import pyra.setup.wizard as wiz
    mock_resp = MagicMock()
    mock_resp.raise_for_status = lambda: None
    monkeypatch.setattr(wiz.httpx, "get", lambda *a, **kw: mock_resp)
    monkeypatch.setattr(wiz.console, "print", _silent_print)
    from pyra.setup.providers import get_provider
    wiz._check_local_server(get_provider("lmstudio"))  # must not raise


def test_check_local_server_retry_then_success(monkeypatch):
    import pyra.setup.wizard as wiz
    from pyra.setup.providers import get_provider

    call_count = {"n": 0}

    def flaky_get(*a, **kw):
        call_count["n"] += 1
        if call_count["n"] == 1:
            raise ConnectionError("refused")
        m = MagicMock()
        m.raise_for_status = lambda: None
        return m

    monkeypatch.setattr(wiz.httpx, "get", flaky_get)
    monkeypatch.setattr(wiz.console, "print", _silent_print)
    monkeypatch.setattr(wiz.questionary, "select",
                        lambda *a, **kw: MagicMock(ask=lambda: "retry"))

    wiz._check_local_server(get_provider("lmstudio"))
    assert call_count["n"] == 2


def test_check_local_server_abort_raises_system_exit(monkeypatch):
    import pyra.setup.wizard as wiz
    from pyra.setup.providers import get_provider

    monkeypatch.setattr(wiz.httpx, "get", MagicMock(side_effect=ConnectionError("refused")))
    monkeypatch.setattr(wiz.console, "print", _silent_print)
    monkeypatch.setattr(wiz.questionary, "select",
                        lambda *a, **kw: MagicMock(ask=lambda: "abort"))

    with pytest.raises(SystemExit):
        wiz._check_local_server(get_provider("lmstudio"))


def test_check_local_server_continue_returns(monkeypatch):
    import pyra.setup.wizard as wiz
    from pyra.setup.providers import get_provider

    monkeypatch.setattr(wiz.httpx, "get", MagicMock(side_effect=ConnectionError("refused")))
    monkeypatch.setattr(wiz.console, "print", _silent_print)
    monkeypatch.setattr(wiz.questionary, "select",
                        lambda *a, **kw: MagicMock(ask=lambda: "continue"))

    wiz._check_local_server(get_provider("lmstudio"))  # must return without raising