add test for utils

better · Sep 1, 2019 · 8402508 · 8402508
1 parent a59d546
commit 8402508
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/test_convoys.py b/test_convoys.py
@@ -288,6 +288,14 @@ def test_convert_dataframe_timedeltas():
         assert 0 <= t2 - t1 < 3.0 / (24*60*60)
 
 
+def test_convert_dataframe_more_args():
+    df = _generate_dataframe()
+    unit, groups, (G, B, T) = convoys.utils.get_arrays(df, max_groups=2)
+    assert len(groups) <= 2
+    unit, groups, (G, B, T) = convoys.utils.get_arrays(df, group_min_size=9999)
+    assert G.shape == (0,)
+
+
 def _test_plot_cohorts(model='weibull', extra_model=None):
     df = _generate_dataframe()
     unit, groups, (G, B, T) = convoys.utils.get_arrays(df)