csaxs_bec/omny_e2e_tests/test_tomo_queue.py

"""Checklist section A -- sim, no crash needed. See TOMO_QUEUE_TESTING.md."""

import datetime

from _queue_helpers import ProgressSampler
from _queue_helpers import add_short_job as _add_short_job
from _queue_helpers import short_params as _short_params


def test_params_restored_per_job(flomni_sim):
    """A1: each job's snapshotted params, not just whatever the global var
    last held, actually reach the running scan.

    job1 and job2 differ only in tomo_angle_stepsize. sub_tomo_scan() derives
    progress["subtomo_total_projections"] = int(tomo_angle_range /
    tomo_angle_stepsize) live, at scan time, from the *live* global var --
    so if the queue's per-job setattr() restore didn't really work, both
    jobs would show whatever stepsize happened to be live last (i.e. job2's,
    since it was set most recently), not each job's own value.
    """
    flomni = flomni_sim

    job1_stepsize = 89.9
    job2_stepsize = 59.9
    expected_n = {"job1": int(180 / job1_stepsize), "job2": int(180 / job2_stepsize)}
    assert expected_n["job1"] != expected_n["job2"], "test setup must pick distinguishable N"

    _add_short_job(flomni, "job1", job1_stepsize)
    _add_short_job(flomni, "job2", job2_stepsize)

    # Reset first: leftover progress state from a previous run could
    # otherwise coincidentally satisfy the "value observed" check below
    # before job1 has actually started.
    flomni.progress.reset()

    with ProgressSampler(flomni) as sampler:
        flomni.tomo_queue_execute()

    observed = sampler.observations
    job1_idx = sampler.first_index_of(expected_n["job1"])
    job2_idx = sampler.first_index_of(expected_n["job2"])

    assert job1_idx is not None, (
        f"job1 (stepsize={job1_stepsize}) should have shown "
        f"subtomo_total_projections={expected_n['job1']} at some point; observed {observed}"
    )
    assert job2_idx is not None, (
        f"job2 (stepsize={job2_stepsize}) should have shown "
        f"subtomo_total_projections={expected_n['job2']} at some point; observed {observed}"
    )
    assert job1_idx < job2_idx, (
        "job1's N should be observed before job2's -- jobs ran out of order "
        f"(job1 first at {job1_idx}, job2 first at {job2_idx})"
    )

    jobs = flomni._tomo_queue_proxy.as_list()
    assert [j["status"] for j in jobs] == ["done", "done"]


def test_empty_and_all_done_queue_are_noops(flomni_sim, monkeypatch, capsys):
    """A3: tomo_queue_execute() on an empty queue, and on a queue whose jobs
    are all "done", must print a no-op message, never prompt, and never
    trigger a scan.
    """
    flomni = flomni_sim

    prompt_calls = []
    from csaxs_bec.bec_ipython_client.plugins.OMNY_shared.omny_general_tools import OMNYTools

    monkeypatch.setattr(
        OMNYTools, "yesno", lambda self, *a, **k: prompt_calls.append((a, k)) or True
    )

    # -- empty queue --
    flomni._tomo_queue_proxy.clear()
    capsys.readouterr()  # discard anything buffered so far
    flomni.tomo_queue_execute()
    out = capsys.readouterr().out
    assert "Tomo queue is empty." in out
    assert prompt_calls == []

    # -- all-done queue --
    _add_short_job(flomni, "already_done")
    job_id = flomni._tomo_queue_proxy.as_list()[0]["id"]
    flomni._tomo_queue_proxy.update_by_id(job_id, status="done")
    capsys.readouterr()
    flomni.tomo_queue_execute()
    out = capsys.readouterr().out
    assert "No pending tomo queue jobs to run." in out
    assert prompt_calls == []

    jobs = flomni._tomo_queue_proxy.as_list()
    assert len(jobs) == 1 and jobs[0]["status"] == "done"  # untouched by the no-op call


def test_start_index_semantics(flomni_sim):
    """A4: start_index is "position below which fresh pending jobs are
    ignored" -- it does NOT apply to resumable (incomplete/running) jobs,
    which always run regardless of position.
    """
    flomni = flomni_sim

    _add_short_job(flomni, "job0")
    _add_short_job(flomni, "job1")
    jobs = flomni._tomo_queue_proxy.as_list()
    job0_id, job1_id = jobs[0]["id"], jobs[1]["id"]

    flomni.tomo_queue_execute(start_index=1)

    jobs = {j["id"]: j for j in flomni._tomo_queue_proxy.as_list()}
    assert jobs[job0_id]["status"] == "pending", "job0 should have been skipped by start_index"
    assert jobs[job1_id]["status"] == "done", "job1 should have run"

    # Mark job0 "incomplete" (as if it had crashed) and drop progress so its
    # resume is a clean no-op ("No tomo scan in progress to resume"). This
    # test is about *selection* (does start_index wrongly skip a resumable
    # job?), not resume accuracy -- that's section B.
    flomni._tomo_queue_proxy.update_by_id(job0_id, status="incomplete")
    flomni.progress.reset()

    flomni.tomo_queue_execute(start_index=1)

    jobs = {j["id"]: j for j in flomni._tomo_queue_proxy.as_list()}
    assert (
        jobs[job0_id]["status"] == "done"
    ), "job0 is resumable (incomplete) and must run despite start_index=1"
    assert jobs[job1_id]["status"] == "done", "job1 (already done) should not have been re-run"


def test_legacy_queue_migration(flomni_sim):
    """A2: a queue written before jobs carried an ``id`` field (e.g. by an
    older client version) must run cleanly, and ensure_ids() must heal every
    job with a unique id as a side effect of tomo_queue_execute().
    """
    flomni = flomni_sim
    bec = flomni.client

    legacy_jobs = []
    for label in ("legacy1", "legacy2"):
        params = _short_params(flomni)
        legacy_jobs.append(
            {
                # deliberately no "id" key -- pre-Step-1 job shape
                "label": label,
                "params": params,
                "status": "pending",
                "added_at": datetime.datetime.now().isoformat(),
            }
        )
    bec.set_global_var("tomo_queue", legacy_jobs)

    flomni.tomo_queue_show()
    flomni.tomo_queue_execute()

    jobs = flomni._tomo_queue_proxy.as_list()
    assert [j["status"] for j in jobs] == ["done", "done"]
    ids = [j.get("id") for j in jobs]
    assert all(ids), f"every job should have been healed with an id, got {ids}"
    assert len(set(ids)) == len(ids), f"healed ids should be unique, got {ids}"


def test_queue_show_survives_missing_param_key(flomni_sim):
    """A job snapshotted before some key (e.g. fovx) existed in
    _TOMO_SCAN_PARAM_NAMES won't carry it in its persisted ``params`` dict.
    tomo_queue_show() must still print a summary line for that job instead
    of raising KeyError -- otherwise a single stale queue entry makes the
    whole queue un-inspectable (and un-deletable via the index tomo_queue_show()
    would have printed).
    """
    flomni = flomni_sim
    bec = flomni.client

    params = _short_params(flomni)
    del params["fovx"]
    del params["fovy"]
    legacy_job = {
        "kind": "tomo",
        "id": "legacy-missing-fov",
        "label": "legacy_no_fov",
        "params": params,
        "status": "pending",
        "added_at": datetime.datetime.now().isoformat(),
    }
    bec.set_global_var("tomo_queue", [legacy_job])

    jobs = flomni.tomo_queue_show()
    assert jobs[0]["label"] == "legacy_no_fov"