ARC_AGI_Wayfinder_Agent / tests /test_perception.py
samyakbayar's picture
Upload 29 files
d3a24e0 verified
Raw
History Blame Contribute Delete
2.28 kB
"""Tests for the perception encoder module."""
from __future__ import annotations
import numpy as np
import pytest
from agents.wayfinder.perception import PerceptionEncoder, GRID_SIZE, NUM_COLORS
class TestPerceptionEncoder:
"""Test cases for PerceptionEncoder."""
@pytest.fixture
def encoder(self) -> PerceptionEncoder:
"""Create a test encoder instance."""
return PerceptionEncoder(latent_dim=64, device="cpu")
@pytest.fixture
def sample_frame(self) -> np.ndarray:
"""Create a sample 64×64 frame."""
return np.random.randint(0, NUM_COLORS, size=(GRID_SIZE, GRID_SIZE), dtype=np.uint8)
def test_encode_returns_correct_shapes(self, encoder: PerceptionEncoder, sample_frame: np.ndarray) -> None:
"""Test that encode returns latent and diff_mask with correct shapes."""
latent, diff_mask = encoder.encode(sample_frame)
assert latent.shape == (64,)
assert diff_mask.shape == (GRID_SIZE, GRID_SIZE)
assert diff_mask.dtype == bool
def test_first_frame_diff_is_zero(self, encoder: PerceptionEncoder, sample_frame: np.ndarray) -> None:
"""Test that the first frame has an all-zero diff mask."""
_, diff_mask = encoder.encode(sample_frame)
assert not np.any(diff_mask)
def test_second_frame_diff_detects_changes(self, encoder: PerceptionEncoder, sample_frame: np.ndarray) -> None:
"""Test that diff mask detects changes between frames."""
encoder.encode(sample_frame)
modified = sample_frame.copy()
modified[0, 0] = (modified[0, 0] + 1) % NUM_COLORS
_, diff_mask = encoder.encode(modified)
assert diff_mask[0, 0]
assert not diff_mask[1, 1] # Unchanged pixel
def test_encode_batch(self, encoder: PerceptionEncoder) -> None:
"""Test batch encoding."""
frames = np.random.randint(0, NUM_COLORS, size=(4, GRID_SIZE, GRID_SIZE), dtype=np.uint8)
latents = encoder.encode_batch(frames)
assert latents.shape == (4, 64)
def test_latent_is_float32(self, encoder: PerceptionEncoder, sample_frame: np.ndarray) -> None:
"""Test that the latent is float32."""
latent, _ = encoder.encode(sample_frame)
assert latent.dtype == np.float32