Wizard/wizard_ai/tests/test_ai.py

import json
import random

import pytest

from wizard_game import GameOptions, WizardGame

from wizard_ai.agents import HeuristicAgent, ModelAgent, RandomAgent, make_agent
from wizard_ai.model import FEATURE_SIZE, ActionSpace, UniformPolicy, encode_observation, load_policy
from wizard_ai.training import SelfPlayTrainer, TrainingConfig, WizardEnv, evaluate


@pytest.mark.parametrize("num_players", [3, 4, 5, 6])
def test_encoding_has_fixed_size_and_masks_match(num_players):
    game = WizardGame(num_players, GameOptions(hidden_bids=True), seed=num_players)
    rng = random.Random(0)
    while not game.is_finished:
        obs = game.observation(game.current_player)
        assert len(encode_observation(obs)) == FEATURE_SIZE
        mask = ActionSpace.mask(obs)
        legal = {ActionSpace.index(a) for a in obs.legal_actions}
        assert {i for i, allowed in enumerate(mask) if allowed} == legal
        game.apply(obs.player, rng.choice(obs.legal_actions))


def test_action_space_round_trip():
    for index in range(ActionSpace.SIZE):
        assert ActionSpace.index(ActionSpace.action(index)) == index


@pytest.mark.parametrize("agent", [HeuristicAgent(), RandomAgent(1), ModelAgent(UniformPolicy(1))])
def test_agents_play_legal_games(agent):
    game = WizardGame(4, seed=3)
    while not game.is_finished:
        obs = game.observation(game.current_player)
        action = agent.act(obs)
        assert action in obs.legal_actions
        game.apply(obs.player, action)


def test_env_rewards_add_up_to_scores():
    env = WizardEnv(3, seed=5)
    obs = env.reset()
    totals = [0.0] * env.num_players
    rng = random.Random(5)
    while obs is not None:
        result = env.step(rng.choice(obs.legal_actions))
        totals = [t + r for t, r in zip(totals, result.rewards)]
        obs = result.observation
    assert [round(t * 100) for t in totals] == env.game.scores


def test_trainer_runs_and_checkpoints(tmp_path):
    config = TrainingConfig(episodes=4, player_counts=(3, 4), episodes_per_update=2,
                            checkpoint_dir=tmp_path / "ckpt", seed=1)
    trainer = SelfPlayTrainer(UniformPolicy(1), config)
    trajectories, scores = trainer.play_episode()
    assert len(trajectories) == len(scores)
    assert all(t[-1].done for t in trajectories)
    assert sum(sum(tr.reward for tr in t) for t in trajectories) == pytest.approx(sum(scores) / 100)
    trainer.train()
    assert json.loads((tmp_path / "ckpt" / "policy.json").read_text())["type"] == "uniform"
    assert isinstance(load_policy(tmp_path / "ckpt"), UniformPolicy)


def test_trainer_against_heuristic_opponents():
    config = TrainingConfig(episodes=2, player_counts=(5,), opponents="heuristic", seed=2)
    trajectories, _ = SelfPlayTrainer(UniformPolicy(2), config).play_episode()
    assert len(trajectories) == 1


def test_make_agent_falls_back_without_checkpoint(tmp_path):
    assert isinstance(make_agent("model", tmp_path / "missing"), HeuristicAgent)
    UniformPolicy().save(tmp_path / "ckpt")
    assert isinstance(make_agent("model", tmp_path / "ckpt"), ModelAgent)


def test_heuristic_beats_random():
    result = evaluate([HeuristicAgent(), RandomAgent(1), RandomAgent(2), RandomAgent(3)], games=40, seed=0)
    assert result.mean_scores[0] > max(result.mean_scores[1:])


# --- Service limits -------------------------------------------------------------

import asyncio  # noqa: E402
import threading  # noqa: E402

from wizard_ai.service import AIService  # noqa: E402


class BlockingAgent(RandomAgent):
    """Blocks in act() until released, to simulate a slow model."""

    name = "slow-model"

    def __init__(self):
        super().__init__(seed=0)
        self.release = threading.Event()

    def act(self, observation):
        self.release.wait(5)
        return super().act(observation)


def decide_message(request_id, timeout=5.0):
    game = WizardGame(3, seed=1)
    return {"type": "decide", "request_id": request_id, "timeout": timeout,
            "observation": game.observation(game.current_player).to_dict()}


def test_service_answers_with_agent_name():
    service = AIService("ws://unused", HeuristicAgent())
    reply = asyncio.run(service.handle_request(decide_message("r1")))
    assert reply["type"] == "action" and reply["request_id"] == "r1" and reply["agent"] == "heuristic"


def test_service_refuses_when_queue_is_full_and_degrades_when_busy():
    async def scenario():
        slow = BlockingAgent()
        service = AIService("ws://unused", slow, fallback=HeuristicAgent(), workers=1, max_queue=3, degrade_queue=1)
        first = asyncio.create_task(service.handle_request(decide_message("a")))
        await asyncio.sleep(0.05)  # "a" is computing with the slow agent
        second = asyncio.create_task(service.handle_request(decide_message("b")))
        third = asyncio.create_task(service.handle_request(decide_message("c")))
        await asyncio.sleep(0.05)
        refused = await service.handle_request(decide_message("d"))
        assert refused == {"type": "error", "request_id": "d", "message": "overloaded"}
        slow.release.set()
        replies = await asyncio.gather(first, second, third)
        # "b" starts while "c" is still waiting -> degraded; "c" runs when nothing waits anymore.
        assert [r["agent"] for r in replies] == ["slow-model", "heuristic", "slow-model"]
        assert service.in_progress == 0

    asyncio.run(scenario())


def test_service_drops_expired_requests():
    async def scenario():
        slow = BlockingAgent()
        service = AIService("ws://unused", slow, workers=1)
        first = asyncio.create_task(service.handle_request(decide_message("a")))
        await asyncio.sleep(0.05)
        waiting = asyncio.create_task(service.handle_request(decide_message("b", timeout=0.1)))
        await asyncio.sleep(0.2)  # "b" expires while "a" blocks the only worker
        slow.release.set()
        await first
        assert (await waiting)["message"] == "expired"

    asyncio.run(scenario())


def test_service_uses_fallback_when_primary_is_too_slow():
    async def scenario():
        service = AIService("ws://unused", RandomAgent(1), fallback=HeuristicAgent())
        service._primary_latency = 3.0  # model needs 3 s, server waits 5 s -> more than half
        reply = await service.handle_request(decide_message("a", timeout=5.0))
        assert reply["agent"] == "heuristic"

    asyncio.run(scenario())