seq_UI_lfmri/test_spectrum/test.py

import os
import httpx

BASE_URL = "http://localhost:8000"
EXPORT_PATH = "downloaded"


def test_full_pipeline():
    with open("sample.csv", "rb") as f:
        upload = httpx.post(f"{BASE_URL}/upload/", files={"file": f})
    session_id = upload.json()["session_id"]

    filter_ = httpx.post(f"{BASE_URL}/filter/", data={
        "session_id": session_id,
        "dt": 1.125e-7,
        "center_freq": 2.97e6,
        "lower_freq": 2.92e6,
        "higher_freq": 3.02e6,
        "low_freq": 600e3
    })
    assert filter_.status_code == 200

    fft_ = httpx.post(f"{BASE_URL}/fft/", data={"session_id": session_id})
    assert fft_.status_code == 200

    result = httpx.get(f"{BASE_URL}/result/", params={"session_id": session_id})
    assert result.status_code == 200
    data = result.json()
    assert "fid" in data and "spectrum" in data and "frequency" in data


def test_export_and_save():
    os.makedirs(EXPORT_PATH, exist_ok=True)

    # upload CSV 
    with open("sample.csv", "rb") as f:
        r = httpx.post(f"{BASE_URL}/upload/", files={"file": f})
    session_id = r.json()["session_id"]

    # filter
    r = httpx.post(f"{BASE_URL}/filter/", data={
        "session_id": session_id,
        "dt": 1.125e-7,
        "center_freq": 2.97e6,
        "lower_freq": 2.92e6,
        "higher_freq": 3.02e6,
        "low_freq": 600e3
    })
    assert r.status_code == 200

    r = httpx.post(f"{BASE_URL}/fft/", data={"session_id": session_id})
    assert r.status_code == 200

    r = httpx.post(f"{BASE_URL}/export/", data={"session_id": session_id})
    assert r.status_code == 200
    data = r.json()

    for name, url in data.items():
        filename = url.split("/")[-1]
        full_url = f"{BASE_URL}{url}"
        response = httpx.get(full_url)
        assert response.status_code == 200
        file_path = os.path.join(EXPORT_PATH, filename)
        with open(file_path, "wb") as f:
            f.write(response.content)
        print(f"Saved: {file_path}")

    for filename in [f"{session_id}_fid.png", f"{session_id}_spectrum.png"]:
        assert os.path.exists(os.path.join(EXPORT_PATH, filename))


def test_export_and_save():
    os.makedirs(EXPORT_PATH, exist_ok=True)

    with open("sample2.csv", "rb") as f:
        r = httpx.post(f"{BASE_URL}/upload/", files={"file": f})
    session_id = r.json()["session_id"]

    r = httpx.post(f"{BASE_URL}/filter/", data={
        "session_id": session_id,
        "dt": 1.125e-7,
        "center_freq": 2.97e6,
        "lower_freq": 2.92e6,
        "higher_freq": 3.02e6,
        "low_freq": 600e3
    })
    assert r.status_code == 200

    r = httpx.post(f"{BASE_URL}/fft/", data={"session_id": session_id})
    assert r.status_code == 200

    r = httpx.post(f"{BASE_URL}/export/", data={"session_id": session_id})
    assert r.status_code == 200
    data = r.json()

    for name, url in data.items():
        filename = url.split("/")[-1]
        full_url = f"{BASE_URL}{url}"
        response = httpx.get(full_url)
        assert response.status_code == 200
        with open(os.path.join(EXPORT_PATH, filename), "wb") as f:
            f.write(response.content)

    for name in ["fid", "spectrum"]:
        assert os.path.exists(os.path.join(EXPORT_PATH, f"{session_id}_{name}.png"))


def test_plot_raw_generation():
    os.makedirs(EXPORT_PATH, exist_ok=True)

    with open("sample2.csv", "rb") as f:
        r = httpx.post(f"{BASE_URL}/upload/", files={"file": f})
    assert r.status_code == 200
    session_id = r.json()["session_id"]

    r = httpx.post(f"{BASE_URL}/plot-raw/", data={"session_id": session_id})
    assert r.status_code == 200
    raw_plot_url = r.json()["raw_plot"]
    filename = raw_plot_url.split("/")[-1]

    full_url = f"{BASE_URL}{raw_plot_url}"
    response = httpx.get(full_url)
    assert response.status_code == 200

    file_path = os.path.join(EXPORT_PATH, filename)
    with open(file_path, "wb") as f:
        f.write(response.content)

    assert os.path.exists(file_path)
    print(f"✓ Plot saved to {file_path}")

def test_export_and_save_json():
    os.makedirs(EXPORT_PATH, exist_ok=True)

    with open("test.json", "rb") as f:
        r = httpx.post(f"{BASE_URL}/upload/", files={"file": f})
    assert r.status_code == 200
    session_id = r.json()["session_id"]

    r = httpx.post(f"{BASE_URL}/filter/", data={
        "session_id": session_id,
        "dt": 1.125e-7,
        "center_freq": 2.97e6,
        "lower_freq": 2.92e6,
        "higher_freq": 3.02e6,
        "low_freq": 600e3
    })
    assert r.status_code == 200

    r = httpx.post(f"{BASE_URL}/fft/", data={"session_id": session_id})
    assert r.status_code == 200

    r = httpx.post(f"{BASE_URL}/export/", data={"session_id": session_id})
    assert r.status_code == 200
    data = r.json()

    for name, url in data.items():
        filename = url.split("/")[-1]
        full_url = f"{BASE_URL}{url}"
        response = httpx.get(full_url)
        assert response.status_code == 200
        with open(os.path.join(EXPORT_PATH, filename), "wb") as f:
            f.write(response.content)

    for name in ["fid", "spectrum"]:
        file_path = os.path.join(EXPORT_PATH, f"{session_id}_{name}.png")
        assert os.path.exists(file_path)
        print(f"✓ Exported: {file_path}")


def test_plot_raw_from_json():
    os.makedirs(EXPORT_PATH, exist_ok=True)

    with open("test.json", "rb") as f:
        r = httpx.post(f"{BASE_URL}/upload/", files={"file": f})
    assert r.status_code == 200
    session_id = r.json()["session_id"]

    r = httpx.post(f"{BASE_URL}/plot-raw/", data={"session_id": session_id})
    assert r.status_code == 200
    plot_url = r.json()["raw_plot"]
    filename = plot_url.split("/")[-1]

    full_url = f"{BASE_URL}{plot_url}"
    response = httpx.get(full_url)
    assert response.status_code == 200

    path = os.path.join(EXPORT_PATH, filename)
    with open(path, "wb") as f:
        f.write(response.content)
    assert os.path.exists(path)
    print(f"✓ Raw plot saved: {path}")


def test_plot_raw_from_json():
    os.makedirs(EXPORT_PATH, exist_ok=True)

    with open("test.json", "rb") as f:
        r = httpx.post(f"{BASE_URL}/upload/", files={"file": f})
    assert r.status_code == 200, f"Upload failed: {r.text}"
    session_id = r.json()["session_id"]

    r = httpx.post(f"{BASE_URL}/plot-raw/", data={"session_id": session_id})
    assert r.status_code == 200, f"Plot raw failed: {r.text}"
    raw_plot_url = r.json()["raw_plot"]
    filename = raw_plot_url.split("/")[-1]

    response = httpx.get(f"{BASE_URL}{raw_plot_url}")
    assert response.status_code == 200, f"Download failed: {response.text}"

    file_path = os.path.join(EXPORT_PATH, filename)
    with open(file_path, "wb") as f:
        f.write(response.content)

    assert os.path.exists(file_path), f"File {file_path} not found"
    print(f"✓ Raw plot saved to {file_path}")

def test_filter_and_fft_from_json():
    os.makedirs(EXPORT_PATH, exist_ok=True)

    with open("test2.json", "rb") as f:
        r = httpx.post(f"{BASE_URL}/upload/", files={"file": f})
    assert r.status_code == 200, f"Upload failed: {r.text}"
    session_id = r.json()["session_id"]

    r = httpx.post(f"{BASE_URL}/filter/", data={
        "session_id": session_id,
        "dt": 1.125e-7,
        "center_freq": 2.97e6,
        "lower_freq": 2.92e6,
        "higher_freq": 3.02e6,
        "low_freq": 600e3
    })
    assert r.status_code == 200, f"Filter failed: {r.text}"

    r = httpx.post(f"{BASE_URL}/fft/", data={"session_id": session_id})
    assert r.status_code == 200, f"FFT failed: {r.text}"

    r = httpx.post(f"{BASE_URL}/export/", data={"session_id": session_id})
    assert r.status_code == 200, f"Export failed: {r.text}"
    data = r.json()

    for name, url in data.items():
        filename = url.split("/")[-1]
        full_url = f"{BASE_URL}{url}"
        response = httpx.get(full_url)
        assert response.status_code == 200, f"Download failed for {filename}"
        file_path = os.path.join(EXPORT_PATH, filename)
        with open(file_path, "wb") as f:
            f.write(response.content)
        print(f"✓ Saved {file_path}")

    for name in ["fid", "spectrum"]:
        file_path = os.path.join(EXPORT_PATH, f"{session_id}_{name}.png")
        assert os.path.exists(file_path), f"{file_path} not found"