PlantAI-dev/tests/unit/test_rnaseq_quant_expression_policy.py

import sys
from pathlib import Path


ROOT = Path(__file__).resolve().parents[2]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))


def _write_matrix(path: Path, header, rows):
    path.write_text(
        "\n".join(
            ["\t".join(header)]
            + ["\t".join([str(row[0]), *[str(v) for v in row[1:]]]) for row in rows]
        )
        + "\n",
        encoding="utf-8",
    )


def test_featurecounts_paired_end_counts_fragments(monkeypatch, tmp_path):
    from apps.cli import run_rnaseq_quant as quant

    captured = {}

    def fake_run(cmd, check, capture_output, text):
        captured["cmd"] = list(cmd)
        return object()

    monkeypatch.setattr(quant.subprocess, "run", fake_run)
    annotation = tmp_path / "genes.gtf"
    annotation.write_text("placeholder\n", encoding="utf-8")

    result = quant.run_featurecounts(
        [tmp_path / "S1.bam"],
        annotation,
        tmp_path / "featurecounts.txt",
        cpu=2,
        library_strandness="non_stranded",
        command_log=tmp_path / "featurecounts.command.txt",
    )

    assert result == annotation
    assert "-p" in captured["cmd"]
    assert "--countReadPairs" in captured["cmd"]
    assert captured["cmd"].index("--countReadPairs") > captured["cmd"].index("-p")


def test_tx2gene_parser_supports_gff3_and_gtf_attributes(tmp_path):
    from apps.cli.run_rnaseq_quant import parse_tx2gene_from_gff

    annotation = tmp_path / "mixed.gff3"
    annotation.write_text(
        "\n".join(
            [
                "chr1\tsrc\tmRNA\t1\t100\t.\t+\t.\tID=tx.gff3;Parent=gene.gff3",
                'chr1\tsrc\ttranscript\t200\t300\t.\t+\t.\tgene_id "gene.gtf"; transcript_id "tx.gtf";',
                "chr1\tsrc\ttranscript\t400\t500\t.\t+\t.\tID=tx.name;Name=gene.name",
                "chr1\tsrc\tCDS\t600\t700\t.\t+\t0\tID=cds1;Parent=tx.parent;gene_id=gene.cds",
            ]
        )
        + "\n",
        encoding="utf-8",
    )

    tx2gene = parse_tx2gene_from_gff(annotation)

    assert tx2gene["tx.gff3"] == "gene.gff3"
    assert tx2gene["tx.gtf"] == "gene.gtf"
    assert tx2gene["tx.name"] == "gene.name"
    assert tx2gene["tx.parent"] == "gene.cds"


def test_matrix_bridge_qc_accepts_salmon_when_samples_and_genes_match(tmp_path):
    from apps.cli.run_rnaseq_quant import compute_matrix_bridge_qc

    count_matrix = tmp_path / "Os.count_matrix.txt"
    salmon_tpm = tmp_path / "Os.salmon.tpm_matrix.txt"
    tx2gene = tmp_path / "Os.salmon.tx2gene.tsv"
    index_dir = tmp_path / "salmon_index"
    index_dir.mkdir()

    rows = [(f"gene{i}", i, i + 1) for i in range(1, 21)]
    _write_matrix(count_matrix, ["gene_id", "S1", "S2"], rows)
    _write_matrix(salmon_tpm, ["gene_id", "S1", "S2"], [(f"gene{i}", 50000, 50000) for i in range(1, 21)])
    tx2gene.write_text(
        "transcript_id\tgene_id\n" + "\n".join(f"tx{i}\tgene{i}" for i in range(1, 21)) + "\n",
        encoding="utf-8",
    )
    (index_dir / "ref_indexing.log").write_text(
        "Removed 2 transcripts that were sequence duplicates of indexed transcripts.\n",
        encoding="utf-8",
    )
    (index_dir / "info.json").write_text('{"keep_duplicates": false}\n', encoding="utf-8")

    qc = compute_matrix_bridge_qc(
        count_matrix=count_matrix,
        abundance_tpm=salmon_tpm,
        tx2gene=tx2gene,
        engine="salmon",
        salmon_index_dir=index_dir,
        output_path=tmp_path / "Os.rnaseq_quant.matrix_bridge_qc.tsv",
    )

    assert qc["status"] == "pass"
    assert qc["sample_columns_match"] is True
    assert qc["gene_overlap_vs_featurecounts_pct"] == 100.0
    assert qc["tpm_column_sums_within_tolerance"] is True
    assert qc["duplicate_transcripts_removed"] == 2
    assert qc["recommended_expression_tpm_source"] == "salmon_tximport_tpm"


def test_matrix_bridge_qc_falls_back_when_tx2gene_fallback_is_high(tmp_path):
    from apps.cli.run_rnaseq_quant import compute_matrix_bridge_qc

    count_matrix = tmp_path / "Os.count_matrix.txt"
    salmon_tpm = tmp_path / "Os.salmon.tpm_matrix.txt"
    tx2gene = tmp_path / "Os.salmon.tx2gene.tsv"

    rows = [(f"gene{i}", i, i + 1) for i in range(1, 21)]
    _write_matrix(count_matrix, ["gene_id", "S1", "S2"], rows)
    _write_matrix(salmon_tpm, ["gene_id", "S1", "S2"], [(f"gene{i}", 50000, 50000) for i in range(1, 21)])
    tx2gene.write_text(
        "transcript_id\tgene_id\n" + "\n".join(f"tx{i}\tgene{i}" for i in range(1, 21)) + "\n",
        encoding="utf-8",
    )

    qc = compute_matrix_bridge_qc(
        count_matrix=count_matrix,
        abundance_tpm=salmon_tpm,
        tx2gene=tx2gene,
        engine="salmon",
        tx2gene_metrics={"tx2gene_fallback_pct": 6.0},
    )

    assert qc["status"] == "fallback"
    assert qc["recommended_expression_tpm_source"] == "featurecounts_tpm"
    assert "tx2gene_fallback_gt_5pct" in qc["failure_reasons"]


def test_expression_resolver_prefers_salmon_tpm_when_bridge_qc_passes(tmp_path):
    from apps.cli.run_expression import resolve_rnaseq_expression_matrices

    final_dir = tmp_path / "final"
    final_dir.mkdir()
    for name in (
        "Os.fpkm_matrix.txt",
        "Os.tpm_matrix.txt",
        "Os.salmon.tpm_matrix.txt",
        "Os.rnaseq_sample_qc.vst_matrix.tsv",
    ):
        _write_matrix(final_dir / name, ["gene_id", "S1", "S2"], [("gene1", 1, 2)])
    (final_dir / "Os.rnaseq_quant.matrix_bridge_qc.tsv").write_text(
        "metric\tvalue\n"
        "status\tpass\n"
        "recommended_expression_tpm_source\tsalmon_tximport_tpm\n"
        "gene_overlap_vs_featurecounts_pct\t99.5\n"
        "sample_columns_match\ttrue\n",
        encoding="utf-8",
    )

    resolved = resolve_rnaseq_expression_matrices(
        final_dir=final_dir,
        prefix="Os",
        fpkm_file=final_dir / "Os.fpkm_matrix.txt",
        tpm_file=final_dir / "Os.tpm_matrix.txt",
        vst_file=final_dir / "Os.rnaseq_sample_qc.vst_matrix.tsv",
    )

    assert resolved["tpm_file"] == final_dir / "Os.salmon.tpm_matrix.txt"
    assert resolved["tpm_source"] == "salmon_tximport_tpm"
    assert "bridge_qc_passed" in resolved["tpm_source_reason"]