Refactor test_suite to be a UserDict class

Author: christopher-wild

causal_testing/testing/causal_test_engine.py

@@ -63,8 +63,7 @@ def execute_test_suite(self, test_suite: CausalTestSuite) -> list[CausalTestResu
         if self.scenario_execution_data_df.empty:
             raise Exception("No data has been loaded. Please call load_data prior to executing a causal test case.")
         test_suite_results = {}
-        test_suite_dict = test_suite.test_suite
-        for edge in test_suite_dict:
+        for edge in test_suite:
             print("edge: ")
             print(edge)
             logger.info("treatment: %s", edge.treatment_variable)
@@ -81,9 +80,9 @@ def execute_test_suite(self, test_suite: CausalTestSuite) -> list[CausalTestResu
                 # TODO: When we implement causal contracts, we should also note the positivity violation there
                 raise Exception("POSITIVITY VIOLATION -- Cannot proceed.")
 
-            estimators = test_suite_dict[edge]["estimators"]
-            tests = test_suite_dict[edge]["tests"]
-            estimate_type = test_suite_dict[edge]["estimate_type"]
+            estimators = test_suite[edge]["estimators"]
+            tests = test_suite[edge]["tests"]
+            estimate_type = test_suite[edge]["estimate_type"]
             results = {}
             for EstimatorClass in estimators:
                 causal_test_results = []

causal_testing/testing/causal_test_suite.py

@@ -1,4 +1,7 @@
-class CausalTestSuite:
+from collections import UserDict
+
+
+class CausalTestSuite(UserDict):
     """
     A CausalTestSuite is a class holding a dictionary, where the keys are base_test_cases representing the treatment
     and outcome Variables, and the values are test objects. Test Objects hold a causal_test_case_list which is a
@@ -8,14 +11,6 @@ class CausalTestSuite:
     base_test_case's and execute each causal_test_case with each iterator.
     """
 
-    def __init__(
-        self,
-    ):
-        self.test_suite = {}
-
     def add_test_object(self, base_test_case, causal_test_case_list, estimators, estimate_type: str = "ate"):
         test_object = {"tests": causal_test_case_list, "estimators": estimators, "estimate_type": estimate_type}
-        self.test_suite[base_test_case] = test_object
-
-    def get_single_test_object(self, base_test_case):
-        return self.test_suite[base_test_case]
+        self.data[base_test_case] = test_object

tests/testing_tests/test_causal_test_suite.py

@@ -71,7 +71,7 @@ def test_adding_test_object(self):
                                    estimators=estimators)
         manual_test_object = {
             self.base_test_case: {"tests": test_list, "estimators": estimators, "estimate_type": "ate"}}
-        self.assertEqual(test_suite.test_suite, manual_test_object)
+        self.assertEqual(test_suite, manual_test_object)
 
     def test_return_single_test_object(self):
         base_test_case = BaseTestCase(self.A, self.D)
@@ -87,7 +87,7 @@ def test_return_single_test_object(self):
 
         manual_test_case = {"tests": test_list, "estimators": estimators, "estimate_type": "ate"}
 
-        test_case = self.test_suite.get_single_test_object(base_test_case)
+        test_case = self.test_suite[base_test_case]
 
         self.assertEqual(test_case, manual_test_case)
 
@@ -103,8 +103,10 @@ def test_execute_test_suite_multiple_base_test_cases(self):
         pass
 
     def test_execute_test_suite_multiple_estimators(self):
-        pass
-
+        #estimators = [LinearRegressionEstimator, CausalForestEstimator]
+        #test_suite_2_estimators = CausalTestSuite
+        #test_suite_2_estimators
+        causal_test_engine = self.create_causal_test_engine()
     def create_causal_test_engine(self):
 
         causal_specification = CausalSpecification(self.scenario, self.causal_dag)