Expand validate tests (#49)

Author: DiogoRibeiro7

tests/test_cli.py

@@ -53,3 +53,96 @@ def fake_generate(model: str, n: int):
     assert out_file.exists()
     content = out_file.read_text()
     assert "time,status,X0,X1" in content
+
+
+def test_dataset_fallback(monkeypatch):
+    """If generate fails with additional kwargs, dataset retries with minimal args."""
+    calls = []
+
+    def fake_generate(**kwargs):
+        calls.append(kwargs)
+        if len(calls) == 1:
+            raise TypeError("bad args")
+        return pd.DataFrame({"time": [0], "status": [1]})
+
+    monkeypatch.setattr("gen_surv.cli.generate", fake_generate)
+    dataset(model="cphm", n=2, output=None)
+    # first call has many parameters, second only model and n
+    assert calls[-1] == {"model": "cphm", "n": 2}
+    assert len(calls) == 2
+
+
+def test_dataset_weibull_parameters(monkeypatch):
+    """Parameters for aft_weibull model are forwarded correctly."""
+    captured = {}
+
+    def fake_generate(**kwargs):
+        captured.update(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)
+    assert captured["model"] == "aft_weibull"
+    assert captured["beta"] == [0.1, 0.2]
+    assert captured["shape"] == 1.1
+    assert captured["scale"] == 2.2
+
+
+def test_dataset_aft_ln(monkeypatch):
+    """aft_ln model should forward beta list and sigma."""
+    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_ln", n=1, beta=[0.3, 0.4], sigma=1.2, output=None)
+    assert captured["beta"] == [0.3, 0.4]
+    assert captured["sigma"] == 1.2
+
+
+def test_dataset_competing_risks(monkeypatch):
+    """competing_risks expands betas and passes hazards."""
+    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",
+        n=1,
+        n_risks=2,
+        baseline_hazards=[0.1, 0.2],
+        beta=0.5,
+        output=None,
+    )
+    assert captured["n_risks"] == 2
+    assert captured["baseline_hazards"] == [0.1, 0.2]
+    assert captured["betas"] == [0.5, 0.5]
+
+
+def test_dataset_mixture_cure(monkeypatch):
+    """mixture_cure passes cure and baseline parameters."""
+    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="mixture_cure",
+        n=1,
+        cure_fraction=0.2,
+        baseline_hazard=0.1,
+        beta=[0.4],
+        output=None,
+    )
+    assert captured["cure_fraction"] == 0.2
+    assert captured["baseline_hazard"] == 0.1
+    assert captured["betas_survival"] == [0.4]
+    assert captured["betas_cure"] == [0.4]
+

tests/test_validate.py

@@ -42,6 +42,31 @@ def test_validate_gen_cmm_inputs_invalid_beta_length():
         )
 
 
+@pytest.mark.parametrize(
+    "n, model_cens, cens_par, cov_range, rate",
+    [
+        (0, "uniform", 0.5, 1.0, [0.1] * 6),
+        (1, "bad", 0.5, 1.0, [0.1] * 6),
+        (1, "uniform", 0.0, 1.0, [0.1] * 6),
+        (1, "uniform", 0.5, 0.0, [0.1] * 6),
+        (1, "uniform", 0.5, 1.0, [0.1] * 3),
+    ],
+)
+def test_validate_gen_cmm_inputs_other_invalid(
+    n, model_cens, cens_par, cov_range, rate
+):
+    with pytest.raises(ValueError):
+        v.validate_gen_cmm_inputs(
+            n, model_cens, cens_par, [0.1, 0.2, 0.3], cov_range, rate
+        )
+
+
+def test_validate_gen_cmm_inputs_valid():
+    v.validate_gen_cmm_inputs(
+        1, "uniform", 1.0, [0.1, 0.2, 0.3], covariate_range=1.0, rate=[0.1] * 6
+    )
+
+
 def test_validate_gen_tdcm_inputs_invalid_lambda():
     """Lambda <= 0 should raise a ValueError."""
     with pytest.raises(ValueError):
@@ -57,6 +82,44 @@ def test_validate_gen_tdcm_inputs_invalid_lambda():
         )
 
 
+@pytest.mark.parametrize(
+    "dist,corr,dist_par",
+    [
+        ("bad", 0.5, [1, 2]),
+        ("weibull", 0.0, [1, 2, 3, 4]),
+        ("weibull", 0.5, [1, 2, -1, 2]),
+        ("weibull", 0.5, [1, 2, 3]),
+        ("exponential", 2.0, [1, 1]),
+        ("exponential", 0.5, [1]),
+    ],
+)
+def test_validate_gen_tdcm_inputs_invalid_dist(dist, corr, dist_par):
+    with pytest.raises(ValueError):
+        v.validate_gen_tdcm_inputs(
+            1,
+            dist,
+            corr,
+            dist_par,
+            "uniform",
+            1.0,
+            beta=[0.1, 0.2, 0.3],
+            lam=1.0,
+        )
+
+
+def test_validate_gen_tdcm_inputs_valid():
+    v.validate_gen_tdcm_inputs(
+        1,
+        "weibull",
+        0.5,
+        [1, 1, 1, 1],
+        "uniform",
+        1.0,
+        beta=[0.1, 0.2, 0.3],
+        lam=1.0,
+    )
+
+
 def test_validate_gen_aft_log_normal_inputs_valid():
     """Valid parameters should not raise an error for AFT log-normal."""
     v.validate_gen_aft_log_normal_inputs(
@@ -68,19 +131,84 @@ def test_validate_gen_aft_log_normal_inputs_valid():
     )
 
 
+@pytest.mark.parametrize(
+    "n,beta,sigma,model_cens,cens_par",
+    [
+        (0, [0.1], 1.0, "uniform", 1.0),
+        (1, "bad", 1.0, "uniform", 1.0),
+        (1, [0.1], 0.0, "uniform", 1.0),
+        (1, [0.1], 1.0, "bad", 1.0),
+        (1, [0.1], 1.0, "uniform", 0.0),
+    ],
+)
+def test_validate_gen_aft_log_normal_inputs_invalid(
+    n, beta, sigma, model_cens, cens_par
+):
+    with pytest.raises(ValueError):
+        v.validate_gen_aft_log_normal_inputs(n, beta, sigma, model_cens, cens_par)
+
+
 def test_validate_dg_biv_inputs_valid_weibull():
     """Valid parameters for a Weibull distribution should pass."""
     v.validate_dg_biv_inputs(5, "weibull", 0.1, [1.0, 1.0, 1.0, 1.0])
 
 
+def test_validate_dg_biv_inputs_invalid_corr_and_params():
+    with pytest.raises(ValueError):
+        v.validate_dg_biv_inputs(1, "exponential", -2.0, [1.0, 1.0])
+    with pytest.raises(ValueError):
+        v.validate_dg_biv_inputs(1, "exponential", 0.5, [1.0])
+    with pytest.raises(ValueError):
+        v.validate_dg_biv_inputs(1, "weibull", 0.5, [1.0, 1.0])
+
+
 def test_validate_gen_aft_weibull_inputs_and_log_logistic():
     with pytest.raises(ValueError):
         v.validate_gen_aft_weibull_inputs(0, [0.1], 1.0, 1.0, "uniform", 1.0)
     with pytest.raises(ValueError):
         v.validate_gen_aft_log_logistic_inputs(1, [0.1], -1.0, 1.0, "uniform", 1.0)
 
 
+@pytest.mark.parametrize(
+    "shape,scale",
+    [(-1.0, 1.0), (1.0, -1.0)],
+)
+def test_validate_gen_aft_weibull_invalid_params(shape, scale):
+    with pytest.raises(ValueError):
+        v.validate_gen_aft_weibull_inputs(1, [0.1], shape, scale, "uniform", 1.0)
+
+
+def test_validate_gen_aft_weibull_valid():
+    v.validate_gen_aft_weibull_inputs(1, [0.1], 1.0, 1.0, "uniform", 1.0)
+
+
+def test_validate_gen_aft_log_logistic_valid():
+    v.validate_gen_aft_log_logistic_inputs(1, [0.1], 1.0, 1.0, "uniform", 1.0)
+
+
 def test_validate_competing_risks_inputs():
     with pytest.raises(ValueError):
         v.validate_competing_risks_inputs(1, 2, [0.1], None, "uniform", 1.0)
     v.validate_competing_risks_inputs(1, 1, [0.5], [[0.1]], "uniform", 0.5)
+
+
+@pytest.mark.parametrize(
+    "n,model_cens,cens_par,beta,cov_range,rate",
+    [
+        (0, "uniform", 1.0, [0.1, 0.2, 0.3], 1.0, [0.1, 0.2, 0.3]),
+        (1, "bad", 1.0, [0.1, 0.2, 0.3], 1.0, [0.1, 0.2, 0.3]),
+        (1, "uniform", 0.0, [0.1, 0.2, 0.3], 1.0, [0.1, 0.2, 0.3]),
+        (1, "uniform", 1.0, [0.1, 0.2], 1.0, [0.1, 0.2, 0.3]),
+        (1, "uniform", 1.0, [0.1, 0.2, 0.3], 0.0, [0.1, 0.2, 0.3]),
+        (1, "uniform", 1.0, [0.1, 0.2, 0.3], 1.0, [0.1]),
+    ],
+)
+def test_validate_gen_thmm_inputs_invalid(
+    n, model_cens, cens_par, beta, cov_range, rate
+):
+    with pytest.raises(ValueError):
+        v.validate_gen_thmm_inputs(n, model_cens, cens_par, beta, cov_range, rate)
+
+
+def test_validate_gen_thmm_inputs_valid():
+    v.validate_gen_thmm_inputs(1, "uniform", 1.0, [0.1, 0.2, 0.3], 1.0, [0.1, 0.2, 0.3])

tests/test_visualization.py

@@ -118,3 +118,55 @@ def test_cli_visualize_missing_column(tmp_path, capsys):
         )
     captured = capsys.readouterr()
     assert "Status column 'status' not found in data" in captured.out
+
+
+def test_cli_visualize_missing_time(tmp_path, capsys):
+    df = pd.DataFrame({"t": [1, 2], "status": [1, 0]})
+    path = tmp_path / "d.csv"
+    df.to_csv(path, index=False)
+    with pytest.raises(typer.Exit):
+        visualize(str(path), time_col="time", status_col="status")
+    captured = capsys.readouterr()
+    assert "Time column 'time' not found in data" in captured.out
+
+
+def test_cli_visualize_missing_group(tmp_path, capsys):
+    df = pd.DataFrame({"time": [1], "status": [1], "x": [0]})
+    path = tmp_path / "d2.csv"
+    df.to_csv(path, index=False)
+    with pytest.raises(typer.Exit):
+        visualize(str(path), time_col="time", status_col="status", group_col="group")
+    captured = capsys.readouterr()
+    assert "Group column 'group' not found in data" in captured.out
+
+
+def test_cli_visualize_import_error(monkeypatch, tmp_path, capsys):
+    """visualize exits when matplotlib is missing."""
+    import builtins
+
+    real_import = builtins.__import__
+
+    def fake_import(name, *args, **kwargs):
+        if name.startswith("matplotlib"):  # simulate missing dependency
+            raise ImportError("no matplot")
+        return real_import(name, *args, **kwargs)
+
+    monkeypatch.setattr(builtins, "__import__", fake_import)
+    csv_path = tmp_path / "d.csv"
+    pd.DataFrame({"time": [1], "status": [1]}).to_csv(csv_path, index=False)
+    with pytest.raises(typer.Exit):
+        visualize(str(csv_path))
+    captured = capsys.readouterr()
+    assert "Visualization requires matplotlib" in captured.out
+
+
+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")))
+    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
+