Increase test coverage (#52)

Author: DiogoRibeiro7

.github/workflows/docs.yml

@@ -21,7 +21,7 @@ jobs:
       - name: Build documentation
         run: poetry run sphinx-build -W -b html docs/source docs/build
       - name: Upload artifacts
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
           name: documentation
           path: docs/build/

tests/test_cli.py

@@ -1,11 +1,12 @@
 import os
 import runpy
 import sys
+import pytest
 
 import pandas as pd
 
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
-from gen_surv.cli import dataset
+from gen_surv.cli import dataset, visualize
 
 
 def test_cli_dataset_stdout(monkeypatch, capsys):
@@ -81,7 +82,9 @@ def fake_generate(**kwargs):
         return pd.DataFrame({"time": [1], "status": [0]})
 
     monkeypatch.setattr("gen_surv.cli.generate", fake_generate)
-    dataset(model="aft_weibull", n=3, beta=[0.1, 0.2], shape=1.1, scale=2.2, output=None)
+    dataset(
+        model="aft_weibull", n=3, beta=[0.1, 0.2], shape=1.1, scale=2.2, output=None
+    )
     assert captured["model"] == "aft_weibull"
     assert captured["beta"] == [0.1, 0.2]
     assert captured["shape"] == 1.1
@@ -146,3 +149,109 @@ def fake_generate(**kwargs):
     assert captured["betas_survival"] == [0.4]
     assert captured["betas_cure"] == [0.4]
 
+
+def test_dataset_invalid_model(monkeypatch):
+    def fake_generate(**kwargs):
+        raise ValueError("bad model")
+
+    monkeypatch.setattr("gen_surv.cli.generate", fake_generate)
+    with pytest.raises(ValueError):
+        dataset(model="nope", n=1, output=None)
+
+
+def test_cli_visualize_basic(monkeypatch, tmp_path):
+    csv = tmp_path / "data.csv"
+    pd.DataFrame({"time": [1, 2], "status": [1, 0]}).to_csv(csv, index=False)
+
+    def fake_plot_survival_curve(**kwargs):
+        import matplotlib.pyplot as plt
+
+        fig, ax = plt.subplots()
+        ax.plot([0, 1], [1, 0])
+        return fig, ax
+
+    monkeypatch.setattr(
+        "gen_surv.visualization.plot_survival_curve", fake_plot_survival_curve
+    )
+
+    saved = []
+
+    def fake_savefig(path, *args, **kwargs):
+        saved.append(path)
+
+    monkeypatch.setattr("matplotlib.pyplot.savefig", fake_savefig)
+
+    visualize(
+        str(csv),
+        time_col="time",
+        status_col="status",
+        group_col=None,
+        output=str(tmp_path / "plot.png"),
+    )
+    assert saved and saved[0].endswith("plot.png")
+
+
+def test_dataset_aft_log_logistic(monkeypatch):
+    captured = {}
+
+    def fake_generate(**kwargs):
+        captured.update(kwargs)
+        return pd.DataFrame({"time": [1], "status": [1]})
+
+    monkeypatch.setattr("gen_surv.cli.generate", fake_generate)
+    dataset(
+        model="aft_log_logistic",
+        n=1,
+        beta=[0.1],
+        shape=1.2,
+        scale=2.3,
+        output=None,
+    )
+    assert captured["model"] == "aft_log_logistic"
+    assert captured["beta"] == [0.1]
+    assert captured["shape"] == 1.2
+    assert captured["scale"] == 2.3
+
+
+def test_dataset_competing_risks_weibull(monkeypatch):
+    captured = {}
+
+    def fake_generate(**kwargs):
+        captured.update(kwargs)
+        return pd.DataFrame({"time": [1], "status": [1]})
+
+    monkeypatch.setattr("gen_surv.cli.generate", fake_generate)
+    dataset(
+        model="competing_risks_weibull",
+        n=1,
+        n_risks=2,
+        shape_params=[0.7, 1.2],
+        scale_params=[2.0, 2.0],
+        beta=0.3,
+        output=None,
+    )
+    assert captured["n_risks"] == 2
+    assert captured["shape_params"] == [0.7, 1.2]
+    assert captured["scale_params"] == [2.0, 2.0]
+    assert captured["betas"] == [0.3, 0.3]
+
+
+def test_dataset_piecewise(monkeypatch):
+    captured = {}
+
+    def fake_generate(**kwargs):
+        captured.update(kwargs)
+        return pd.DataFrame({"time": [1], "status": [1]})
+
+    monkeypatch.setattr("gen_surv.cli.generate", fake_generate)
+    dataset(
+        model="piecewise_exponential",
+        n=1,
+        breakpoints=[1.0],
+        hazard_rates=[0.2, 0.3],
+        beta=[0.4],
+        output=None,
+    )
+    assert captured["breakpoints"] == [1.0]
+    assert captured["hazard_rates"] == [0.2, 0.3]
+    assert captured["betas"] == [0.4]

tests/test_competing_risks.py

@@ -8,6 +8,7 @@
 from hypothesis import given
 from hypothesis import strategies as st
 
+import gen_surv.competing_risks as cr
 from gen_surv.competing_risks import (
     cause_specific_cumulative_incidence,
     gen_competing_risks,
@@ -183,3 +184,33 @@ def test_reproducibility():
 
     with pytest.raises(AssertionError):
         pd.testing.assert_frame_equal(df1, df3)
+
+
+def test_competing_risks_summary_basic():
+    df = gen_competing_risks(n=10, n_risks=2, seed=1)
+    summary = cr.competing_risks_summary(df)
+    assert summary["n_subjects"] == 10
+    assert summary["n_causes"] == 2
+    assert set(summary["events_by_cause"]) <= {1, 2}
+    assert "time_stats" in summary
+
+
+def test_competing_risks_summary_with_categorical():
+    df = gen_competing_risks(n=8, n_risks=2, seed=2)
+    df["group"] = ["A", "B"] * 4
+    summary = cr.competing_risks_summary(df, covariate_cols=["X0", "group"])
+    assert summary["covariate_stats"]["group"]["categories"] == 2
+    assert "distribution" in summary["covariate_stats"]["group"]
+
+
+import matplotlib
+
+matplotlib.use("Agg")
+
+
+def test_plot_cause_specific_hazards_runs():
+    df = gen_competing_risks(n=30, n_risks=2, seed=3)
+    fig, ax = cr.plot_cause_specific_hazards(df, time_points=np.linspace(0, 5, 5))
+    assert hasattr(fig, "savefig")
+    assert len(ax.get_lines()) >= 1
+    matplotlib.pyplot.close(fig)

tests/test_piecewise.py

@@ -19,6 +19,7 @@ def test_piecewise_invalid_lengths():
             n=5, breakpoints=[1.0, 2.0], hazard_rates=[0.5], seed=42
         )
 
+
 def test_piecewise_invalid_hazard_and_breakpoints():
     with pytest.raises(ValueError):
         gen_piecewise_exponential(
@@ -34,3 +35,38 @@ def test_piecewise_invalid_hazard_and_breakpoints():
             hazard_rates=[0.5, -1.0],
             seed=42,
         )
+
+
+def test_piecewise_covariate_distributions():
+    for dist, params in [
+        ("uniform", {"low": 0.0, "high": 1.0}),
+        ("binary", {"p": 0.7}),
+    ]:
+        df = gen_piecewise_exponential(
+            n=5,
+            breakpoints=[1.0],
+            hazard_rates=[0.2, 0.4],
+            covariate_dist=dist,
+            covariate_params=params,
+            seed=1,
+        )
+        assert len(df) == 5
+        assert {"X0", "X1"}.issubset(df.columns)
+
+
+def test_piecewise_custom_betas_reproducible():
+    df1 = gen_piecewise_exponential(
+        n=5,
+        breakpoints=[1.0],
+        hazard_rates=[0.1, 0.2],
+        betas=[0.5, -0.2],
+        seed=2,
+    )
+    df2 = gen_piecewise_exponential(
+        n=5,
+        breakpoints=[1.0],
+        hazard_rates=[0.1, 0.2],
+        betas=[0.5, -0.2],
+        seed=2,
+    )
+    pd.testing.assert_frame_equal(df1, df2)

tests/test_summary_extra.py

@@ -1,6 +1,10 @@
 import pandas as pd
 import pytest
-from gen_surv.summary import check_survival_data_quality, compare_survival_datasets
+from gen_surv.summary import (
+    check_survival_data_quality,
+    compare_survival_datasets,
+    _print_summary,
+)
 from gen_surv import generate
 
 
@@ -35,17 +39,65 @@ def test_check_survival_data_quality_no_fix():
 
 
 def test_compare_survival_datasets_basic():
-    ds1 = generate(model="cphm", n=5, model_cens="uniform", cens_par=1.0, beta=0.5, covariate_range=1.0)
-    ds2 = generate(model="cphm", n=5, model_cens="uniform", cens_par=1.0, beta=1.0, covariate_range=1.0)
+    ds1 = generate(
+        model="cphm",
+        n=5,
+        model_cens="uniform",
+        cens_par=1.0,
+        beta=0.5,
+        covariate_range=1.0,
+    )
+    ds2 = generate(
+        model="cphm",
+        n=5,
+        model_cens="uniform",
+        cens_par=1.0,
+        beta=1.0,
+        covariate_range=1.0,
+    )
     comparison = compare_survival_datasets({"A": ds1, "B": ds2})
     assert set(["A", "B"]).issubset(comparison.columns)
     assert "n_subjects" in comparison.index
 
 
 def test_compare_survival_datasets_with_covariates_and_empty_error():
-    ds = generate(model="cphm", n=3, model_cens="uniform", cens_par=1.0, beta=0.5, covariate_range=1.0)
+    ds = generate(
+        model="cphm",
+        n=3,
+        model_cens="uniform",
+        cens_par=1.0,
+        beta=0.5,
+        covariate_range=1.0,
+    )
     comparison = compare_survival_datasets({"only": ds}, covariate_cols=["X0"])
     assert "only" in comparison.columns
     assert "X0_mean" in comparison.index
     with pytest.raises(ValueError):
         compare_survival_datasets({})
+
+
+def test_check_survival_data_quality_min_and_max():
+    df = pd.DataFrame({"time": [-1.0, 3.0], "status": [1, 1]})
+    fixed, issues = check_survival_data_quality(
+        df, min_time=0.0, max_time=2.0, fix_issues=True
+    )
+    assert (fixed["time"] <= 2.0).all()
+    assert issues["modifications"]["values_fixed"] > 0
+
+
+def test_print_summary_with_issues(capsys):
+    summary = {
+        "dataset_info": {"n_subjects": 2, "n_unique_ids": 2, "n_covariates": 0},
+        "event_info": {"n_events": 1, "n_censored": 1, "event_rate": 0.5},
+        "time_info": {"min": 0.0, "max": 2.0, "mean": 1.0, "median": 1.0},
+        "data_quality": {
+            "missing_time": 0,
+            "missing_status": 0,
+            "negative_time": 1,
+            "invalid_status": 0,
+        },
+        "covariates": {},
+    }
+    _print_summary(summary, "time", "status", None, [])
+    out = capsys.readouterr().out
+    assert "Issues detected" in out

tests/test_visualization.py

@@ -162,11 +162,12 @@ def fake_import(name, *args, **kwargs):
 
 def test_cli_visualize_read_error(monkeypatch, tmp_path, capsys):
     """visualize handles CSV read failures gracefully."""
-    monkeypatch.setattr("pandas.read_csv", lambda *a, **k: (_ for _ in ()).throw(Exception("boom")))
+    monkeypatch.setattr(
+        "pandas.read_csv", lambda *a, **k: (_ for _ in ()).throw(Exception("boom"))
+    )
     csv_path = tmp_path / "x.csv"
     csv_path.write_text("time,status\n1,1\n")
     with pytest.raises(typer.Exit):
         visualize(str(csv_path))
     captured = capsys.readouterr()
     assert "Error loading CSV file" in captured.out
-