scanengine-3/tests/test_sras_format.py

"""core.sras_format vs the pre-refactor golden fixtures.

The goldens were produced by the original sc3_aui_app implementation; the
extracted module must reproduce them byte-for-byte (writer) and
field-for-field (parser + frontier walk).
"""
import json
from dataclasses import asdict
from pathlib import Path

import numpy as np
import pytest

from core.scan_geometry import build_plan
from core.sras_format import SrasFile, create_scan_file
from golden_util import (
    BACKGROUND, CHANNELS, LASER_FREQ_HZ, PREAMBLES, SAMPLE_RATE, SPF,
    TINY_PLAN_ARGS, VELOCITY_MM_S, synthetic_frame,
)

GOLDEN = Path(__file__).parent / "golden"


@pytest.fixture(scope="module")
def expected():
    with open(GOLDEN / "sras_expected.json") as f:
        return json.load(f)


def _tiny_plan():
    return build_plan(**TINY_PLAN_ARGS, laser_freq_hz=LASER_FREQ_HZ,
                      velocity_mm_s=VELOCITY_MM_S)


def _write_complete(path):
    plan = _tiny_plan()
    f = create_scan_file(path, plan, SPF, SAMPLE_RATE, PREAMBLES, BACKGROUND)
    try:
        for ai, pa in enumerate(plan.per_angle):
            for ri in range(pa.n_rows):
                for ci in range(len(CHANNELS)):
                    for fi in range(pa.n_frames):
                        f.write(synthetic_frame(ai, ri, ci, fi))
    finally:
        f.close()


def test_writer_byte_identical_to_golden(tmp_path):
    out = tmp_path / "rewrite.sras"
    _write_complete(out)
    assert out.read_bytes() == (GOLDEN / "complete.sras").read_bytes()


def test_header_matches_golden(expected):
    sras = SrasFile(GOLDEN / "complete.sras")
    h = expected["header"]
    assert asdict(sras.header) == {
        "n_angles": h["n_angles"],
        "x_start_nominal": h["x_start_nominal"], "y_start_nominal": h["y_start_nominal"],
        "x_delta_nominal": h["x_delta_nominal"], "y_delta_nominal": h["y_delta_nominal"],
        "row_spacing": h["row_spacing"], "velocity": h["velocity"],
        "laser_freq": h["laser_freq"],
        "samples_per_frame": h["samples_per_frame"], "sample_rate": h["sample_rate"],
        "bytes_per_sample": h["bytes_per_sample"], "n_channels": h["n_channels"],
    }
    assert sras.data_start_offset == h["data_start_offset"]
    assert [pa.angle_deg for pa in sras.per_angle] == h["angles"]
    for pa, exp in zip(sras.per_angle, h["per_angle"], strict=True):
        assert pa.angle_deg == exp["angle"]
        assert pa.x_start == exp["x_start"]
        assert pa.x_delta == exp["x_delta"]
        assert pa.n_frames == exp["n_frames"]
        assert pa.n_rows == exp["n_rows"]
        assert pa.y_positions == exp["y_positions"]


def test_frontier_all_truncation_variants(expected):
    for name, exp_statuses in expected["statuses"].items():
        statuses = SrasFile(GOLDEN / name).angle_status()
        assert [asdict(s) for s in statuses] == exp_statuses, f"mismatch for {name}"


def test_preambles_and_background_roundtrip():
    sras = SrasFile(GOLDEN / "complete.sras")
    assert sras.preambles == PREAMBLES
    assert sras.background == BACKGROUND


def test_load_angle_memmap_equals_eager():
    with SrasFile(GOLDEN / "complete.sras") as sras:
        raw = (GOLDEN / "complete.sras").read_bytes()
        for ai, pa in enumerate(sras.per_angle):
            view = sras.load_angle(ai)
            h = sras.header
            assert view.shape == (pa.n_rows, h.n_channels, pa.n_frames, h.samples_per_frame)
            start = sras.angle_data_offset(ai)
            eager = np.frombuffer(
                raw, dtype=np.int8, offset=start, count=view.size
            ).reshape(view.shape)
            assert np.array_equal(view, eager)
            assert not view.flags.writeable


def test_load_row_matches_synthetic_pattern():
    with SrasFile(GOLDEN / "complete.sras") as sras:
        for ai in range(2):
            for ri in range(3):
                for ci in range(3):
                    row = sras.load_row(ai, ri, ci)
                    expected_bytes = b"".join(
                        synthetic_frame(ai, ri, ci, fi)
                        for fi in range(sras.per_angle[ai].n_frames)
                    )
                    assert row.tobytes() == expected_bytes


def test_truncated_load_angle_partial_rows():
    sras = SrasFile(GOLDEN / "trunc_rowboundary_a1.sras")
    st = sras.angle_status()[1]
    assert st.status == "TRUNCATED" and st.n_rows_available == 2
    view = sras.load_angle(1, n_rows=st.n_rows_available)
    assert view.shape[0] == 2
    sras.close()


def test_bad_magic_and_version_rejected(tmp_path):
    bad = tmp_path / "bad.sras"
    bad.write_bytes(b"XXXX" + bytes(60))
    with pytest.raises(ValueError, match="bad magic"):
        SrasFile(bad)

    data = bytearray((GOLDEN / "complete.sras").read_bytes())
    data[4] = 5  # version byte
    v5 = tmp_path / "v5.sras"
    v5.write_bytes(bytes(data))
    with pytest.raises(ValueError, match="version 5"):
        SrasFile(v5)