55 lines
1.6 KiB
Python
55 lines
1.6 KiB
Python
import pytest
|
|
from generals.agents import ExpanderAgent, HunterAgent, RandomAgent
|
|
|
|
from general_bots_training.opponents import ModelOpponent, SelfPlayOpponent, make_opponent
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("name", "expected_type"),
|
|
[
|
|
("random", RandomAgent),
|
|
("expander", ExpanderAgent),
|
|
("hunter", HunterAgent),
|
|
],
|
|
)
|
|
def test_make_opponent(name, expected_type):
|
|
assert isinstance(make_opponent(name), expected_type)
|
|
|
|
|
|
def test_make_opponent_is_case_insensitive():
|
|
assert isinstance(make_opponent("HUNTER"), HunterAgent)
|
|
|
|
|
|
def test_make_opponent_supports_self_play_aliases():
|
|
assert isinstance(make_opponent("self_play"), SelfPlayOpponent)
|
|
assert isinstance(make_opponent("self-play"), SelfPlayOpponent)
|
|
|
|
|
|
def test_make_opponent_rejects_unknown_name():
|
|
with pytest.raises(ValueError, match="unknown opponent"):
|
|
make_opponent("turtle")
|
|
|
|
|
|
def test_make_opponent_model_requires_checkpoint():
|
|
with pytest.raises(ValueError, match="requires a checkpoint path"):
|
|
make_opponent("model")
|
|
|
|
|
|
def test_make_opponent_model_loads_checkpoint(tmp_path):
|
|
import equinox as eqx
|
|
import jax.random as jrandom
|
|
|
|
from general_bots_training.network import PolicyValueNetwork
|
|
|
|
network = PolicyValueNetwork(jrandom.PRNGKey(0), in_channels=14)
|
|
checkpoint_path = tmp_path / "model.eqx"
|
|
eqx.tree_serialise_leaves(checkpoint_path, network)
|
|
|
|
opponent = make_opponent(f"model:{checkpoint_path}")
|
|
|
|
assert isinstance(opponent, ModelOpponent)
|
|
|
|
|
|
def test_make_opponent_rejects_checkpoint_on_non_model():
|
|
with pytest.raises(ValueError, match="does not take a checkpoint"):
|
|
make_opponent("random:some.eqx")
|