tests/test_clients.py

from clients import art_client_from_env, boss_client_from_env, card_client_from_env
from art import DiffusersImageClient, LazyArtClient
from boss import NemotronBossClient
from generator import LlamaCppCardClient, MiniCPMCardClient
from local_llm import LocalCompletionClient, LocalJsonChatClient, nemotron_prompt


class FakeChat:
    # Complete one fake prompt.
    def complete(self, system: str, user: str) -> str:
        return "{}"


class FakeArtBackend:
    # Return a deterministic image URI.
    def create_art(self, prompt: str) -> str:
        return f"uri:{prompt}"


# Verify card client is disabled without an endpoint.
def test_card_client_from_env_disabled(monkeypatch) -> None:
    monkeypatch.delenv("TABRAS_CARD_ENDPOINT", raising=False)
    assert card_client_from_env() is None


# Verify card client reads endpoint environment.
def test_card_client_from_env(monkeypatch) -> None:
    monkeypatch.setenv("TABRAS_CARD_ENDPOINT", "http://cards")
    monkeypatch.setenv("TABRAS_CARD_MODEL", "mini")
    client = card_client_from_env()
    assert isinstance(client, MiniCPMCardClient)
    assert client.chat.endpoint == "http://cards"
    assert client.chat.model == "mini"


# Verify card client can use the direct MiniCPM Transformers backend.
def test_card_client_from_env_transformers(monkeypatch) -> None:
    monkeypatch.setenv("TABRAS_CARD_BACKEND", "transformers")
    monkeypatch.setenv("TABRAS_CARD_MODEL", "mini-local")
    monkeypatch.setattr("clients.MiniCPMTransformersChatClient.load", lambda model_path, **kwargs: FakeChat())
    client = card_client_from_env()
    assert isinstance(client, MiniCPMCardClient)
    assert isinstance(client.chat, FakeChat)


# Verify card client can use a llama.cpp MiniCPM backend.
def test_card_client_from_env_llamacpp(monkeypatch) -> None:
    monkeypatch.setenv("TABRAS_CARD_BACKEND", "llamacpp")
    monkeypatch.setenv("TABRAS_CARD_ENDPOINT", "http://cards")
    monkeypatch.setenv("TABRAS_CARD_MODEL", "mini-q4")
    client = card_client_from_env()
    assert isinstance(client, LlamaCppCardClient)
    assert isinstance(client.chat, LocalJsonChatClient)
    assert client.chat.endpoint == "http://cards"
    assert client.chat.model == "mini-q4"
    assert client.chat.max_tokens == 320


# Verify boss client is disabled unless enabled.
def test_boss_client_from_env_disabled(monkeypatch) -> None:
    monkeypatch.delenv("TABRAS_AI_BOSS", raising=False)
    assert boss_client_from_env() is None


# Verify boss client reads endpoint environment.
def test_boss_client_from_env(monkeypatch) -> None:
    monkeypatch.setenv("TABRAS_AI_BOSS", "1")
    monkeypatch.setenv("TABRAS_BOSS_ENDPOINT", "http://boss")
    monkeypatch.setenv("TABRAS_BOSS_MODEL", "nemotron")
    client = boss_client_from_env()
    assert isinstance(client, NemotronBossClient)
    assert client.chat.endpoint == "http://boss"
    assert client.chat.model == "nemotron"


# Verify boss client can use a Nemotron formatted completion endpoint.
def test_boss_client_from_env_completion(monkeypatch) -> None:
    monkeypatch.setenv("TABRAS_AI_BOSS", "1")
    monkeypatch.setenv("TABRAS_BOSS_BACKEND", "completion")
    monkeypatch.setenv("TABRAS_BOSS_ENDPOINT", "http://boss-completion")
    client = boss_client_from_env()
    assert isinstance(client, NemotronBossClient)
    assert isinstance(client.chat, LocalCompletionClient)
    assert client.chat.endpoint == "http://boss-completion"
    assert client.chat.prompt_template is nemotron_prompt


# Verify boss client can use the direct Nemotron Transformers backend.
def test_boss_client_from_env_transformers(monkeypatch) -> None:
    monkeypatch.setenv("TABRAS_AI_BOSS", "1")
    monkeypatch.setenv("TABRAS_BOSS_BACKEND", "transformers")
    monkeypatch.setenv("TABRAS_BOSS_MODEL", "nemotron-local")
    monkeypatch.setattr("clients.NemotronTransformersChatClient.load", lambda *_args, **_kwargs: FakeChat())
    client = boss_client_from_env()
    assert isinstance(client, NemotronBossClient)
    assert isinstance(client.chat, FakeChat)


# Verify boss client can use a quantized MLX backend.
def test_boss_client_from_env_mlx(monkeypatch) -> None:
    monkeypatch.setenv("TABRAS_AI_BOSS", "1")
    monkeypatch.setenv("TABRAS_BOSS_BACKEND", "mlx")
    monkeypatch.setenv("TABRAS_BOSS_MODEL", "nemotron-mlx")
    monkeypatch.setattr("clients.MLXChatClient.load", lambda *_args, **_kwargs: FakeChat())
    client = boss_client_from_env()
    assert isinstance(client, NemotronBossClient)
    assert isinstance(client.chat, FakeChat)


# Verify art client is disabled unless explicitly enabled.
def test_art_client_from_env_disabled(monkeypatch) -> None:
    monkeypatch.delenv("TABRAS_ART_BACKEND", raising=False)
    monkeypatch.setattr("clients._art_client_cache", None)
    assert art_client_from_env() is None


# Verify art client loads and caches the SD-Turbo diffusers backend.
def test_art_client_from_env_diffusers(monkeypatch) -> None:
    monkeypatch.setenv("TABRAS_ART_BACKEND", "diffusers")
    monkeypatch.setenv("TABRAS_ART_MODEL", "sdxl-local")
    monkeypatch.setenv("TABRAS_ART_STEPS", "1")
    monkeypatch.setenv("TABRAS_ART_WIDTH", "128")
    monkeypatch.setenv("TABRAS_ART_HEIGHT", "96")
    monkeypatch.setattr("clients._art_client_cache", None)

    calls = []

    def fake_load(model_id: str, **kwargs):
        calls.append(model_id)
        assert model_id == "sdxl-local"
        assert kwargs["steps"] == 1
        assert kwargs["width"] == 128
        assert kwargs["height"] == 96
        return FakeArtBackend()

    monkeypatch.setattr(DiffusersImageClient, "load", staticmethod(fake_load))
    client = art_client_from_env()
    assert isinstance(client, LazyArtClient)
    assert art_client_from_env() is client
    assert calls == []
    assert client.create_art("spell") == "uri:spell"
    assert calls == ["sdxl-local"]


# Verify art client defaults use the lazy SDXL-Turbo backend.
def test_art_client_from_env_fast_defaults(monkeypatch) -> None:
    monkeypatch.setenv("TABRAS_ART_BACKEND", "diffusers")
    monkeypatch.delenv("TABRAS_ART_MODEL", raising=False)
    monkeypatch.delenv("TABRAS_ART_STEPS", raising=False)
    monkeypatch.delenv("TABRAS_ART_WIDTH", raising=False)
    monkeypatch.delenv("TABRAS_ART_HEIGHT", raising=False)
    monkeypatch.setattr("clients._art_client_cache", None)

    def fake_load(model_id: str, **kwargs):
        assert model_id == "stabilityai/sdxl-turbo"
        assert kwargs["steps"] == 1
        assert kwargs["width"] == 512
        assert kwargs["height"] == 320
        return FakeArtBackend()

    monkeypatch.setattr(DiffusersImageClient, "load", staticmethod(fake_load))
    assert art_client_from_env().create_art("spark") == "uri:spark"