PlantAI-dev/tests/unit/test_string_species_snapshot.py

from pathlib import Path
import json

import pytest


def _write_species(path: Path, taxa: list[tuple[str, str]]) -> None:
    rows = [
        "#taxon_id\tSTRING_type\tSTRING_name_compact\tofficial_name_NCBI\tdomain"
    ]
    rows.extend(
        f"{taxon_id}\tcore\t{name}\t{name}\tEukaryota"
        for taxon_id, name in taxa
    )
    path.write_text("\n".join(rows) + "\n", encoding="utf-8")


def _write_tree(path: Path, rows: list[tuple[str, str, str, str]]) -> None:
    content = ["#taxon_id\tparent_taxon_id\ttaxon_name\tis_STRING_species"]
    content.extend("\t".join(row) for row in rows)
    path.write_text("\n".join(content) + "\n", encoding="utf-8")


def test_snapshot_rejects_missing_parent_and_missing_required_nodes(tmp_path: Path) -> None:
    from scripts.prepare_string_species_snapshot import validate_snapshot

    species = tmp_path / "species.txt"
    tree = tmp_path / "tree.txt"
    _write_species(species, [("3702", "Arabidopsis thaliana")])
    _write_tree(
        tree,
        [
            ("1", "1", "root", "f"),
            ("3702", "999999", "Arabidopsis thaliana", "t"),
        ],
    )

    with pytest.raises(ValueError, match="parent|33090|3702"):
        validate_snapshot(species, tree)


def test_snapshot_accepts_official_root_self_loop_and_strict_tf(tmp_path: Path) -> None:
    from scripts.prepare_string_species_snapshot import validate_snapshot

    species = tmp_path / "species.txt"
    tree = tmp_path / "tree.txt"
    _write_species(species, [("3702", "Arabidopsis thaliana")])
    _write_tree(
        tree,
        [
            ("1", "1", "root", "f"),
            ("33090", "1", "Viridiplantae", "f"),
            ("3702", "33090", "Arabidopsis thaliana", "t"),
        ],
    )

    result = validate_snapshot(species, tree)

    assert result["viridiplantae_taxon_id"] == "33090"
    assert result["arabidopsis_taxon_id"] == "3702"
    assert result["tree_node_count"] == 3
    assert result["string_species_count"] == 1


def test_snapshot_rejects_non_root_cycle(tmp_path: Path) -> None:
    from scripts.prepare_string_species_snapshot import validate_snapshot

    species = tmp_path / "species.txt"
    tree = tmp_path / "tree.txt"
    _write_species(species, [("3702", "Arabidopsis thaliana")])
    _write_tree(
        tree,
        [
            ("1", "1", "root", "f"),
            ("33090", "3702", "Viridiplantae", "f"),
            ("3702", "33090", "Arabidopsis thaliana", "t"),
        ],
    )

    with pytest.raises(ValueError, match="cycle"):
        validate_snapshot(species, tree)


def test_snapshot_rejects_registry_tree_string_species_mismatch(tmp_path: Path) -> None:
    from scripts.prepare_string_species_snapshot import validate_snapshot

    species = tmp_path / "species.txt"
    tree = tmp_path / "tree.txt"
    _write_species(
        species,
        [("3702", "Arabidopsis thaliana"), ("3708", "Brassica napus")],
    )
    _write_tree(
        tree,
        [
            ("1", "1", "root", "f"),
            ("33090", "1", "Viridiplantae", "f"),
            ("3702", "33090", "Arabidopsis thaliana", "t"),
            ("3708", "33090", "Brassica napus", "f"),
        ],
    )

    with pytest.raises(ValueError, match="registry|STRING species"):
        validate_snapshot(species, tree)


def test_prepare_snapshot_writes_validated_local_files_and_manifest(
    tmp_path: Path,
) -> None:
    import scripts.prepare_string_species_snapshot as snapshot_module

    sources = tmp_path / "sources"
    sources.mkdir()
    species = sources / "species.v12.0.txt"
    tree = sources / "species.tree.v12.0.txt"
    _write_species(species, [("3702", "Arabidopsis thaliana")])
    _write_tree(
        tree,
        [
            ("1", "1", "root", "f"),
            ("33090", "1", "Viridiplantae", "f"),
            ("3702", "33090", "Arabidopsis thaliana", "t"),
        ],
    )
    destination = tmp_path / "snapshot"
    manifest = snapshot_module.prepare_snapshot(
        destination,
        source_directory=sources,
    )

    saved = json.loads((destination / "source_manifest.json").read_text(encoding="utf-8"))
    assert saved == manifest
    assert saved["string_version"] == "12.0"
    assert saved["transfer_mode"] == "local_files"
    assert saved["snapshot_id"]
    assert saved["validation"]["string_species_count"] == 1
    for name in snapshot_module.SOURCE_URLS:
        assert (destination / name).is_file()
        assert saved["files"][name]["size"] > 0
        assert len(saved["files"][name]["sha256"]) == 64


@pytest.mark.parametrize("value", ["false", "0", "yes", "T"])
def test_snapshot_rejects_non_official_boolean_tokens(tmp_path: Path, value: str) -> None:
    from scripts.prepare_string_species_snapshot import validate_snapshot

    species = tmp_path / "species.txt"
    tree = tmp_path / "tree.txt"
    _write_species(species, [("3702", "Arabidopsis thaliana")])
    _write_tree(
        tree,
        [
            ("1", "1", "root", "f"),
            ("33090", "1", "Viridiplantae", "f"),
            ("3702", "33090", "Arabidopsis thaliana", value),
        ],
    )

    with pytest.raises(ValueError, match="is_STRING_species"):
        validate_snapshot(species, tree)