tests/test_integrate_datasets_tool.py

"""
Tests for the decoupler_integrate_datasets MCP tool (ADR-0001 T9, Mode A).

The tool's value is the routing gate: it pools + runs a batch-aware DE ONLY on a
plan verdict of 'early', and refuses/reroutes otherwise. These tests monkeypatch
the plan, the combined-AnnData builder, and the DE tool so the routing is checked
deterministically with no network or compute. The actual pooling/DE mechanics
are covered by test_integration_mode_a.py and test_de_batch_covariate.py.
"""

from __future__ import annotations

import asyncio
import sys
from pathlib import Path

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

import src.tools.integration_tools as it  # noqa: E402
from src.tools.integration_tools import (  # noqa: E402
    decoupler_integrate_datasets,
    integration_mcp,
)


def _plan(mode, **extra):
    base = {"mode": mode, "reason": f"{mode} reason.", "refusal_rules_triggered": []}
    base.update(extra)
    return lambda *a, **k: base


def test_tool_registered():
    names = {t.name for t in asyncio.run(integration_mcp.list_tools())}
    assert "decoupler_integrate_datasets" in names


def test_arity_guard():
    out = decoupler_integrate_datasets(
        dataset_ids=["solo"], design_factor="g", test_group="t", control_group="c"
    )
    assert out["refused"] is True
    assert ">=2" in out["error"]


def test_late_reroutes_to_meta_analyze(monkeypatch):
    monkeypatch.setattr(it, "get_integration_plan", _plan("late"))
    out = decoupler_integrate_datasets(
        dataset_ids=["a", "b"], design_factor="g", test_group="t", control_group="c"
    )
    assert out["refused"] and out["mode"] == "late"
    assert "decoupler_meta_analyze" in out["error"]


def test_concordance_reroutes(monkeypatch):
    monkeypatch.setattr(it, "get_integration_plan", _plan("concordance"))
    out = decoupler_integrate_datasets(
        dataset_ids=["a", "b"], design_factor="g", test_group="t", control_group="c"
    )
    assert out["refused"] and out["mode"] == "concordance"
    assert "decoupler_normalization_concordance" in out["error"]


def test_refuse_passes_reason_and_rules(monkeypatch):
    monkeypatch.setattr(
        it,
        "get_integration_plan",
        _plan(
            "refuse", reason="CONFOUNDED_DESIGN: ...", refusal_rules_triggered=["CONFOUNDED_DESIGN"]
        ),
    )
    out = decoupler_integrate_datasets(
        dataset_ids=["a", "b"], design_factor="g", test_group="t", control_group="c"
    )
    assert out["refused"] and out["mode"] == "refuse"
    assert "CONFOUNDED_DESIGN" in out["error"]
    assert out["refusal_rules_triggered"] == ["CONFOUNDED_DESIGN"]


def test_plan_exception_refuses(monkeypatch):
    def boom(*a, **k):
        raise ValueError("dataset 'nope' not registered")

    monkeypatch.setattr(it, "get_integration_plan", boom)
    out = decoupler_integrate_datasets(
        dataset_ids=["a", "nope"], design_factor="g", test_group="t", control_group="c"
    )
    assert out["refused"] and "not registered" in out["error"]


def test_early_builds_and_runs_batch_aware_de(monkeypatch, tmp_path):
    monkeypatch.setattr(
        it, "get_integration_plan", _plan("early", poolable_data_level="raw_counts")
    )
    built = {
        "output_path": str(tmp_path / "combined.h5ad"),
        "n_obs": 42,
        "n_vars": 1500,
        "per_batch_n": {"a": 20, "b": 22},
        "dataset_ids": ["a", "b"],
    }
    monkeypatch.setattr(it, "build_combined_anndata", lambda *a, **k: built)
    captured = {}

    def fake_de(**kwargs):
        captured.update(kwargs)
        return {
            "method_used": "DESeq2",
            "n_significant": 7,
            "output_path": str(tmp_path / "de.csv"),
            "sanity_warnings": {"n_warnings": 0},
        }

    monkeypatch.setattr(it, "decoupler_differential_expression", fake_de)
    out = decoupler_integrate_datasets(
        dataset_ids=["a", "b"],
        design_factor="subtype",
        test_group="basal",
        control_group="classical",
    )
    assert out["mode"] == "early" and out["batch_modeled"] is True
    assert out["n_significant"] == 7 and out["method_used"] == "DESeq2"
    assert out["per_batch_n"] == {"a": 20, "b": 22}
    # auto -> deseq2 for raw counts; batch covariate + contrast wired correctly
    assert captured["method"] == "deseq2"
    assert captured["batch_column"] == "batch"
    assert captured["contrast"] == ["subtype", "basal", "classical"]
    assert captured["adata_path"] == built["output_path"]


def test_early_auto_method_limma_for_non_raw(monkeypatch, tmp_path):
    monkeypatch.setattr(
        it, "get_integration_plan", _plan("early", poolable_data_level="log_expression")
    )
    monkeypatch.setattr(
        it,
        "build_combined_anndata",
        lambda *a, **k: {
            "output_path": str(tmp_path / "c.h5ad"),
            "n_obs": 10,
            "n_vars": 900,
            "per_batch_n": {"a": 5, "b": 5},
        },
    )
    captured = {}

    def fake_de(**kwargs):
        captured.update(kwargs)
        return {"method_used": "limma", "n_significant": 3, "output_path": str(tmp_path / "de.csv")}

    monkeypatch.setattr(it, "decoupler_differential_expression", fake_de)
    decoupler_integrate_datasets(
        dataset_ids=["a", "b"], design_factor="g", test_group="t", control_group="c"
    )
    assert captured["method"] == "limma"


def test_early_build_failure_refuses(monkeypatch):
    monkeypatch.setattr(
        it, "get_integration_plan", _plan("early", poolable_data_level="raw_counts")
    )

    def boom(*a, **k):
        raise ValueError("only 12 shared gene symbols")

    monkeypatch.setattr(it, "build_combined_anndata", boom)
    out = decoupler_integrate_datasets(
        dataset_ids=["a", "b"], design_factor="g", test_group="t", control_group="c"
    )
    assert out["refused"] and "shared gene symbols" in out["error"]