tests/test_engine_worker.py - master - keorlov/aichess

tests/test_engine_worker.py
229 строк8 KB
Konstantin Orlov
wip
06 июн 2026, 13:19
06 июн 2026, 13:19cb645b7
"""Tests for engine subprocess isolation and security model."""

import os
import tempfile
import textwrap

import pytest

from src.engine_discovery import discover_engines
from src.match import EngineWorker, _call_engine, play_match


def test_discover_engines_does_not_execute_engine_code(tmp_path, capsys):
    """Discovery must not exec engine.py — a side-effecting engine is only
    run when actually invoked, never at startup. We assert the side effect
    never appears, confirming no execution happened.
    """
    engine_dir = tmp_path / "EvilBot"
    engine_dir.mkdir()
    (engine_dir / "engine.py").write_text(
        "SIDE_EFFECT_MARKER = 'pwned'\n"
        "def choose_move(fen):\n"
        "    return 'e2e4'\n"
    )
    discover_engines(str(tmp_path))
    captured = capsys.readouterr()
    assert "pwned" not in captured.out


def test_discover_engines_warns_on_missing_choose_move(tmp_path, capsys):
    engine_dir = tmp_path / "NoFuncBot"
    engine_dir.mkdir()
    (engine_dir / "engine.py").write_text("x = 1\n")
    engines = discover_engines(str(tmp_path))
    assert "NoFuncBot" in engines
    assert "no choose_move" in capsys.readouterr().out


def test_discover_engines_warns_on_syntax_error(tmp_path, capsys):
    engine_dir = tmp_path / "BrokenBot"
    engine_dir.mkdir()
    (engine_dir / "engine.py").write_text("def choose_move(:\n  pass\n")
    discover_engines(str(tmp_path))
    assert "syntax error" in capsys.readouterr().out


def test_engine_worker_reuses_subprocess_across_moves():
    """A single EngineWorker instance should answer multiple FENs without
    spawning a new process per call. We verify by checking the underlying
    process is the same object across two requests.
    """
    engines = discover_engines("engines")
    path = engines["FirstMoveBot"]

    worker = EngineWorker(path, "FirstMoveBot")
    try:
        proc_id_1 = id(worker._proc)
        m1, e1 = worker.get_move(
            "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1", 5.0
        )
        assert e1 is None
        assert isinstance(m1, str)

        proc_id_2 = id(worker._proc)
        m2, e2 = worker.get_move(
            "rnbqkbnr/pppppppp/8/8/4P3/8/PPPP1PPP/RNBQKBNR b KQkq e3 0 1", 5.0
        )
        assert e2 is None
        assert isinstance(m2, str)

        # Same worker object, same underlying subprocess.
        assert proc_id_1 == proc_id_2
        assert worker._proc.is_alive()
    finally:
        worker.close()


def test_engine_worker_recovers_from_timeout():
    """After a timeout the worker must be restarted and the next call
    must succeed normally.
    """
    engines = discover_engines("engines")
    sleepy = engines["SleepyBot"]
    first_move = engines["FirstMoveBot"]

    worker = EngineWorker(sleepy, "SleepyBot")
    try:
        m, err = worker.get_move(
            "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1", 0.1
        )
        assert m is None
        assert err == "timeout"

        # The worker should have restarted itself; next call should still time out
        # (SleepyBot always times out), but the error should be a clean "timeout"
        # rather than a worker death.
        m2, err2 = worker.get_move(
            "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1", 0.1
        )
        assert m2 is None
        assert err2 == "timeout"
    finally:
        worker.close()


def test_play_match_respects_colors_parameter():
    """When colors=('a','a') is passed, engine_a plays white in BOTH games.
    Default colors=('a','b') gives the historical colour-swap behaviour.
    """
    engines = discover_engines("engines")
    result = play_match(
        "FirstMoveBot", engines["FirstMoveBot"],
        "RandomBot", engines["RandomBot"],
        move_timeout=5.0,
        colors=("a", "a"),
    )
    # With colors=('a','a'), engine_a (FirstMoveBot) is white in both games
    assert result.games[0].white == "FirstMoveBot"
    assert result.games[1].white == "FirstMoveBot"
    assert result.games[0].black == "RandomBot"
    assert result.games[1].black == "RandomBot"


def test_play_match_default_colors_swap():
    """Default colors=('a','b') means engine_a white in game 1, black in game 2.
    """
    engines = discover_engines("engines")
    result = play_match(
        "FirstMoveBot", engines["FirstMoveBot"],
        "RandomBot", engines["RandomBot"],
        move_timeout=5.0,
    )
    assert result.games[0].white == "FirstMoveBot"
    assert result.games[0].black == "RandomBot"
    assert result.games[1].white == "RandomBot"
    assert result.games[1].black == "FirstMoveBot"


def test_engine_receives_move_timeout(tmp_path):
    """Engine's choose_move must receive the move_timeout the host passed."""
    engine_dir = tmp_path / "TimeoutSpyBot"
    engine_dir.mkdir()
    (engine_dir / "engine.py").write_text(textwrap.dedent("""\
        import json, os
        def choose_move(fen, move_timeout):
            with open(os.path.join(os.path.dirname(__file__), "seen.txt"), "w") as f:
                f.write(f"{move_timeout:.4f}")
            return "e2e4"
    """))
    worker = EngineWorker(str(engine_dir / "engine.py"), "TimeoutSpyBot")
    try:
        m, err = worker.get_move(
            "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1", 2.5
        )
        assert err is None
        assert m == "e2e4"
        seen = (engine_dir / "seen.txt").read_text()
        assert abs(float(seen) - 2.5) < 1e-6
    finally:
        worker.close()


def test_engine_backward_compatible_old_signature(tmp_path):
    """Engines with the old choose_move(fen) signature must still work."""
    engine_dir = tmp_path / "LegacyBot"
    engine_dir.mkdir()
    (engine_dir / "engine.py").write_text(textwrap.dedent("""\
        def choose_move(fen):
            return "e2e4"
    """))
    worker = EngineWorker(str(engine_dir / "engine.py"), "LegacyBot")
    try:
        m, err = worker.get_move(
            "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1", 1.0
        )
        assert err is None
        assert m == "e2e4"
    finally:
        worker.close()


def test_engine_receives_fraction_of_timeout(tmp_path):
    """With engine_timeout_fraction=0.8 and host timeout=2.0, engine must
    receive 1.6 (0.8 * 2.0), not 2.0."""
    engine_dir = tmp_path / "FractionSpyBot"
    engine_dir.mkdir()
    (engine_dir / "engine.py").write_text(textwrap.dedent("""\
        import os
        def choose_move(fen, move_timeout):
            with open(os.path.join(os.path.dirname(__file__), "seen.txt"), "w") as f:
                f.write(f"{move_timeout:.4f}")
            return "e2e4"
    """))
    worker = EngineWorker(
        str(engine_dir / "engine.py"), "FractionSpyBot",
        engine_timeout_fraction=0.8,
    )
    try:
        m, err = worker.get_move(
            "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1", 2.0
        )
        assert err is None
        assert m == "e2e4"
        seen = (engine_dir / "seen.txt").read_text()
        assert abs(float(seen) - 1.6) < 1e-6
    finally:
        worker.close()


def test_engine_default_fraction_is_one(tmp_path):
    """With no fraction specified, engine receives the full timeout (1.0 * t)."""
    engine_dir = tmp_path / "FullTimeoutBot"
    engine_dir.mkdir()
    (engine_dir / "engine.py").write_text(textwrap.dedent("""\
        import os
        def choose_move(fen, move_timeout):
            with open(os.path.join(os.path.dirname(__file__), "seen.txt"), "w") as f:
                f.write(f"{move_timeout:.4f}")
            return "e2e4"
    """))
    worker = EngineWorker(str(engine_dir / "engine.py"), "FullTimeoutBot")
    try:
        worker.get_move(
            "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1", 1.5
        )
        seen = (engine_dir / "seen.txt").read_text()
        assert abs(float(seen) - 1.5) < 1e-6
    finally:
        worker.close()