tests/test_dataset_load_tool.py

"""
Tests for dataset_load — the registry-aware load-by-dataset_id executor.

The 2026-08-04 step-budget measurement showed runs burning 3-5 steps picking a
loading entry point (and 2 more collapsing probes when the uncollapsed URL was
used despite a precomputed collapsed_url). dataset_load executes the manifest's
loading plan in ONE tool call: correct loader for the source type, collapsed
URL when present, clinical join, curated-sample filter — and returns the
analysis-ready h5ad path plus routing facts.

Loader executors are monkeypatched (no network); the plans are built from the
real registry manifests so the wiring under test is the deployed wiring.
"""

from __future__ import annotations

import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))

import src.tools.dataset_tools.catalog as catalog  # noqa: E402
from src.datasets.registry import load_manifest  # noqa: E402
from src.tools.dataset_tools._base import _LOADING_PLAN_TOOLS, _build_loading_plan  # noqa: E402


def _fake_executors(calls, fail_on=None):
    """Executor map whose loaders record calls and return sequential paths."""

    def make(tool):
        def run(**kwargs):
            if tool == fail_on:
                raise RuntimeError("boom")
            calls.append((tool, kwargs))
            return {"output_path": f"/tmp/fake/{tool}_{len(calls)}.h5ad"}

        return run

    return {t: make(t) for t in _LOADING_PLAN_TOOLS}


def test_unknown_dataset_returns_available_list():
    res = catalog.dataset_load("no_such_dataset")
    assert "error" in res
    assert "gse71729_moffitt" in res["available_datasets"]


def test_executes_plan_in_order_and_threads_paths(monkeypatch):
    calls = []
    monkeypatch.setattr(catalog, "_loading_step_executors", lambda: _fake_executors(calls))
    res = catalog.dataset_load("gse71729_moffitt")

    assert "error" not in res
    plan_tools = [
        s["tool"]
        for s in _build_loading_plan(load_manifest("gse71729_moffitt"))
        if s["tool"] in _LOADING_PLAN_TOOLS
    ]
    assert [t for t, _ in calls] == plan_tools
    assert [s["tool"] for s in res["steps_executed"]] == plan_tools
    # The final adata_path is the last step's output.
    assert res["adata_path"] == f"/tmp/fake/{plan_tools[-1]}_{len(calls)}.h5ad"
    # No placeholder strings may reach a loader.
    for _, kwargs in calls:
        for v in kwargs.values():
            assert not (isinstance(v, str) and v.startswith("<")), (
                f"placeholder leaked into loader args: {v}"
            )
    # Routing facts come from the manifest.
    assert res["analysis_path"] in ("A", "B", "P")
    assert res["data_level"]
    assert "design_factor" in res["default_contrast"]
    assert "decoupler_inspect_data" in res["next_step"]  # explicitly discouraged


def test_curated_dataset_includes_curation_step(monkeypatch):
    """tcga_paad has a curated_sample_list — the executor must run the filter."""
    calls = []
    monkeypatch.setattr(catalog, "_loading_step_executors", lambda: _fake_executors(calls))
    res = catalog.dataset_load("tcga_paad")
    tools = [t for t, _ in calls]
    assert "dataset_filter_to_curated_samples" in tools
    assert res["analysis_path"] == "A"
    assert "deseq2" in res["recommended_de_method"]


def test_path_b_dataset_recommends_non_deseq2(monkeypatch):
    calls = []
    monkeypatch.setattr(catalog, "_loading_step_executors", lambda: _fake_executors(calls))
    res = catalog.dataset_load("gse71729_moffitt")
    assert res["analysis_path"] == "B"
    assert "deseq2" not in res["recommended_de_method"].split(" ")[0]
    assert "limma" in res["recommended_de_method"]


def test_step_failure_returns_partial_progress(monkeypatch):
    calls = []
    manifest = load_manifest("tcga_paad")
    plan_tools = [
        s["tool"] for s in _build_loading_plan(manifest) if s["tool"] in _LOADING_PLAN_TOOLS
    ]
    fail_tool = plan_tools[-1]  # fail the last loading step
    monkeypatch.setattr(
        catalog, "_loading_step_executors", lambda: _fake_executors(calls, fail_on=fail_tool)
    )
    res = catalog.dataset_load("tcga_paad")
    assert "error" in res and fail_tool in res["error"]
    # earlier steps are reported, and the last good intermediate is surfaced
    assert [s["tool"] for s in res["steps_executed"]] == plan_tools[:-1]
    assert res["adata_path"] == f"/tmp/fake/{plan_tools[-2]}_{len(calls)}.h5ad"


def test_error_dict_from_loader_is_surfaced(monkeypatch):
    def bad_executors():
        ex = _fake_executors([])
        first = [
            s["tool"]
            for s in _build_loading_plan(load_manifest("gse71729_moffitt"))
            if s["tool"] in _LOADING_PLAN_TOOLS
        ][0]
        ex[first] = lambda **kw: {"error": "404 not found"}
        return ex

    monkeypatch.setattr(catalog, "_loading_step_executors", bad_executors)
    res = catalog.dataset_load("gse71729_moffitt")
    assert "error" in res and "404 not found" in res["error"]


def test_single_cell_dataset_flags_pseudobulk(monkeypatch):
    calls = []
    monkeypatch.setattr(catalog, "_loading_step_executors", lambda: _fake_executors(calls))
    res = catalog.dataset_load("gse155698_steele")
    assert res["analysis_path"] == "P"
    assert "pseudobulk" in res["recommended_de_method"].lower()
    assert "pseudobulk" in res["next_step"].lower()