move some tests to test_input_validation

Author: drbenvincent

causalpy/tests/test_input_validation.py

@@ -2,7 +2,11 @@
 import pytest
 
 import causalpy as cp
-from causalpy.custom_exceptions import DataException, FormulaException
+from causalpy.custom_exceptions import (
+    BadIndexException,
+    DataException,
+    FormulaException,
+)
 
 sample_kwargs = {"tune": 20, "draws": 20, "chains": 2, "cores": 2}
 
@@ -96,3 +100,44 @@ def test_did_validation_group_dummy_coded():
             group_variable_name="group",
             model=cp.pymc_models.LinearRegression(sample_kwargs=sample_kwargs),
         )
+
+
+# Synthetic Control
+
+
+def test_sc_input_error():
+    """Confirm that a BadIndexException is raised treatment_time is pd.Timestamp
+    and df.index is not pd.DatetimeIndex."""
+    with pytest.raises(BadIndexException):
+        df = cp.load_data("sc")
+        treatment_time = pd.to_datetime("2016 June 24")
+        _ = cp.pymc_experiments.SyntheticControl(
+            df,
+            treatment_time,
+            formula="actual ~ 0 + a + b + c + d + e + f + g",
+            model=cp.pymc_models.WeightedSumFitter(sample_kwargs=sample_kwargs),
+        )
+
+
+def test_sc_brexit_input_error():
+    """Confirm a BadIndexException is raised if the data index is datetime and the
+    treatment time is not pd.Timestamp."""
+    with pytest.raises(BadIndexException):
+        df = cp.load_data("brexit")
+        df["Time"] = pd.to_datetime(df["Time"])
+        df.set_index("Time", inplace=True)
+        df = df.iloc[df.index > "2009", :]
+        treatment_time = "2016 June 24"  # NOTE This is not of type pd.Timestamp
+        df = df.drop(["Japan", "Italy", "US", "Spain"], axis=1)
+        target_country = "UK"
+        all_countries = df.columns
+        other_countries = all_countries.difference({target_country})
+        all_countries = list(all_countries)
+        other_countries = list(other_countries)
+        formula = target_country + " ~ " + "0 + " + " + ".join(other_countries)
+        _ = cp.pymc_experiments.SyntheticControl(
+            df,
+            treatment_time,
+            formula=formula,
+            model=cp.pymc_models.WeightedSumFitter(sample_kwargs=sample_kwargs),
+        )

causalpy/tests/test_integration_pymc_examples.py

@@ -2,7 +2,6 @@
 import pytest
 
 import causalpy as cp
-from causalpy.custom_exceptions import BadIndexException
 
 sample_kwargs = {"tune": 20, "draws": 20, "chains": 2, "cores": 2}
 
@@ -196,21 +195,6 @@ def test_sc():
     assert len(result.idata.posterior.coords["draw"]) == sample_kwargs["draws"]
 
 
-@pytest.mark.integration
-def test_sc_input_error():
-    """Confirm that a BadIndexException is raised treatment_time is pd.Timestamp
-    and df.index is not pd.DatetimeIndex."""
-    with pytest.raises(BadIndexException):
-        df = cp.load_data("sc")
-        treatment_time = pd.to_datetime("2016 June 24")
-        _ = cp.pymc_experiments.SyntheticControl(
-            df,
-            treatment_time,
-            formula="actual ~ 0 + a + b + c + d + e + f + g",
-            model=cp.pymc_models.WeightedSumFitter(sample_kwargs=sample_kwargs),
-        )
-
-
 @pytest.mark.integration
 def test_sc_brexit():
     df = (
@@ -239,31 +223,6 @@ def test_sc_brexit():
     assert len(result.idata.posterior.coords["draw"]) == sample_kwargs["draws"]
 
 
-@pytest.mark.integration
-def test_sc_brexit_input_error():
-    """Confirm a BadIndexException is raised if the data index is datetime and the
-    treatment time is not pd.Timestamp."""
-    with pytest.raises(BadIndexException):
-        df = cp.load_data("brexit")
-        df["Time"] = pd.to_datetime(df["Time"])
-        df.set_index("Time", inplace=True)
-        df = df.iloc[df.index > "2009", :]
-        treatment_time = "2016 June 24"  # NOTE This is not of type pd.Timestamp
-        df = df.drop(["Japan", "Italy", "US", "Spain"], axis=1)
-        target_country = "UK"
-        all_countries = df.columns
-        other_countries = all_countries.difference({target_country})
-        all_countries = list(all_countries)
-        other_countries = list(other_countries)
-        formula = target_country + " ~ " + "0 + " + " + ".join(other_countries)
-        _ = cp.pymc_experiments.SyntheticControl(
-            df,
-            treatment_time,
-            formula=formula,
-            model=cp.pymc_models.WeightedSumFitter(sample_kwargs=sample_kwargs),
-        )
-
-
 @pytest.mark.integration
 def test_ancova():
     df = cp.load_data("anova1")