chore: fixes unit test

Author: DiogoRibeiro7

gen_surv/visualization.py

@@ -1,11 +1,13 @@
 """
 Visualization utilities for survival data.
 
-This module provides functions to visualize survival data generated by gen_surv,
-including Kaplan-Meier survival curves and other commonly used plots in survival analysis.
+This module provides functions to visualize survival data generated by 
+gen_surv,
+including Kaplan-Meier survival curves and other commonly used plots in 
+survival analysis.
 """
 
-from typing import Dict, List, Optional, Tuple, Union, Any
+from typing import Dict, Optional, Tuple, Union, Any
 import numpy as np
 import pandas as pd
 import matplotlib.pyplot as plt
@@ -217,7 +219,8 @@ def plot_covariate_effect(
     ci_alpha: float = 0.2,
 ) -> Tuple[Figure, Axes]:
     """
-    Visualize the effect of a continuous covariate on survival by discretizing it.
+    Visualize the effect of a continuous covariate on survival by discretizing 
+    it.
 
     Parameters
     ----------

tests/test_aft.py

@@ -1,11 +1,13 @@
 """
 Tests for Accelerated Failure Time (AFT) models.
 """
-
+import os
+import sys
 import pandas as pd
 import pytest
 from hypothesis import given, strategies as st
 
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
 from gen_surv.aft import gen_aft_log_normal, gen_aft_weibull, gen_aft_log_logistic
 
 

tests/test_bivariate.py

@@ -1,4 +1,8 @@
+import os
+import sys
 import numpy as np
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
 from gen_surv.bivariate import sample_bivariate_distribution
 import pytest
 

tests/test_cli.py

@@ -1,10 +1,19 @@
+import sys
+import os
+import runpy
+
 import pandas as pd
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
 from gen_surv.cli import dataset
-import runpy
 
 
 def test_cli_dataset_stdout(monkeypatch, capsys):
-    """Dataset command prints CSV to stdout when no output file is given."""
+    """
+    Test that the 'dataset' CLI command prints the generated CSV data to stdout when no output file is specified.
+    This test patches the 'generate' function to return a simple DataFrame, invokes the CLI command directly,
+    and asserts that the expected CSV header appears in the captured standard output.
+    """
 
     def fake_generate(model: str, n: int):
         return pd.DataFrame({"time": [1.0], "status": [1], "X0": [0.1], "X1": [0.2]})