add test to cover the opposite situation

Author: drbenvincent

causalpy/pymc_experiments.py

@@ -60,6 +60,17 @@ def __init__(
         **kwargs,
     ) -> None:
         super().__init__(model=model, **kwargs)
+
+        # Input validation
+        if isinstance(data.index, pd.DatetimeIndex):
+            assert isinstance(
+                treatment_time, pd.Timestamp
+            ), "If data.index is DatetimeIndex, treatment_time must be pd.Timestamp."
+        else:
+            assert (
+                isinstance(treatment_time, pd.Timestamp) is False
+            ), "If treatment_time is pd.Timestamp, this only makese sense if data.index is DatetimeIndex."  # noqa: E501
+
         self.treatment_time = treatment_time
         # split data in to pre and post intervention
         self.datapre = data[data.index <= self.treatment_time]
@@ -81,12 +92,6 @@ def __init__(
         self.post_X = np.asarray(new_x)
         self.post_y = np.asarray(new_y)
 
-        # Input validation
-        if isinstance(data.index, pd.DatetimeIndex):
-            assert isinstance(
-                treatment_time, pd.Timestamp
-            ), "If data.index is DatetimeIndex, treatment_time must be pd.Timestamp."
-
         # DEVIATION FROM SKL EXPERIMENT CODE =============================
         # fit the model to the observed (pre-intervention) data
         COORDS = {"coeffs": self.labels, "obs_indx": np.arange(self.pre_X.shape[0])}

causalpy/tests/test_integration_pymc_examples.py

@@ -191,6 +191,21 @@ def test_sc():
     assert len(result.idata.posterior.coords["draw"]) == sample_kwargs["draws"]
 
 
+@pytest.mark.integration
+def test_sc_input_error():
+    """Test that an error is raised if the data index is not datetime and the
+    treatment time is pd.Timestamp."""
+    with pytest.raises(AssertionError):
+        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 = cp.load_data("brexit")