tests/test_boss.py

from typing import Any

from boss import (
    HeuristicBossClient,
    NemotronBossClient,
    affordable_indexes,
    boss_chooser,
    boss_indexes,
    boss_payload,
    boss_prompt,
    boss_system_prompt,
    hand_payload,
    parse_card_indexes,
    player_payload,
    safe_boss_indexes,
)
from budget import Card
from game import DuelState, create_player
from play import run_text_duel
from primitives import Effect


class FakeBossClient:
    # Initialize the fake boss client with a fixed raw response.
    def __init__(self, raw: dict[str, Any]) -> None:
        self.raw = raw
        self.payload: dict[str, Any] | None = None

    # Capture the payload and return a fixed boss response.
    def choose_cards(self, payload: dict[str, Any]) -> dict[str, Any]:
        self.payload = payload
        return self.raw


class FakeChat:
    # Initialize the fake chat client with fixed text.
    def __init__(self, text: str) -> None:
        self.text = text
        self.calls: list[tuple[str, str]] = []

    # Capture prompts and return fixed text.
    def complete(self, system: str, user: str) -> str:
        self.calls.append((system, user))
        return self.text


class BrokenBossClient:
    # Raise like a malformed model response.
    def choose_cards(self, payload: dict[str, Any]) -> dict[str, Any]:
        del payload
        raise ValueError("bad model output")


# Build one boss test card.
def make_card(name: str, cost: int, effect: Effect) -> Card:
    return Card(name, cost, "fire", "wuxia", (effect,))


# Build one boss test duel.
def make_duel() -> DuelState:
    return DuelState(create_player("You", []), create_player("Enemy", []))


# Verify boss payload exposes visible state.
def test_boss_payload() -> None:
    state = make_duel()
    state.enemy.energy = 3
    state.enemy.shield_charge = 2
    state.enemy.hand = [make_card("Strike", 1, Effect("deal", amount=2))]
    payload = boss_payload(state, state.enemy)
    assert payload["actor"]["shield_charge"] == 2
    assert payload["opponent"]["name"] == "You"
    assert payload["hand"][0]["rules_text"] == "Deal 2 damage."


# Verify heuristic boss client returns an empty model choice.
def test_heuristic_boss_client() -> None:
    assert HeuristicBossClient().choose_cards({}) == {"card_indexes": []}


# Verify boss prompts ask for strict JSON and playable indexes.
def test_boss_prompts() -> None:
    payload = {"playable_indexes": (0,), "actor": {}, "opponent": {}, "hand": []}
    assert "strict JSON" in boss_system_prompt()
    assert "playable_indexes" in boss_prompt(payload)


# Verify Nemotron boss client uses local chat completions.
def test_nemotron_boss_client() -> None:
    chat = FakeChat('{"card_indexes": [0], "reason": "lethal"}')
    client = NemotronBossClient(chat)  # type: ignore[arg-type]
    assert client.choose_cards({"playable_indexes": [0]})["card_indexes"] == [0]
    assert "strict JSON" in chat.calls[0][0]
    assert "card_indexes" in chat.calls[0][1]


# Verify player payload includes zone counts.
def test_player_payload() -> None:
    player = create_player("Enemy", [make_card("Deck", 1, Effect("deal", amount=1))])
    player.hand = [make_card("Hand", 1, Effect("deal", amount=1))]
    player.discard = [make_card("Discard", 1, Effect("deal", amount=1))]
    assert player_payload(player)["deck_count"] == 1
    assert player_payload(player)["hand_count"] == 1
    assert player_payload(player)["discard_count"] == 1


# Verify hand payload indexes playable cards.
def test_hand_payload() -> None:
    player = create_player("Enemy", [])
    player.hand = [make_card("Strike", 1, Effect("deal", amount=2))]
    assert hand_payload(player) == [{"index": 0, "name": "Strike", "cost": 1, "rules_text": "Deal 2 damage."}]


# Verify model indexes parse defensively.
def test_parse_card_indexes() -> None:
    assert parse_card_indexes({"card_indexes": [2, "x", 0]}) == (2, 0)
    assert parse_card_indexes({"card_indexes": "bad"}) == ()
    assert parse_card_indexes({}) == ()


# Verify affordability validation preserves sequence.
def test_affordable_indexes() -> None:
    actor = create_player("Enemy", [])
    actor.energy = 3
    actor.hand = [
        make_card("One", 1, Effect("deal", amount=1)),
        make_card("Three", 3, Effect("deal", amount=1)),
        make_card("Two", 2, Effect("deal", amount=1)),
    ]
    assert affordable_indexes(actor, (2, 0, 1, 2, 9)) == (2, 0)


# Verify boss indexes go through payload and validation.
def test_boss_indexes() -> None:
    state = make_duel()
    state.enemy.energy = 1
    state.enemy.hand = [make_card("Strike", 1, Effect("deal", amount=2))]
    client = FakeBossClient({"card_indexes": [0]})
    assert boss_indexes(client, state, state.enemy) == (0,)
    assert client.payload is not None
    assert client.payload["actor"]["name"] == "Enemy"


# Verify boss index safety converts model failures to no choice.
def test_safe_boss_indexes_handles_model_failure() -> None:
    state = make_duel()
    assert safe_boss_indexes(BrokenBossClient(), state, state.enemy) == ()


# Verify boss chooser falls back when model output is invalid.
def test_boss_chooser_falls_back() -> None:
    state = make_duel()
    state.enemy.energy = 1
    state.enemy.hand = [make_card("Strike", 1, Effect("deal", amount=2))]
    chooser = boss_chooser(FakeBossClient({"card_indexes": [9]}))
    assert chooser(state, state.enemy) == 0


# Verify boss chooser falls back when the model fails.
def test_boss_chooser_falls_back_on_model_failure() -> None:
    state = make_duel()
    state.enemy.energy = 1
    state.enemy.hand = [make_card("Strike", 1, Effect("deal", amount=2))]
    chooser = boss_chooser(BrokenBossClient())
    assert chooser(state, state.enemy) == 0


# Verify boss chooser uses a valid model card first.
def test_boss_chooser_uses_model_card() -> None:
    state = make_duel()
    state.enemy.energy = 2
    state.enemy.hand = [
        make_card("A", 1, Effect("deal", amount=1)),
        make_card("B", 2, Effect("block", amount=5)),
    ]
    chooser = boss_chooser(FakeBossClient({"card_indexes": [0]}))
    assert chooser(state, state.enemy) == 0


# Verify run_text_duel can use a boss client for enemy turns.
def test_run_text_duel_uses_boss_client() -> None:
    outputs: list[str] = []
    player_deck = tuple(make_card("Wait", 5, Effect("block", amount=1)) for _ in range(4))
    enemy_deck = tuple(make_card("Strike", 1, Effect("deal", amount=25)) for _ in range(4))
    result = run_text_duel(player_deck, enemy_deck, lambda _: "pass", outputs.append, max_rounds=1, boss_client=FakeBossClient({"card_indexes": [0]}))
    assert result == "Enemy"
    assert any("Enemy plays Strike" in line for line in outputs)