Harden: CI, uniform selector contract, property tests, provenance

.github/workflows/ci.yml

@@ -0,0 +1,27 @@
+name: CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+          cache: pip
+
+      - name: Install package with dev extras
+        run: pip install -e ".[dev]"
+
+      - name: Lint (ruff)
+        run: ruff check .
+
+      - name: Test (pytest)
+        run: pytest -q

CLAUDE.md

@@ -280,6 +280,8 @@ pytest tests/              # unit tests (all should pass without external tools
 
 **Idempotent stages**: Each stage writes its output to a predictable path under `--outdir`. Re-running a stage skips existing outputs (file-existence check). The `--force` flag disables this.
 
+**Run provenance**: `RefRoverPipeline.run()` writes `params.json` to `--outdir` up front (selector, k, min-similarity, binners, threads, sample count, refrover version, UTC timestamp), so a run's configuration is recoverable even if a later stage fails.
+
 **No Nextflow dependency**: RefRover is pure Python + subprocess calls. It has no Nextflow or nf-core dependency and can run in any environment where the external tools are available.
 
 ---

pyproject.toml

@@ -28,3 +28,5 @@ where = ["src"]
 [tool.ruff]
 line-length = 100
 target-version = "py311"
+# Exploratory research/analysis scripts are not held to package lint standards.
+extend-exclude = ["data"]

src/refrover/pipeline.py

@@ -51,11 +51,16 @@ def run(self) -> PipelineResults:
         from refrover.formatters import format_for_binner
 
         outdir = self.outdir
+        outdir.mkdir(parents=True, exist_ok=True)
         sketch_dir = outdir / "sketches"
         bam_dir = outdir / "bams"
         cov_dir = outdir / "coverage"
         fmt_dir = outdir / "formatted"
 
+        # 0. Record run provenance up front, so a params.json exists even if a
+        #    later stage fails.
+        self._write_provenance(outdir / "params.json")
+
         # 1-2. Build the selection matrix.
         # Containment selectors use a precomputed reads-vs-assembly containment
         # matrix; Jaccard selectors sketch the assemblies and compare them.
@@ -100,3 +105,25 @@ def run(self) -> PipelineResults:
             coverage_tables[binner] = out if isinstance(out, Path) else out[0]
 
         return PipelineResults(coverage_tables=coverage_tables, assignments=assignments)
+
+    def _write_provenance(self, path: Path) -> None:
+        """Write a params.json capturing the run configuration for reproducibility."""
+        import json
+        from datetime import datetime, timezone
+
+        from refrover import __version__
+
+        sel = self.selector
+        params = {
+            "refrover_version": __version__,
+            "timestamp_utc": datetime.now(timezone.utc).isoformat(),
+            "selector": type(sel).__name__,
+            "k": getattr(sel, "k", None),
+            "min_similarity": getattr(sel, "min_containment", getattr(sel, "min_jaccard", None)),
+            "binners": list(self.binners),
+            "threads": self.threads,
+            "n_samples": int(len(self.manifest)),
+            "uses_containment_matrix": self.containment_matrix is not None,
+            "force": self.force,
+        }
+        path.write_text(json.dumps(params, indent=2) + "\n")

src/refrover/selectors/random.py

@@ -4,7 +4,14 @@
 
 
 class RandomSelector(BaseSelector):
-    """Baseline: random k prototypes within Jaccard threshold."""
+    """
+    Baseline: the query's own assembly plus random prototypes within threshold.
+
+    Like the other selectors, the query is always the first prototype (it is the
+    reference its own reads are guaranteed to map to), which keeps random a fair
+    control: every selector is handed the same guaranteed anchor and differs only
+    in how it picks the remaining k-1.
+    """
 
     def __init__(self, k: int, min_jaccard: float = 0.1, seed: int | None = None):
         super().__init__(k, min_jaccard)
@@ -14,5 +21,9 @@ def select(self, sim_matrix: pd.DataFrame, query_id: str) -> list[str]:
         candidates = self._candidates(sim_matrix, query_id)
         if len(candidates) <= self.k:
             return self._trim_to_k(candidates)
+
         rng = random.Random(self.seed)
+        if query_id in candidates:
+            others = [c for c in candidates if c != query_id]
+            return [query_id] + rng.sample(others, self.k - 1)
         return rng.sample(candidates, self.k)

src/refrover/sketch.py

@@ -3,7 +3,6 @@
 """
 
 import subprocess
-import sys
 from pathlib import Path
 from typing import Sequence
 
@@ -43,7 +42,6 @@ def sketch_assemblies(
 
     # sourmash sketch dna processes multiple files in one call
     fasta_list = [str(fa) for fa, _ in to_sketch]
-    sig_list = [str(sig) for _, sig in to_sketch]
 
     cmd = [
         sourmash_path, "sketch", "dna",

tests/conftest.py

@@ -1,7 +1,6 @@
 import numpy as np
 import pandas as pd
 import pytest
-from pathlib import Path
 
 
 def _make_sim(matrix: np.ndarray, ids: list[str]) -> pd.DataFrame:

tests/test_containment.py

@@ -1,6 +1,5 @@
 """Tests for the package containment-matrix computation."""
 
-from pathlib import Path
 
 import pytest
 

tests/test_coverage.py

@@ -1,7 +1,6 @@
 """Tests for CoverM output normalization into RefRover's internal schema."""
 
 import pandas as pd
-import pytest
 
 from refrover.coverage import normalize_coverage, load_coverage
 

tests/test_formatters/test_generic.py

@@ -20,9 +20,10 @@ def test_generic_skips_existing(coverage_df, tmp_path):
 
 
 def test_generic_force_overwrites(coverage_df, tmp_path):
+    import time
     out1 = format_for_binner(coverage_df, binner="generic", outdir=tmp_path)
     mtime1 = out1.stat().st_mtime
-    import time; time.sleep(0.01)
+    time.sleep(0.01)
     out2 = format_for_binner(coverage_df, binner="generic", outdir=tmp_path, force=True)
     assert out2.stat().st_mtime >= mtime1
 

tests/test_provenance.py

@@ -0,0 +1,54 @@
+"""Test that the pipeline records run provenance."""
+
+import json
+
+import pandas as pd
+
+from refrover.pipeline import RefRoverPipeline
+from refrover.selectors import ContainmentSelector, MaxMinSelector
+
+
+def test_provenance_records_containment_run(tmp_path):
+    manifest = pd.DataFrame({
+        "sample_id": ["s0", "s1"],
+        "assembly": ["a0.fasta", "a1.fasta"],
+        "r1": ["s0_R1.fq", "s1_R1.fq"],
+    })
+    cont = pd.DataFrame([[1.0, 0.5], [0.5, 1.0]], index=["s0", "s1"], columns=["s0", "s1"])
+
+    pipe = RefRoverPipeline(
+        manifest=manifest,
+        selector=ContainmentSelector(k=2, min_containment=0.05),
+        binners=["metabat2", "semibin2"],
+        threads=4,
+        outdir=tmp_path,
+        containment_matrix=cont,
+    )
+    pipe._write_provenance(tmp_path / "params.json")
+
+    data = json.loads((tmp_path / "params.json").read_text())
+    assert data["selector"] == "ContainmentSelector"
+    assert data["k"] == 2
+    assert data["min_similarity"] == 0.05
+    assert data["binners"] == ["metabat2", "semibin2"]
+    assert data["n_samples"] == 2
+    assert data["uses_containment_matrix"] is True
+    assert "refrover_version" in data
+    assert "timestamp_utc" in data
+
+
+def test_provenance_records_jaccard_selector(tmp_path):
+    manifest = pd.DataFrame({"sample_id": ["s0"], "assembly": ["a0.fasta"], "r1": ["x.fq"]})
+    pipe = RefRoverPipeline(
+        manifest=manifest,
+        selector=MaxMinSelector(k=5, min_jaccard=0.2),
+        binners=["generic"],
+        outdir=tmp_path,
+    )
+    pipe._write_provenance(tmp_path / "params.json")
+
+    data = json.loads((tmp_path / "params.json").read_text())
+    assert data["selector"] == "MaxMinSelector"
+    assert data["k"] == 5
+    assert data["min_similarity"] == 0.2
+    assert data["uses_containment_matrix"] is False

tests/test_selectors/test_archetype.py

@@ -22,8 +22,7 @@ def test_spans_clusters(clustered_sim):
     """Archetypes should select samples from all 3 distinct clusters."""
     sel = ArchetypeSelector(k=3, min_jaccard=0.0)
     result = sel.select(clustered_sim, "s00")
-    cluster_of = lambda sid: int(sid[1:]) // 4
-    assert len({cluster_of(r) for r in result}) == 3
+    assert len({int(r[1:]) // 4 for r in result}) == 3
 
 
 def test_valid_sample_ids(clustered_sim):

tests/test_selectors/test_contract.py

@@ -0,0 +1,66 @@
+"""
+Shared contract tests applied to every selector in SELECTOR_REGISTRY.
+
+These lock the invariants the pipeline relies on, so a new selector that
+violates the contract fails here rather than producing subtly wrong assignments.
+"""
+
+import pytest
+
+from refrover.selectors import SELECTOR_REGISTRY
+
+# Selectors runnable on a square similarity/containment matrix without extra deps.
+# feedback is an unimplemented stub; archetype needs the optional `archetypes`
+# package and is exercised in its own test module.
+_CONTRACT_SELECTORS = ["random", "maxmin", "kmedoids", "greedy_var", "containment"]
+
+
+def _make(name, **kw):
+    return SELECTOR_REGISTRY[name](**kw)
+
+
+@pytest.mark.parametrize("name", _CONTRACT_SELECTORS)
+def test_returns_at_most_k(name, clustered_sim):
+    sel = _make(name, k=3, min_jaccard=0.1)
+    result = sel.select(clustered_sim, "s00")
+    assert len(result) <= 3
+
+
+@pytest.mark.parametrize("name", _CONTRACT_SELECTORS)
+def test_results_unique(name, clustered_sim):
+    sel = _make(name, k=4, min_jaccard=0.1)
+    result = sel.select(clustered_sim, "s00")
+    assert len(result) == len(set(result))
+
+
+@pytest.mark.parametrize("name", _CONTRACT_SELECTORS)
+def test_results_are_valid_candidates(name, clustered_sim):
+    """Every returned id clears the threshold for the query."""
+    sel = _make(name, k=3, min_jaccard=0.1)
+    result = sel.select(clustered_sim, "s00")
+    for rid in result:
+        assert rid in clustered_sim.columns
+        assert clustered_sim.loc["s00", rid] >= 0.1
+
+
+@pytest.mark.parametrize("name", _CONTRACT_SELECTORS)
+def test_query_included_first(name, clustered_sim):
+    """The query's own assembly is always the first prototype."""
+    sel = _make(name, k=3, min_jaccard=0.1)
+    result = sel.select(clustered_sim, "s00")
+    assert result[0] == "s00"
+
+
+@pytest.mark.parametrize("name", _CONTRACT_SELECTORS)
+def test_unknown_query_raises(name, clustered_sim):
+    sel = _make(name, k=3, min_jaccard=0.1)
+    with pytest.raises(KeyError):
+        sel.select(clustered_sim, "not_a_sample")
+
+
+@pytest.mark.parametrize("name", _CONTRACT_SELECTORS)
+def test_fewer_candidates_than_k_warns(name, sparse_sim):
+    sel = _make(name, k=6, min_jaccard=0.1)
+    with pytest.warns(UserWarning):
+        result = sel.select(sparse_sim, "s00")
+    assert len(result) <= 6

tests/test_selectors/test_kmedoids.py

@@ -1,4 +1,3 @@
-import pytest
 from refrover.selectors import KMedoidsSelector
 
 
@@ -29,5 +28,4 @@ def test_covers_clusters(clustered_sim):
     """k-medoids should select one representative from each cluster."""
     sel = KMedoidsSelector(k=3, min_jaccard=0.0)
     result = sel.select(clustered_sim, "s00")
-    cluster_of = lambda sid: int(sid[1:]) // 4
-    assert len({cluster_of(r) for r in result}) == 3
+    assert len({int(r[1:]) // 4 for r in result}) == 3

tests/test_selectors/test_maxmin.py

@@ -1,6 +1,3 @@
-import numpy as np
-import pandas as pd
-import pytest
 from refrover.selectors import MaxMinSelector
 
 
@@ -19,8 +16,7 @@ def test_selects_across_clusters(clustered_sim):
     """With k=3 and 3 clusters, MaxMin should pick one from each cluster."""
     sel = MaxMinSelector(k=3, min_jaccard=0.0)
     result = sel.select(clustered_sim, "s00")
-    cluster_of = lambda sid: int(sid[1:]) // 4
-    clusters_hit = {cluster_of(r) for r in result}
+    clusters_hit = {int(r[1:]) // 4 for r in result}
     assert len(clusters_hit) == 3