PlantAI-dev/tests/unit/test_rnaseq_context.py

from pathlib import Path
import sys

import pytest

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

from gfa_core.rnaseq_context import (
    RnaseqLibraryStrandnessError,
    apply_rnaseq_context,
    build_rnaseq_context,
    featurecounts_strandness_value,
    find_rnaseq_hisat2_index,
    find_rnaseq_reference_file,
    hisat2_rna_strandness_args,
    require_rnaseq_library_strandness,
    resolve_rnaseq_library_strandness,
    resolve_rnaseq_project_context,
)


def test_rnaseq_context_prefers_species_output_over_generic_or_reference_prefix(tmp_path):
    input_dir = tmp_path / "input"
    input_dir.mkdir()
    (input_dir / "samples.yaml").write_text(
        """species:
  code: Os
  speciesName: Oryza sativa
reference:
  id: osa1_r7
rnaseq:
  outputPrefix: Os
samples:
  S1:
    r1: S1_1.fastq
    r2: S1_2.fastq
    condition: T0
""",
        encoding="utf-8",
    )
    (input_dir / "osa1_r7.gene_models.repr.cds.1.ht2").write_text("cds index", encoding="utf-8")
    species_dir = input_dir / "Os"
    species_dir.mkdir()
    genome = species_dir / "Os.genome.fasta"
    genome.write_text(">chr1\nACGT\n", encoding="utf-8")
    (species_dir / "Os.genome.1.ht2").write_text("genome index", encoding="utf-8")

    context = resolve_rnaseq_project_context(
        input_dir=input_dir,
        requested_prefix="genome_prepare",
        sample_config=input_dir / "samples.yaml",
    )

    assert context["output_prefix"] == "Os"
    assert context["species_code"] == "Os"
    assert context["reference_id"] == "osa1_r7"
    assert find_rnaseq_reference_file(input_dir, "genome", "Os", "osa1_r7") == genome
    assert find_rnaseq_hisat2_index(input_dir, "Os", "osa1_r7") == species_dir / "Os.genome"


def test_rnaseq_reference_lookup_rejects_cds_as_genome(tmp_path):
    input_dir = tmp_path / "input"
    input_dir.mkdir()
    (input_dir / "osa1_r7.gene_models.repr.cds.fa").write_text(">cds\nACGT\n", encoding="utf-8")
    (input_dir / "osa1_r7.gene_models.repr.cds.1.ht2").write_text("cds index", encoding="utf-8")

    assert find_rnaseq_reference_file(input_dir, "genome", "Os", "osa1_r7") is None
    assert find_rnaseq_hisat2_index(input_dir, "Os", "osa1_r7") is None


def test_rnaseq_library_strandness_project_level_is_inherited(tmp_path):
    sample_config = tmp_path / "samples.yaml"
    sample_config.write_text(
        """rnaseq:
  library_strandness: reverse
samples:
  S1: {r1: S1_R1.fq.gz, r2: S1_R2.fq.gz}
  S2: {r1: S2_R1.fq.gz, r2: S2_R2.fq.gz}
""",
        encoding="utf-8",
    )

    info = require_rnaseq_library_strandness(sample_config, ["S1", "S2"])

    assert info["project_library_strandness"] == "reverse"
    assert info["library_strandness_by_sample"] == {"S1": "reverse", "S2": "reverse"}
    assert info["library_strandness_values"] == ["reverse"]
    assert hisat2_rna_strandness_args("reverse") == ["--rna-strandness", "RF"]
    assert featurecounts_strandness_value("reverse") == "2"


def test_rnaseq_library_strandness_sample_level_overrides_project_level(tmp_path):
    sample_config = tmp_path / "samples.yaml"
    sample_config.write_text(
        """rnaseq:
  library_strandness: forward
samples:
  S1: {r1: S1_R1.fq.gz, r2: S1_R2.fq.gz}
  S2: {r1: S2_R1.fq.gz, r2: S2_R2.fq.gz, library_strandness: non_stranded}
""",
        encoding="utf-8",
    )

    info = require_rnaseq_library_strandness(sample_config, ["S1", "S2"])

    assert info["library_strandness_by_sample"] == {"S1": "forward", "S2": "non_stranded"}
    assert info["library_strandness_values"] == ["forward", "non_stranded"]
    assert info["library_strandness_mixed"] is True
    assert hisat2_rna_strandness_args("forward") == ["--rna-strandness", "FR"]
    assert featurecounts_strandness_value("non_stranded") == "0"


def test_rnaseq_library_strandness_missing_is_required_for_align_quant(tmp_path):
    sample_config = tmp_path / "samples.yaml"
    sample_config.write_text(
        """samples:
  S1: {r1: S1_R1.fq.gz, r2: S1_R2.fq.gz}
""",
        encoding="utf-8",
    )

    with pytest.raises(RnaseqLibraryStrandnessError, match="缺少 library_strandness"):
        require_rnaseq_library_strandness(sample_config, ["S1"])


def test_rnaseq_library_strandness_invalid_value_is_rejected(tmp_path):
    sample_config = tmp_path / "samples.yaml"
    sample_config.write_text(
        """samples:
  S1: {r1: S1_R1.fq.gz, r2: S1_R2.fq.gz, library_strandness: unknown}
""",
        encoding="utf-8",
    )

    with pytest.raises(RnaseqLibraryStrandnessError, match="允许值"):
        require_rnaseq_library_strandness(sample_config, ["S1"])


def test_rnaseq_library_strandness_context_is_written_to_manifest(tmp_path):
    sample_config = tmp_path / "samples.yaml"
    count_matrix = tmp_path / "count_matrix.txt"
    sample_config.write_text(
        """species:
  code: Os
rnaseq:
  outputPrefix: Os
samples:
  S1: {condition: CK, library_strandness: reverse}
  S2: {condition: salt, library_strandness: reverse}
groups:
  S1: CK
  S2: salt
""",
        encoding="utf-8",
    )
    count_matrix.write_text("gene\tS1\tS2\nGene1\t10\t20\n", encoding="utf-8")

    context = build_rnaseq_context(
        prefix="Os",
        sample_config=sample_config,
        count_matrix=count_matrix,
        stage="rnaseq_quant",
    )
    manifest = apply_rnaseq_context({"parameters": {}}, context)

    assert manifest["rnaseq_library_strandness_by_sample"] == {"S1": "reverse", "S2": "reverse"}
    assert manifest["rnaseq_library_strandness_values"] == ["reverse"]
    assert manifest["parameters"]["rnaseq_library_strandness_signature"]
    assert resolve_rnaseq_library_strandness(sample_config, ["S1", "S2"])["library_strandness_signature"]