tests.py

import random
from typing import Dict, List, Optional, Tuple

from engine import (
    AMMO_RESTORE,
    BULLET_DAMAGE,
    HEALTH_RESTORE,
    HP_CAP,
    MAX_TURNS,
    STARTING_AMMO,
    STARTING_HP,
    MoveAction,
    PassAction,
    ShootAction,
    apply_actions,
    is_valid_action,
    make_snapshot,
    parse_map,
    render_board,
    run_match,
)

def create_custom_map(p0_pos: Tuple[int, int], p1_pos: Tuple[int, int], powerups: Optional[Dict[Tuple[int, int], str]] = None, walls: Optional[List[Tuple[int, int]]] = None) -> str:
    grid = [["." for _ in range(15)] for _ in range(15)]

    for r, c in walls or []: grid[r][c] = "#"
    for (r, c), powerup_type in (powerups or {}).items(): grid[r][c] = powerup_type

    grid[p0_pos[0]][p0_pos[1]] = "1"
    grid[p1_pos[0]][p1_pos[1]] = "2"
    return "\n".join("".join(row) for row in grid)

def assert_tank(state, player_id: int, *, hp=None, ammo=None, pos=None, facing=None) -> None:
    tank = state.tanks[player_id]
    if hp is not None:
        assert tank.hp == hp, f"P{player_id} HP: expected {hp}, got {tank.hp}"
    if ammo is not None:
        assert tank.ammo == ammo, f"P{player_id} ammo: expected {ammo}, got {tank.ammo}"
    if pos is not None:
        assert tank.pos == pos, f"P{player_id} position: expected {pos}, got {tank.pos}"
    if facing is not None:
        assert tank.facing == facing, f"P{player_id} facing: expected {facing}, got {tank.facing}"

def assert_value_error(function, message: str) -> None:
    try:
        function()
    except ValueError:
        return
    raise AssertionError(message)

def test_action_validation() -> None:
    # 1.1 PASS is always valid
    state = parse_map(create_custom_map((7, 7), (0, 0)))
    assert is_valid_action(state, 0, PassAction())

    # 1.2 Moving into a wall counts as PASS
    state = parse_map(create_custom_map((7, 7), (0, 0), walls=[(7, 8)]))
    updated, _ = apply_actions(
        state, MoveAction("right"), PassAction(), resolution_order=(0, 1)
    )
    assert_tank(updated, 0, pos=(7, 7))

    # 1.3 Moving outside the map counts as PASS
    state = parse_map(create_custom_map((0, 0), (14, 14)))
    updated, _ = apply_actions(
        state, MoveAction("up"), PassAction(), resolution_order=(0, 1)
    )
    assert_tank(updated, 0, pos=(0, 0))

    # 1.4 Moving into enenmy tile counts as PASS
    state = parse_map(create_custom_map((7, 7), (7, 8)))
    updated, _ = apply_actions(
        state, MoveAction("right"), PassAction(), resolution_order=(0, 1)
    )
    assert_tank(updated, 0, pos=(7, 7))

    # 1.5 Garbage input results in PASS
    state = parse_map(create_custom_map((7, 7), (0, 0)))
    updated, _ = apply_actions(
        state, MoveAction("diagonal"), PassAction(), resolution_order=(0, 1)
    )
    assert_tank(updated, 0, pos=(7, 7))

    state = parse_map(create_custom_map((7, 0), (7, 7)))
    updated, _ = apply_actions(
        state, ShootAction("diagonal"), PassAction(), resolution_order=(0, 1)
    )
    assert_tank(updated, 0, ammo=STARTING_AMMO)
    assert_tank(updated, 1, hp=STARTING_HP)

    # 1.6 Shooting with zero ammo results in PASS
    state = parse_map(create_custom_map((7, 0), (7, 7)))
    state.tanks[0].ammo = 0
    updated, _ = apply_actions(
        state, ShootAction("right"), PassAction(), resolution_order=(0, 1)
    )
    assert_tank(updated, 0, ammo=0)
    assert_tank(updated, 1, hp=STARTING_HP)    

def test_movement_resolution() -> None:
    # 2.1 P0 vacates tile, and P1 enters tile (P0 resolved first)
    state = parse_map(create_custom_map((7, 7), (7, 8)))
    updated, _ = apply_actions(
        state, MoveAction("left"), MoveAction("left"), resolution_order=(0, 1)
    )
    assert_tank(updated, 0, pos=(7, 6), facing="left")
    assert_tank(updated, 1, pos=(7, 7), facing="left")

    # 2.2 P0 and P1 try to move to same tile, first resolved player succeeds
    # 2.2.0 P0 first resolved.
    state = parse_map(create_custom_map((7, 6), (7, 8)))
    updated, _ = apply_actions(
        state, MoveAction("right"), MoveAction("left"), resolution_order=(0, 1)
    )
    assert_tank(updated, 0, pos=(7, 7))
    assert_tank(updated, 1, pos=(7, 8))

    # 2.2.1 P1 first resolved.
    state = parse_map(create_custom_map((7, 6), (7, 8)))
    updated, _ = apply_actions(
        state, MoveAction("right"), MoveAction("left"), resolution_order=(1, 0)
    )
    assert_tank(updated, 0, pos=(7, 6))
    assert_tank(updated, 1, pos=(7, 7))

    # 2.3 Both tanks try to swap position
    state = parse_map(create_custom_map((7, 7), (7, 8)))
    updated, _ = apply_actions(
        state, MoveAction("right"), MoveAction("left"), resolution_order=(0, 1)
    )
    assert_tank(updated, 0, pos=(7, 7))
    assert_tank(updated, 1, pos=(7, 8))

    state = parse_map(create_custom_map((7, 7), (7, 8)))
    updated, _ = apply_actions(
        state, MoveAction("right"), MoveAction("left"), resolution_order=(1, 0)
    )
    assert_tank(updated, 0, pos=(7, 7))
    assert_tank(updated, 1, pos=(7, 8))

def test_bullet_behaviour() -> None:
    # 3.1 Bullet is shot, costs one ammo, hits enemy and deals BULLET_DAMAGE
    state = parse_map(create_custom_map((7, 0), (7, 7)))
    updated, _ = apply_actions(
        state, ShootAction("right"), PassAction(), resolution_order=(0, 1)
    )
    assert_tank(updated, 0, ammo=STARTING_AMMO-1, facing="right")
    assert_tank(updated, 1, hp=STARTING_HP-BULLET_DAMAGE)

    # 3.2 Bullet is fired but stopped by wall. Costs 1 ammo 
    state = parse_map(create_custom_map((7, 0), (7, 7), walls=[(7, 3)]))
    updated, _ = apply_actions(
        state, ShootAction("right"), PassAction(), resolution_order=(0, 1)
    )
    assert_tank(updated, 0, ammo=STARTING_AMMO-1)
    assert_tank(updated, 1, hp=STARTING_HP)

def test_dodge_mechanics() -> None:
    # 4.1 P1 shoots in old location of P0. Bullet misses
    state = parse_map(create_custom_map((7, 7), (7, 10)))
    updated, _ = apply_actions(
        state, MoveAction("up"), ShootAction("left"), resolution_order=(0, 1)
    )
    assert_tank(updated, 0, hp=STARTING_HP, pos=(6, 7))
    assert_tank(updated, 1, ammo=STARTING_AMMO-1)

    state = parse_map(create_custom_map((7, 7), (7, 10)))
    updated, _ = apply_actions(
        state, MoveAction("up"), ShootAction("left"), resolution_order=(1, 0)
    )
    assert_tank(updated, 0, hp=STARTING_HP-BULLET_DAMAGE, pos=(6, 7))

    # 4.2 P0 shoots before P1 moves. Bullet hits
    state = parse_map(create_custom_map((7, 7), (7, 10)))
    updated, _ = apply_actions(
        state, ShootAction("right"), MoveAction("up"), resolution_order=(0, 1)
    )
    assert_tank(updated, 1, hp=STARTING_HP-BULLET_DAMAGE, pos=(6, 10))

    state = parse_map(create_custom_map((7, 7), (7, 10)))
    updated, _ = apply_actions(
        state, ShootAction("right"), MoveAction("up"), resolution_order=(1, 0)
    )
    assert_tank(updated, 0, ammo=STARTING_AMMO-1)
    assert_tank(updated, 1, hp=STARTING_HP, pos=(6, 10))

def test_kill_shots() -> None:
    # 5.1 P0 shoots P1. P1 dies and its shot is not fired.
    state = parse_map(create_custom_map((7, 7), (7, 10)))
    state.tanks[1].hp = BULLET_DAMAGE
    updated, result = apply_actions(
        state, ShootAction("right"), ShootAction("left"), resolution_order=(0, 1)
    )
    assert result == "PLAYER 0 WINS!"
    assert_tank(updated, 0, hp=STARTING_HP, ammo=STARTING_AMMO - 1)
    assert_tank(updated, 1, hp=0, ammo=STARTING_AMMO)

    # Mirror of 5.1
    state = parse_map(create_custom_map((7, 7), (7, 10)))
    state.tanks[0].hp = BULLET_DAMAGE
    updated, result = apply_actions(
        state, ShootAction("right"), ShootAction("left"), resolution_order=(1, 0)
    )
    assert result == "PLAYER 1 WINS!"
    assert_tank(updated, 0, hp=0, ammo=STARTING_AMMO)
    assert_tank(updated, 1, hp=STARTING_HP, ammo=STARTING_AMMO - 1)

    # 5.2 Killing shot invalidates move onto health pack
    state = parse_map(create_custom_map((7, 5), (7, 7), powerups={(6, 7): "H"}))
    state.tanks[1].hp = BULLET_DAMAGE
    updated, result = apply_actions(
        state, ShootAction("right"), MoveAction("up"), resolution_order=(0, 1)
    )
    assert result == "PLAYER 0 WINS!"
    assert_tank(updated, 1, hp=0, pos=(7, 7))
    assert (6, 7) in updated.powerups

def test_powerup_mechanics() -> None:
    # 6.1 Player picks up health pack
    state = parse_map(create_custom_map((7, 7), (0, 0), powerups={(7, 8): "H"}))
    state.tanks[0].hp = 50
    updated, _ = apply_actions(
        state, MoveAction("right"), PassAction(), resolution_order=(0, 1)
    )
    assert_tank(updated, 0, hp=50 + HEALTH_RESTORE, pos=(7, 8))
    assert (7, 8) not in updated.powerups

    # 6.2 HP Cap enforcement
    state = parse_map(create_custom_map((7, 7), (0, 0), powerups={(7, 8): "H"}))
    state.tanks[0].hp = 90
    updated, _ = apply_actions(
        state, MoveAction("right"), PassAction(), resolution_order=(0, 1)
    )
    assert_tank(updated, 0, hp=HP_CAP)

    # 6.3 No Ammo Cap enforcement
    state = parse_map(create_custom_map((7, 7), (0, 0), powerups={(7, 8): "A"}))
    state.tanks[0].ammo = 15
    updated, _ = apply_actions(
        state, MoveAction("right"), PassAction(), resolution_order=(0, 1)
    )
    assert_tank(updated, 0, ammo=15 + AMMO_RESTORE)
    assert (7, 8) not in updated.powerups

    # 6.4 Contested health pack (resolution order dependent)
    state = parse_map(create_custom_map((7, 6), (7, 8), powerups={(7, 7): "H"}))
    state.tanks[0].hp = 50
    updated, _ = apply_actions(
        state, MoveAction("right"), MoveAction("left"), resolution_order=(0, 1)
    )
    assert_tank(updated, 0, hp=50 + HEALTH_RESTORE, pos=(7, 7))
    assert_tank(updated, 1, pos=(7, 8))
    assert (7, 7) not in updated.powerups

    state = parse_map(create_custom_map((7, 6), (7, 8), powerups={(7, 7): "H"}))
    state.tanks[1].hp = 50
    updated, _ = apply_actions(
        state, MoveAction("right"), MoveAction("left"), resolution_order=(1, 0)
    )
    assert_tank(updated, 0, pos=(7, 6))
    assert_tank(updated, 1, hp=50 + HEALTH_RESTORE, pos=(7, 7))
    assert (7, 7) not in updated.powerups

def test_turn_limit_rules() -> None:
    # 7.1 Winner after 150 turns decided by greater HP. If same HP, match is drawn
    state = parse_map(create_custom_map((0, 0), (14, 14)))
    state.turn = MAX_TURNS
    state.tanks[0].hp, state.tanks[1].hp = 80, 60
    _, result = apply_actions(
        state, PassAction(), PassAction(), resolution_order=(0, 1)
    )
    assert result == "PLAYER 0 WINS!"

    state = parse_map(create_custom_map((0, 0), (14, 14)))
    state.turn = MAX_TURNS
    _, result = apply_actions(
        state, PassAction(), PassAction(), resolution_order=(0, 1)
    )
    assert result == "DRAW"

def test_injection_protection() -> None:
    # 8.1 Both bots correctly receive snapshot, and cannot modify the state
    state = parse_map(
        create_custom_map((7, 1), (7, 13), powerups={(5, 5): "H", (9, 9): "A"})
    )
    snapshot = make_snapshot(state, 0)
    assert snapshot["grid"] == state.grid
    assert snapshot["powerups"] == state.powerups
    assert snapshot["opponent_pos"] == state.tanks[1].pos

    snapshot["my_hp"] = 9999
    snapshot["grid"][0][1] = "#"
    snapshot["powerups"].clear()
    assert state.tanks[0].hp == STARTING_HP
    assert state.grid[0][1] == "."
    assert state.powerups == {(5, 5): "H", (9, 9): "A"}

    original = parse_map(create_custom_map((7, 7), (0, 0), powerups={(7, 8): "A"}))
    updated, _ = apply_actions(
        original, MoveAction("right"), PassAction(), resolution_order=(0, 1)
    )
    assert_tank(original, 0, pos=(7, 7), ammo=STARTING_AMMO)
    assert original.powerups == {(7, 8): "A"}
    assert_tank(updated, 0, pos=(7, 8), ammo=STARTING_AMMO + AMMO_RESTORE)
    assert updated.powerups == {}

def test_bot_information() -> None:
    observations = []

    def p0_bot(snapshot):
        observations.append(("p0", snapshot["my_pos"], snapshot["opponent_pos"]))
        return MoveAction("left")

    def p1_bot(snapshot):
        observations.append(("p1", snapshot["my_pos"], snapshot["opponent_pos"]))
        return PassAction()

    run_match(
        p0_bot,
        p1_bot,
        create_custom_map((7, 7), (7, 8)),
        resolution_seed=0,
    )

    # 9.1 No bot can see the other bot's current turn.
    assert observations[0] == ("p0", (7, 7), (7, 8))
    assert observations[1] == ("p1", (7, 8), (7, 7))

def test_resolution_rng() -> None:
    def pass_bot(_snapshot):
        return PassAction()

    map_str = create_custom_map((7, 1), (7, 13))

    first = run_match(pass_bot, pass_bot, map_str, resolution_seed=69420)
    second = run_match(pass_bot, pass_bot, map_str, resolution_seed=69420)
    assert first == second

    # 10.1 Resolution does not use global RNG
    random.seed(676767)
    expected_first = random.random()
    expected_second = random.random()

    random.seed(676767)
    actual_first = random.random()
    run_match(pass_bot, pass_bot, map_str, resolution_seed=998244353)
    actual_second = random.random()

    assert actual_first == expected_first
    assert actual_second == expected_second

def run_tests() -> None:
    test_action_validation()
    test_movement_resolution()
    test_bullet_behaviour()
    test_dodge_mechanics()
    test_kill_shots()
    test_powerup_mechanics()
    test_turn_limit_rules()
    test_injection_protection()
    test_bot_information()
    test_resolution_rng()
    print("All tests passed.")


if __name__ == "__main__":
    run_tests()