test_get_formulae unittest and fix bug with covariates

Author: christopher-wild

causal_testing/testing/estimators.py

@@ -142,9 +142,12 @@ def get_terms_from_formula(self) -> tuple[str, str, list[str]]:
                 rhs_terms.add(term.factors[0].code)
         if self.treatment not in rhs_terms:
             raise ValueError(f"Treatment variable '{self.treatment}' not found in formula")
-        covariates = rhs_terms.remove(self.treatment)
+        rhs_terms.remove(self.treatment)
+        covariates = rhs_terms
         if covariates is None:
             covariates = []
+        else:
+            covariates = list(covariates)
         return outcome, self.treatment, covariates
 
     def validate_formula(self, causal_dag: CausalDAG):

tests/testing_tests/test_estimators.py

@@ -7,6 +7,7 @@
     CausalForestEstimator,
     LogisticRegressionEstimator,
     InstrumentalVariableEstimator,
+    RegressionEstimator,
 )
 from causal_testing.specification.variable import Input
 from causal_testing.utils.validation import CausalValidator
@@ -124,15 +125,15 @@ def test_ate_adjustment(self):
         logistic_regression_estimator = LogisticRegressionEstimator(
             "length_in", 65, 55, {"large_gauge"}, "completed", df
         )
-        ate, _ = logistic_regression_estimator.estimate_ate(adjustment_config = {"large_gauge": 0})
+        ate, _ = logistic_regression_estimator.estimate_ate(adjustment_config={"large_gauge": 0})
         self.assertEqual(round(ate, 4), -0.3388)
 
     def test_ate_invalid_adjustment(self):
         df = self.scarf_df.copy()
         logistic_regression_estimator = LogisticRegressionEstimator("length_in", 65, 55, {}, "completed", df)
         with self.assertRaises(ValueError):
             ate, _ = logistic_regression_estimator.estimate_ate(
-                adjustment_config = {"large_gauge": 0}
+                adjustment_config={"large_gauge": 0}
             )
 
     def test_ate_effect_modifiers(self):
@@ -394,7 +395,7 @@ def test_program_15_no_interaction_ate_calculated(self):
         # for term_to_square in terms_to_square:
 
         ate, [ci_low, ci_high] = linear_regression_estimator.estimate_ate_calculated(
-            adjustment_config = {k: self.nhefs_df.mean()[k] for k in covariates}
+            adjustment_config={k: self.nhefs_df.mean()[k] for k in covariates}
         )
         self.assertEqual(round(ate, 1), 3.5)
         self.assertEqual([round(ci_low, 1), round(ci_high, 1)], [1.9, 5])
@@ -491,3 +492,21 @@ def test_X1_effect(self):
         test_results = lr_model.estimate_ate()
         ate = test_results[0]
         self.assertAlmostEqual(ate, 2.0)
+
+
+class TestRegressionEstimator(unittest.TestCase):
+    """Test the extended functionality of the TestRegressionEstimator"""
+
+    @classmethod
+    def setUpClass(cls):
+        class RegressionEstimatorTesting(RegressionEstimator):
+            def add_modelling_assumptions(self):
+                pass
+
+        cls.regression_estimator = RegressionEstimatorTesting("X", 1, 0, {"Z"}, "Y", formula="Y ~ X + Z")
+
+    def test_get_formulae(self):
+        outcome, treatment, covariates = self.regression_estimator.get_terms_from_formula()
+        self.assertEqual(outcome, "Y")
+        self.assertEqual(treatment, "X")
+        self.assertEqual(covariates, ["Z"])