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