tests/test_oasis_player.py

from __future__ import annotations

import json
from pathlib import Path

from adam.oasis_player import capture_path, discover_models, frame_size, is_action_model, model_info, random_roblox_dataset_frame


def _model(root: Path, name: str = "Playable") -> Path:
    folder = root / "output_action_flow_models" / name
    (folder / "unet").mkdir(parents=True)
    (folder / "unet" / "config.json").write_text("{}", encoding="utf-8")
    (folder / "unet" / "diffusion_pytorch_model.safetensors").write_bytes(b"weights")
    (folder / "action_flow_model_info.json").write_text(json.dumps({
        "model_type": "action_conditioned_rectified_flow_video", "resolution": "256x144",
    }), encoding="utf-8")
    return folder


def _temporal_model(root: Path, name: str = "Temporal") -> Path:
    folder = _model(root, name)
    (folder / "vae").mkdir()
    (folder / "vae" / "config.json").write_text("{}", encoding="utf-8")
    (folder / "vae" / "pytorch_model.bin").write_bytes(b"weights")
    (folder / "action_flow_model_info.json").write_text(json.dumps({
        "model_type": "action_conditioned_temporal_latent_flow", "resolution": "256x144",
    }), encoding="utf-8")
    return folder


def test_native_player_discovers_only_complete_models(tmp_path: Path) -> None:
    playable = _model(tmp_path)
    incomplete = tmp_path / "output_action_flow_models" / "Incomplete"
    incomplete.mkdir()

    assert is_action_model(playable)
    assert not is_action_model(incomplete)
    assert discover_models(tmp_path) == [("Playable", playable)]


def test_native_player_reads_legacy_square_resolution_as_16_by_9(tmp_path: Path) -> None:
    playable = _model(tmp_path)
    payload = model_info(playable)
    assert frame_size(payload) == (256, 144)


def test_native_player_accepts_complete_temporal_model(tmp_path: Path) -> None:
    assert is_action_model(_temporal_model(tmp_path))


def test_native_player_rejects_temporal_model_without_vae(tmp_path: Path) -> None:
    temporal = _model(tmp_path)
    (temporal / "action_flow_model_info.json").write_text(json.dumps({
        "model_type": "action_conditioned_temporal_latent_flow", "resolution": "256x144",
    }), encoding="utf-8")
    assert not is_action_model(temporal)


def test_player_captures_are_adam_owned_and_unique(tmp_path: Path) -> None:
    first = capture_path(tmp_path)
    second = capture_path(tmp_path)
    assert first.parent == tmp_path / "data" / "oasis_captures"
    assert first != second


def test_player_finds_a_frame_through_nested_roblox_dataset_folders(tmp_path: Path) -> None:
    frame = tmp_path / "OldDatasets" / "Roblox Dataset" / "recording one" / "frames" / "frame_0001.png"
    frame.parent.mkdir(parents=True)
    frame.write_bytes(b"image")
    (frame.parent / "notes.txt").write_text("not a frame", encoding="utf-8")

    assert random_roblox_dataset_frame(tmp_path, chooser=lambda _count: 0) == frame