Unittest to check estimator type when formula is used

Author: christopher-wild

causal_testing/json_front/json_class.py

@@ -309,7 +309,7 @@ def _setup_test(self, causal_test_case: CausalTestCase, test: Mapping) -> Estima
         """
         estimator_kwargs = {}
         if "formula" in test:
-            if test["estimator"] != LinearRegressionEstimator or LogisticRegressionEstimator:
+            if test["estimator"] != (LinearRegressionEstimator or LogisticRegressionEstimator):
                 raise TypeError(
                     "Currently only LinearRegressionEstimator and LogisticRegressionEstimator supports the use of formulas"
                 )

tests/json_front_tests/test_json_class.py

@@ -4,7 +4,7 @@
 import scipy
 import os
 
-from causal_testing.testing.estimators import LinearRegressionEstimator
+from causal_testing.testing.estimators import LinearRegressionEstimator, CausalForestEstimator
 from causal_testing.testing.causal_test_outcome import NoEffect, Positive
 from tests.test_helpers import remove_temp_dir_if_existent
 from causal_testing.json_front.json_class import JsonUtility, CausalVariables
@@ -291,6 +291,31 @@ def test_no_data_provided(self):
         with self.assertRaises(ValueError):
             json_class.setup(self.scenario)
 
+    def test_estimator_formula_type_check(self):
+        example_test = {
+            "tests": [
+                {
+                    "name": "test1",
+                    "mutations": {"test_input": "Increase"},
+                    "estimator": "CausalForestEstimator",
+                    "estimate_type": "ate",
+                    "effect_modifiers": [],
+                    "expected_effect": {"test_output": "Positive"},
+                    "skip": False,
+                    "formula": "test_output ~ test_input",
+                }
+            ]
+        }
+        self.json_class.test_plan = example_test
+        effects = {"Positive": Positive()}
+        mutates = {
+            "Increase": lambda x: self.json_class.scenario.treatment_variables[x].z3
+                                  > self.json_class.scenario.variables[x].z3
+        }
+        estimators = {"CausalForestEstimator": CausalForestEstimator}
+        with self.assertRaises(TypeError):
+            self.json_class.run_json_tests(effects=effects, mutates=mutates, estimators=estimators, f_flag=False)
+
     def tearDown(self) -> None:
         remove_temp_dir_if_existent()
         if os.path.exists("temp_out.txt"):