Feat/add unit tests (#88)

* tests: remove regression baselines

* tests: add deterministic generation and visualize CLI integration

Author: DiogoRibeiro7

pyproject.toml

@@ -28,6 +28,7 @@ python = ">=3.10,<3.13"
 numpy = "^1.26"
 pandas = "^2.2.3"
 typer = "^0.12.3"
+click = ">=8.1,<8.2"
 matplotlib = "~3.8"
 lifelines = "^0.30"
 pyarrow = "^14"
@@ -36,10 +37,11 @@ pyreadr = "^0.5"
 [tool.poetry.group.dev.dependencies]
 pytest = "^8.3.5"
 pytest-cov = "^6.1.1"
+pytest-benchmark = "^4.0"
 python-semantic-release = "^9.21.0"
 mypy = "^1.15.0"
 invoke = "^2.2.0"
-hypothesis = "^6.98"
+hypothesis = "^6.108"
 tomli = "^2.2.1"
 black = "^24.1.0"
 isort = "^5.13.2"
@@ -51,6 +53,7 @@ pre-commit = "^3.8"
 dev = [
     "pytest",
     "pytest-cov",
+    "pytest-benchmark",
     "python-semantic-release",
     "mypy",
     "invoke",
@@ -72,6 +75,11 @@ sphinx-copybutton = "^0.5.2"
 sphinx-design = "^0.5.0"
 linkify-it-py = ">=2.0"
 
+[tool.pytest.ini_options]
+markers = [
+    "slow: slow performance tests",
+]
+
 [tool.poetry.scripts]
 gen_surv = "gen_surv.cli:app"
 

tests/conftest.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Callable
+
+import numpy as np
+import pandas as pd
+import pytest
+
+BASELINE_DIR = Path(__file__).parent / "baselines"
+BASELINE_DIR.mkdir(exist_ok=True)
+
+
+def pytest_addoption(parser: pytest.Parser) -> None:
+    parser.addoption(
+        "--update-baselines",
+        action="store_true",
+        default=False,
+        help="Refresh stored baselines for regression tests.",
+    )
+
+
+@pytest.fixture(scope="session")
+def rng() -> np.random.Generator:
+    return np.random.default_rng(seed=42)
+
+
+@pytest.fixture(scope="session")
+def save_baseline() -> Callable[[pd.DataFrame, str], None]:
+    def _save(df: pd.DataFrame, name: str) -> None:
+        (BASELINE_DIR / f"{name}.parquet").write_bytes(df.to_parquet(index=False))
+
+    return _save
+
+
+@pytest.fixture(scope="session")
+def load_baseline() -> Callable[[str], pd.DataFrame]:
+    def _load(name: str) -> pd.DataFrame:
+        path = BASELINE_DIR / f"{name}.parquet"
+        if not path.exists():
+            pytest.skip(
+                f"Missing baseline {path}; run with --update-baselines to refresh."
+            )
+        return pd.read_parquet(path)
+
+    return _load
+
+
+def assert_frame_numeric_equal(
+    got: pd.DataFrame,
+    expected: pd.DataFrame,
+    *,
+    rtol: float = 1e-6,
+    atol: float = 1e-8,
+) -> None:
+    assert list(got.columns) == list(expected.columns), "Column order/name changed."
+    assert got.shape == expected.shape, "Shape changed."
+    for col in got.columns:
+        g = pd.to_numeric(got[col], errors="coerce")
+        e = pd.to_numeric(expected[col], errors="coerce")
+        if g.notna().all() and e.notna().all():
+            np.testing.assert_allclose(g.to_numpy(), e.to_numpy(), rtol=rtol, atol=atol)
+        else:
+            assert (
+                got[col].astype(str).values == expected[col].astype(str).values
+            ).all(), f"Mismatch in column {col!r}"

tests/test_api_contract.py

@@ -0,0 +1,10 @@
+from __future__ import annotations
+
+import inspect
+
+import gen_surv
+
+
+def test_generate_signature_stable() -> None:
+    sig = inspect.signature(gen_surv.generate)
+    assert "model:" in str(sig) and "**kwargs" in str(sig)

tests/test_cli_snapshot.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+import subprocess
+import sys
+
+
+def run_cli(args: list[str]) -> subprocess.CompletedProcess[str]:
+    return subprocess.run(
+        [sys.executable, "-m", "gen_surv", *args],
+        text=True,
+        capture_output=True,
+    )
+
+
+def test_cli_help_shows_usage() -> None:
+    cp = run_cli(["--help"])
+    assert cp.returncode == 0, cp.stderr
+    assert "Generate synthetic survival datasets." in cp.stdout

tests/test_cli_visualize_integration.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+from typer.testing import CliRunner
+
+from gen_surv import generate
+from gen_surv.cli import app
+
+
+def test_visualize_cli_generates_plot(tmp_path) -> None:
+    csv_path = tmp_path / "data.csv"
+    plot_path = tmp_path / "plot.png"
+    df = generate(
+        model="cphm",
+        n=10,
+        beta=0.5,
+        covariate_range=1.0,
+        model_cens="uniform",
+        cens_par=0.7,
+        seed=1234,
+    )
+    df.to_csv(csv_path, index=False)
+    runner = CliRunner()
+    result = runner.invoke(
+        app, ["visualize", str(csv_path), "--output", str(plot_path)]
+    )
+    assert result.exit_code == 0
+    assert plot_path.exists()

tests/test_generate_deterministic.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+from gen_surv import generate
+from tests.conftest import assert_frame_numeric_equal
+
+
+def test_generate_reproducible_with_seed() -> None:
+    cfg = dict(
+        model="cphm",
+        n=32,
+        beta=0.5,
+        covariate_range=2.0,
+        model_cens="uniform",
+        cens_par=0.7,
+        seed=1234,
+    )
+    df1 = generate(**cfg)
+    df2 = generate(**cfg)
+    assert_frame_numeric_equal(df1, df2)

tests/test_generate_export_integration.py

@@ -0,0 +1,42 @@
+import pandas as pd
+import pytest
+
+from gen_surv import generate
+from gen_surv.export import export_dataset
+from gen_surv.integration import to_sksurv
+
+
+def test_generate_export_roundtrip(tmp_path):
+    """Integration test for generate and export_dataset."""
+    df = generate(
+        model="cphm",
+        n=10,
+        model_cens="uniform",
+        cens_par=1.0,
+        beta=0.5,
+        covariate_range=1.0,
+        seed=42,
+    )
+    out = tmp_path / "data.json"
+    export_dataset(df, out)
+    loaded = pd.read_json(out, orient="table")
+    pd.testing.assert_frame_equal(df, loaded)
+
+
+def test_generate_export_to_sksurv_roundtrip(tmp_path):
+    """Full pipeline from generation to scikit-survival array."""
+    pytest.importorskip("sksurv.util")
+    df = generate(
+        model="cphm",
+        n=8,
+        model_cens="uniform",
+        cens_par=1.0,
+        beta=0.5,
+        covariate_range=1.0,
+        seed=0,
+    )
+    out = tmp_path / "data.json"
+    export_dataset(df, out)
+    loaded = pd.read_json(out, orient="table")
+    arr = to_sksurv(loaded)
+    assert arr.shape[0] == 8

tests/test_generate_regression.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+from typing import Any, Dict
+
+import pandas as pd
+import pytest
+
+from gen_surv import generate
+
+MODEL_CONFIGS: Dict[str, Dict[str, Any]] = {
+    "cphm": dict(
+        model="cphm",
+        n=256,
+        beta=0.5,
+        covariate_range=2.0,
+        model_cens="uniform",
+        cens_par=0.7,
+        seed=1234,
+    ),
+    "aft_ln": dict(
+        model="aft_ln",
+        n=256,
+        beta=[0.5],
+        sigma=0.8,
+        model_cens="uniform",
+        cens_par=0.8,
+        seed=1234,
+    ),
+    "aft_log_logistic": dict(
+        model="aft_log_logistic",
+        n=256,
+        beta=[0.5],
+        shape=1.3,
+        scale=1.7,
+        model_cens="uniform",
+        cens_par=0.8,
+        seed=1234,
+    ),
+    "aft_weibull": dict(
+        model="aft_weibull",
+        n=256,
+        beta=[0.5],
+        shape=1.4,
+        scale=1.1,
+        model_cens="uniform",
+        cens_par=0.8,
+        seed=1234,
+    ),
+}
+
+
+@pytest.mark.parametrize("model_key", sorted(MODEL_CONFIGS.keys()))
+def test_generate_matches_baseline(
+    model_key: str,
+    request: pytest.FixtureRequest,
+    load_baseline,
+    save_baseline,
+) -> None:
+    cfg = MODEL_CONFIGS[model_key]
+    df: pd.DataFrame = generate(**cfg)
+    assert "time" in df.columns and (
+        "event" in df.columns or "status" in df.columns
+    ), "Missing core survival columns."
+    baseline_name = f"gen_{model_key}"
+    if request.config.getoption("--update-baselines"):
+        save_baseline(df, baseline_name)
+        pytest.skip(
+            f"Baseline {baseline_name} updated; re-run without --update-baselines."
+        )
+    expected = load_baseline(baseline_name)
+    from conftest import assert_frame_numeric_equal
+
+    assert_frame_numeric_equal(df[expected.columns], expected)

tests/test_integration_sksurv.py

@@ -1,13 +1,68 @@
+import sys
+import types
+
 import pandas as pd
 import pytest
 
 from gen_surv.integration import to_sksurv
+from gen_surv.interface import generate
 
 
 def test_to_sksurv():
-    # Optional integration test; skipped when scikit-survival is not installed.
+    """Basic conversion with default column names."""
     pytest.importorskip("sksurv.util")
     df = pd.DataFrame({"time": [1.0, 2.0], "status": [1, 0]})
     arr = to_sksurv(df)
     assert arr.dtype.names == ("status", "time")
     assert arr.shape[0] == 2
+
+
+def test_to_sksurv_custom_columns():
+    """Unit test for custom time/event column names."""
+    pytest.importorskip("sksurv.util")
+    df = pd.DataFrame({"T": [1.0, 2.0], "E": [1, 0]})
+    arr = to_sksurv(df, time_col="T", event_col="E")
+    assert arr.dtype.names == ("E", "T")
+
+
+def test_to_sksurv_missing_dependency(monkeypatch):
+    """Regression test ensuring a helpful ImportError is raised."""
+    fake_mod = types.ModuleType("sksurv")
+    monkeypatch.setitem(sys.modules, "sksurv", fake_mod)
+    monkeypatch.delitem(sys.modules, "sksurv.util", raising=False)
+    df = pd.DataFrame({"time": [1.0], "status": [1]})
+    with pytest.raises(ImportError, match="scikit-survival is required"):
+        to_sksurv(df)
+
+
+def test_to_sksurv_missing_columns():
+    """Regression test: missing required columns should raise KeyError."""
+    pytest.importorskip("sksurv.util")
+    df = pd.DataFrame({"status": [1, 0]})
+    with pytest.raises(KeyError):
+        to_sksurv(df)
+
+
+def test_to_sksurv_empty_dataframe():
+    """Unit test for handling empty DataFrames."""
+    pytest.importorskip("sksurv.util")
+    df = pd.DataFrame({"time": [], "status": []})
+    arr = to_sksurv(df)
+    assert arr.shape == (0,)
+    assert arr.dtype.names == ("status", "time")
+
+
+def test_generate_to_sksurv_pipeline():
+    """Integration test covering generation and conversion."""
+    pytest.importorskip("sksurv.util")
+    df = generate(
+        model="cphm",
+        n=5,
+        model_cens="uniform",
+        cens_par=1.0,
+        beta=0.5,
+        covariate_range=1.0,
+        seed=0,
+    )
+    arr = to_sksurv(df)
+    assert arr.shape[0] == 5

tests/test_perf_smoke.py

@@ -0,0 +1,22 @@
+from __future__ import annotations
+
+import pytest
+
+from gen_surv import generate
+
+
+@pytest.mark.slow
+def test_generate_perf_smoke(benchmark) -> None:
+    def _run():
+        return generate(
+            model="cphm",
+            n=50_000,
+            beta=0.5,
+            covariate_range=2.0,
+            model_cens="uniform",
+            cens_par=0.7,
+            seed=123,
+        )
+
+    df = benchmark(_run)
+    assert len(df) == 50_000