tests/test_integration_plan_tool.py

"""
Tests for the dataset_get_integration_plan MCP tool (ADR-0001 T4).

The tool is a thin agent-side wrapper over biodata_registry.get_integration_plan,
re-exposed through src/datasets/registry.py. These tests confirm it is registered
on dataset_mcp, delegates unchanged to the registry, returns the full plan
contract, forwards the optional contrast args to the confound gate, and degrades
to an error dict on unknown ids. Pure metadata — no network, no data loading.
"""

from __future__ import annotations

import asyncio
import sys
from pathlib import Path

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

import biodata_registry as bdr  # noqa: E402

from src.datasets.registry import get_integration_plan  # noqa: E402
from src.tools import dataset_tools  # noqa: E402
from src.tools.dataset_tools import dataset_get_integration_plan  # noqa: E402

PLAN_KEYS = {
    "mode",
    "reason",
    "shared_feature_space",
    "requires_ortholog_mapping",
    "requires_probe_collapse",
    "batch_key",
    "poolable_data_level",
    "per_dataset",
    "refusal_rules_triggered",
}


def test_tool_registered_on_dataset_mcp():
    """The wrapper is discoverable as an MCP tool, like the other dataset tools."""
    tools = asyncio.run(dataset_tools.dataset_mcp.list_tools())
    names = {t.name for t in tools}
    assert "dataset_get_integration_plan" in names


def test_known_raw_counts_pair_is_early():
    """Two raw_counts RNA-seq cohorts pool early; the full contract is returned."""
    plan = dataset_get_integration_plan(["paca_au_rnaseq", "tcga_paad"])
    assert set(plan) >= PLAN_KEYS
    assert plan["mode"] == "early"
    assert plan["poolable_data_level"] == "raw_counts"
    assert plan["batch_key"] == "dataset_id"
    assert plan["reason"]  # non-empty; the agent surfaces it verbatim
    assert plan["refusal_rules_triggered"] == []


def test_mixed_levels_fall_back_to_late():
    """Microarray (log_expression) + RNA-seq (raw_counts) must NOT pool -> late."""
    plan = dataset_get_integration_plan(["gse71729_moffitt", "tcga_paad"])
    assert plan["mode"] == "late"
    assert plan["reason"]


def test_single_dataset_refuses_not_multi():
    plan = dataset_get_integration_plan(["tcga_paad"])
    assert plan["mode"] == "refuse"
    assert "NOT_MULTI" in plan["refusal_rules_triggered"]


def test_unknown_dataset_returns_error_dict():
    """Unknown ids degrade to an error dict (mirrors dataset_describe), not a raise."""
    plan = dataset_get_integration_plan(["tcga_paad", "nope_not_real"])
    assert "error" in plan
    assert "available_datasets" in plan
    assert "nope_not_real" in plan["error"]


def test_contrast_args_reach_confound_gate():
    """A contrast no cohort can express forwards to Gate 6 -> CONFOUNDED_DESIGN."""
    plan = dataset_get_integration_plan(
        ["paca_au_rnaseq", "tcga_paad"],
        design_factor="tumor_subtype",
        test_group="a",
        control_group="b",
    )
    assert plan["mode"] == "refuse"
    assert "CONFOUNDED_DESIGN" in plan["refusal_rules_triggered"]


def test_registry_passthrough_is_thin():
    """src.datasets.registry.get_integration_plan delegates unchanged to the package."""
    pair = ["paca_au_rnaseq", "tcga_paad"]
    assert get_integration_plan(pair) == bdr.get_integration_plan(pair)