SciMesh/tests/test_cli_workload.py

"""Tests for the generic ``scimesh workload`` CLI."""

from __future__ import annotations

from pathlib import Path

import pytest

from scimesh.cli import main


def test_workload_cli_lists_sdk_workloads(
    capsys: pytest.CaptureFixture[str],
) -> None:
    assert main(["workload", "list"]) == 0
    output = capsys.readouterr().out
    assert "descriptor-batch" in output
    assert "similarity-graph" in output
    assert "similarity-search" in output
    assert "enabled" in output


def test_workload_cli_runs_descriptor_batch(
    tmp_path: Path, capsys: pytest.CaptureFixture[str]
) -> None:
    dataset = tmp_path / "molecules.tsv"
    dataset.write_text(
        "chembl_id\tcanonical_smiles\nA\tCCO\nB\tCCCC\nC\tCCN\n",
        encoding="utf-8",
    )
    output = tmp_path / "descriptors.csv"

    code = main(
        [
            "workload",
            "run",
            "descriptor-batch",
            "--input",
            str(dataset),
            "--params",
            '{"skip_invalid": true}',
            "--shard-rows",
            "2",
            "-o",
            str(output),
        ]
    )

    assert code == 0
    lines = output.read_text(encoding="utf-8").splitlines()
    assert lines[0].startswith("chembl_id,canonical_smiles,ExactMolWt")
    assert len(lines) == 4
    assert "rows_emitted" in capsys.readouterr().out


def test_workload_cli_runs_similarity_search(tmp_path: Path) -> None:
    dataset = tmp_path / "molecules.tsv"
    dataset.write_text(
        "chembl_id\tcanonical_smiles\nQUERY\tCCO\nMATCH\tCCCO\n",
        encoding="utf-8",
    )
    output = tmp_path / "search.csv"

    code = main(
        [
            "workload",
            "run",
            "similarity-search",
            "--input",
            str(dataset),
            "--params",
            '{"query_smiles": "CCO", "top_k": 5, "progress_every": 0}',
            "--shard-rows",
            "2",
            "-o",
            str(output),
        ]
    )

    assert code == 0
    lines = output.read_text(encoding="utf-8").splitlines()
    assert lines[0] == "rank,chembl_id,canonical_smiles,similarity"
    assert len(lines) == 2


def test_workload_cli_rejects_unknown_or_missing_workload(tmp_path: Path) -> None:
    import pytest

    dataset = tmp_path / "molecules.tsv"
    dataset.write_text("chembl_id\tcanonical_smiles\nA\tCCO\n", encoding="utf-8")
    assert main(["workload", "run", "no-such-workload", "--input", str(dataset)]) == 1
    with pytest.raises(SystemExit):
        main(["workload", "run", "descriptor-batch"])


def test_workload_cli_rejects_invalid_params_json(tmp_path: Path) -> None:
    dataset = tmp_path / "molecules.tsv"
    dataset.write_text("chembl_id\tcanonical_smiles\nA\tCCO\n", encoding="utf-8")
    assert (
        main(
            [
                "workload",
                "run",
                "descriptor-batch",
                "--input",
                str(dataset),
                "--params",
                "{broken",
            ]
        )
        == 1
    )


def test_workload_cli_runs_an_allowlisted_custom_workload(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
) -> None:
    import csv

    from scimesh.sdk import (
        ArtifactSchema,
        ComponentRef,
        MapReduceWorkload,
        PortSpec,
        SchemaRef,
        WorkloadId,
    )
    from scimesh.sdk.registry import WorkloadRegistry
    from scimesh.workloads.environment import (
        current_environment_digest,
        current_scimesh_package_digest,
    )

    class CountRowsWorkload(MapReduceWorkload):
        workload_id = WorkloadId("count-rows", "1.0.0")
        description = "Count TSV data rows per shard."
        parameters_schema = {
            "type": "object",
            "additionalProperties": False,
            "properties": {},
        }
        input_port = PortSpec(
            ArtifactSchema(
                SchemaRef("molecule-table", 1),
                "text/tab-separated-values",
                "utf-8",
                10**9,
                ComponentRef("delimited-table", 1),
                validator_configuration={
                    "required_columns": ["canonical_smiles", "chembl_id"]
                },
                max_records=10**8,
            )
        )
        partial_port = output_port = PortSpec(
            ArtifactSchema(
                SchemaRef("count-table", 1),
                "text/csv",
                "utf-8",
                10**9,
                ComponentRef("delimited-table", 1),
                validator_configuration={"columns": ["id", "rows"]},
                max_records=10**8,
            )
        )

        def partition_input(self, input_path, parameters, workspace):
            paths = []
            with input_path.open(encoding="utf-8", newline="") as source:
                for index, row in enumerate(csv.DictReader(source, delimiter="\t")):
                    path = workspace / f"shard-{index}.tsv"
                    path.write_text(
                        "chembl_id\tcanonical_smiles\n"
                        + row["chembl_id"]
                        + "\t"
                        + row["canonical_smiles"]
                        + "\n",
                        encoding="utf-8",
                    )
                    paths.append(path)
            return paths

        def compute_shard(self, inputs, parameters, output_path):
            lines = inputs["input"].read_text(encoding="utf-8").splitlines()
            rows = max(len(lines) - 1, 0)
            output_path.write_text(
                "id,rows\nshard," + str(rows) + "\n", encoding="utf-8"
            )
            return {"rows": rows}

        def reduce_partials(self, partial_paths, parameters, output_path):
            with output_path.open("w", encoding="utf-8") as destination:
                destination.write("id,rows\n")
                total = 0
                for partial in partial_paths:
                    rows = partial.read_text(encoding="utf-8").splitlines()[1:]
                    destination.write("".join(row + "\n" for row in rows))
                    total += len(rows)
            return {"rows_total": total, "partial_count": len(partial_paths)}

    definition = CountRowsWorkload(
        package_digest=current_scimesh_package_digest(),
        environment_digest=current_environment_digest(),
    ).definition()

    def fake_discover(self: WorkloadRegistry, allowlist) -> None:
        self.register(definition, enabled=True)

    monkeypatch.setattr(WorkloadRegistry, "discover_installed", fake_discover)
    monkeypatch.setenv(
        "SCIMESH_WORKLOAD_ALLOWLIST",
        '[{"distribution": "scimesh", "name": "count-rows", "version": "1.0.0", '
        '"digest": "' + current_scimesh_package_digest() + '"}]',
    )

    dataset = tmp_path / "molecules.tsv"
    dataset.write_text(
        "chembl_id\tcanonical_smiles\nA\tCCO\nB\tCCCC\n", encoding="utf-8"
    )
    output = tmp_path / "counts.csv"
    assert (
        main(
            [
                "workload",
                "run",
                "count-rows",
                "--input",
                str(dataset),
                "-o",
                str(output),
            ]
        )
        == 0
    )
    assert output.read_text(encoding="utf-8") == "id,rows\nshard,1\nshard,1\n"


def test_workload_cli_exports_the_library_as_json(tmp_path: Path) -> None:
    import json

    output = tmp_path / "workloads.json"
    assert main(["workload", "export", "-o", str(output)]) == 0
    payload = json.loads(output.read_text(encoding="utf-8"))
    assert payload["schema_version"] == 2
    names = [item["name"] for item in payload["workloads"]]
    assert names == sorted(
        [
            "descriptor-batch",
            "molwt-filter",
            "similarity-graph",
            "similarity-search",
            "similarity-search-parallel",
        ]
    )
    for item in payload["workloads"]:
        assert item["version"] == "1.0.0"
        assert item["enabled"] is True
        assert item["determinism"] == "byte_exact"
        assert item["verifier"] == "exact-artifact@1"
        assert "input" in item["inputs"]
        assert "result" in item["outputs"]
    molwt = next(
        item for item in payload["workloads"] if item["name"] == "molwt-filter"
    )
    assert molwt["parameters_schema"]["properties"] == {
        "min_molwt": {
            "type": "number",
            "minimum": 0,
            "description": "Keep molecules with MolWt >= this value",
        },
        "max_molwt": {
            "type": "number",
            "minimum": 0,
            "description": "Keep molecules with MolWt <= this value",
        },
        "skip_invalid": {
            "type": "boolean",
            "default": True,
            "description": "Skip rows with invalid SMILES instead of failing",
        },
    }
    assert molwt["ui_elements"] == [
        {
            "field": "min_molwt",
            "widget": "number",
            "label": "Minimum molecular weight",
            "help": "Keep molecules with MolWt at least this value. Optional.",
            "placeholder": "e.g. 100",
            "options": [],
            "default": None,
            "order": 1,
            "group": "",
        },
        {
            "field": "max_molwt",
            "widget": "number",
            "label": "Maximum molecular weight",
            "help": "Keep molecules with MolWt at most this value. Optional.",
            "placeholder": "e.g. 600",
            "options": [],
            "default": None,
            "order": 2,
            "group": "",
        },
        {
            "field": "skip_invalid",
            "widget": "checkbox",
            "label": "Skip invalid molecules",
            "help": "Skip rows with invalid SMILES instead of failing the shard.",
            "placeholder": "",
            "options": [],
            "default": True,
            "order": 3,
            "group": "",
        },
    ]