tests/test_state_codec.py

"""Unit tests for the state codec (no torch / GPU required)."""

import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from qwenjev.state_codec import (  # noqa: E402
    canonical_state, encode_observation, frame_hash, grid_diff, grid_to_image,
    grid_to_text,
)

GRID = [[0, 0, 3], [0, 5, 3], [1, 0, 0]]
GRID2 = [[0, 0, 3], [2, 5, 3], [1, 0, 0]]  # one cell changed (0,0)->? (1,0)


def test_grid_to_text_roundtrip_shape():
    text = grid_to_text(GRID)
    lines = text.split("\n")
    assert len(lines) == 3
    assert lines[0] == "003"


def test_diff_marks_changed_cells():
    diff = grid_diff(GRID, GRID2)
    changed_row = diff.split("\n")[1]
    assert changed_row[0] == "*"
    # unchanged cells keep their hex value
    assert diff.split("\n")[0] == "003"


def test_canonical_state_is_deterministic():
    a = canonical_state("003\n053", {"L": 2, "H": 3})
    b = canonical_state("003\n053", {"L": 2, "H": 3})
    assert a == b
    assert a.startswith("qwenjev-state-v1")


def test_frame_hash_stable():
    assert frame_hash(GRID) == frame_hash([list(r) for r in GRID])
    assert frame_hash(GRID) != frame_hash(GRID2)


def test_encode_observation_changed_flag():
    obs = encode_observation([GRID2], [GRID], levels_completed=1, win_levels=3)
    assert obs.changed is True
    assert obs.levels_completed == 1
    assert len(obs.frame_hashes) == 1


def test_grid_to_image_dimensions():
    img = grid_to_image(GRID, cell_px=2)
    assert img.size == (6, 6)