lf_mri_platform/services/reconstructor/test.py

import importlib
import io
import json
import os

import pytest
from fastapi.testclient import TestClient


@pytest.fixture()
def client(tmp_path, monkeypatch):
    app_mod = importlib.import_module("app")

    app_mod.STORE_DIR = str(tmp_path)
    os.makedirs(app_mod.STORE_DIR, exist_ok=True)

    class DummyRecoApp:
        def __init__(self, name, digit, shift):
            self.name = name
            self.digit = digit
            self.shift = shift

        def start_reconstruction(self, path_raw_data, path_np_data_json, path_order_json):
            import matplotlib.pyplot as plt
            plt.figure()
            plt.plot([0, 1], [0, 1])
            plt.title(f"{self.name}-{self.digit}")
            plt.savefig("data/test.png")
            plt.close()

    monkeypatch.setattr(app_mod, "ReconstructionApp", DummyRecoApp, raising=True)

    return TestClient(app_mod.app)


def _upload(client: TestClient, name: str, content: bytes, mime: str = "application/octet-stream"):
    files = {"file": (name, io.BytesIO(content), mime)}
    r = client.post("/upload", files=files)
    assert r.status_code == 200, r.text
    obj = r.json()
    assert "file_id" in obj
    return obj["file_id"]


def test_upload_json(client: TestClient):
    params = {
        "Np": 8, "Nf": 8, "sl_nb": 1,
        "contrasts": 1, "RF_spoil": 0,
        "ETL": 0, "N_TE": 0, "phi_wing": 0, "N_wings": 0,
        "D_scans": 1
    }
    fid = _upload(client, "params.json", json.dumps(params).encode("utf-8"), "application/json")
    assert isinstance(fid, str) and len(fid) > 0


def test_upload_h5(client: TestClient):
    fake = b"\x89HDF\r\n\x1a\n" + b"\x00" * 64
    fid = _upload(client, "raw.h5", fake, "application/octet-stream")
    assert isinstance(fid, str) and len(fid) > 0


def test_full_reconstruction_flow(client: TestClient, tmp_path):
    raw_bytes = b"\x89HDF\r\n\x1a\n" + b"\x00" * 64
    raw_id = _upload(client, "raw.h5", raw_bytes)

    params = {
        "Np": 8, "Nf": 8, "sl_nb": 1,
        "contrasts": 1, "RF_spoil": 0,
        "ETL": 0, "N_TE": 0, "phi_wing": 0, "N_wings": 0,
        "D_scans": 1
    }
    json_id = _upload(client, "params.json", json.dumps(params).encode("utf-8"), "application/json")

    payload = {
        "file_raw_id": raw_id,
        "file_json_id": json_id,
        "file_order_id": None,
        "sequence_name": "linear_decart",
        "digit": "2d",
        "phase_shift": True
    }
    r = client.post("/reconstruct", json=payload)
    assert r.status_code == 200, r.text
    job = r.json()
    job_id = job["job_id"]

    for _ in range(50):
        s = client.get(f"/jobs/{job_id}")
        assert s.status_code == 200
        st = s.json()
        if st["status"] in ("done", "error"):
            break

    assert st["status"] == "done", f"job failed: {st.get('error_traceback')}"

    lf = client.get(f"/jobs/{job_id}/files")
    assert lf.status_code == 200
    files = lf.json()["files"]
    assert any(name.endswith(".png") for name in files), f"no png files in {files}"

    dz = client.get(f"/jobs/{job_id}/archive.zip")
    assert dz.status_code == 200
    assert dz.headers.get("content-type") == "application/zip"


def test_reconstruct_bad_ids(client: TestClient):
    payload = {
        "file_raw_id": "nope",
        "file_json_id": "nope",
        "file_order_id": None,
        "sequence_name": "linear_decart",
        "digit": "2d",
        "phase_shift": False
    }
    r = client.post("/reconstruct", json=payload)
    assert r.status_code == 400