Add concrete test params to test json

Author: christopher-wild

causal_testing/json_front/json_class.py

@@ -104,48 +104,7 @@ def _create_abstract_test_case(self, test, mutates, effects):
             effect=test.get("effect", "total"),
         )
         return abstract_test
-    def run_coefficient_test(self, test, f_flag):
-        base_test_case = BaseTestCase(
-            treatment_variable=next(self.scenario.variables[v] for v in test["mutations"]),
-            outcome_variable=next(self.scenario.variables[v] for v in test["expected_effect"]),
-            effect=test.get("effect", "direct"),
-        )
-        assert len(test["expected_effect"]) == 1, "Can only have one expected effect."
-        causal_test_case = CausalTestCase(
-            base_test_case=base_test_case,
-            expected_causal_effect=next(
-                self.effects[effect] for variable, effect in test["expected_effect"].items()
-            ),
-            estimate_type="coefficient",
-            effect_modifier_configuration={
-                self.scenario.variables[v] for v in test.get("effect_modifiers", [])
-            },
-        )
-        result = self._execute_test_case(causal_test_case=causal_test_case, test=test, f_flag=f_flag)
-        msg = (
-                f"Executing test: {test['name']} \n"
-                + f"  {causal_test_case} \n"
-                + "  "
-                + ("\n  ").join(str(result[1]).split("\n"))
-                + "==============\n"
-                + f"  Result: {'FAILED' if result[0] else 'Passed'}"
-        )
-        return msg
-    def run_ate_test(self, test, f_flag):
-        abstract_test = self._create_abstract_test_case(test, self.mutates, self.effects)
-        concrete_tests, _ = abstract_test.generate_concrete_tests(sample_size=5, target_ks_score=0.05)
-        failures, _ = self._execute_tests(concrete_tests, test, f_flag)
 
-        msg = (
-                f"Executing test: {test['name']} \n"
-                + "  abstract_test \n"
-                + f"  {abstract_test} \n"
-                + f"  {abstract_test.treatment_variable.name},"
-                + f"  {abstract_test.treatment_variable.distribution} \n"
-                + f"  Number of concrete tests for test case: {str(len(concrete_tests))} \n"
-                + f"  {failures}/{len(concrete_tests)} failed for {test['name']}"
-        )
-        return msg
     def run_json_tests(self, effects: dict, estimators: dict, f_flag: bool = False, mutates: dict = None):
         """Runs and evaluates each test case specified in the JSON input
 
@@ -164,7 +123,7 @@ def run_json_tests(self, effects: dict, estimators: dict, f_flag: bool = False,
                 if test["estimate_type"] == "coefficient":
                     msg = self.run_coefficient_test(test=test, f_flag=f_flag)
                 else:
-                    msg =  self.run_ate_test(test=test, f_flag=f_flag)
+                    msg = self.run_ate_test(test=test, f_flag=f_flag)
                 self._append_to_file(msg, logging.INFO)
             else:
                 outcome_variable = next(
@@ -185,14 +144,66 @@ def run_json_tests(self, effects: dict, estimators: dict, f_flag: bool = False,
                 failed, _ = self._execute_test_case(causal_test_case=causal_test_case, test=test, f_flag=f_flag)
 
                 msg = (
-                    f"Executing concrete test: {test['name']} \n"
-                    + f"treatment variable: {test['treatment_variable']} \n"
-                    + f"outcome_variable = {outcome_variable} \n"
-                    + f"control value = {test['control_value']}, treatment value = {test['treatment_value']} \n"
-                    + f"Result: {'FAILED' if failed else 'Passed'}"
+                        f"Executing concrete test: {test['name']} \n"
+                        + f"treatment variable: {test['treatment_variable']} \n"
+                        + f"outcome_variable = {outcome_variable} \n"
+                        + f"control value = {test['control_value']}, treatment value = {test['treatment_value']} \n"
+                        + f"Result: {'FAILED' if failed else 'Passed'}"
                 )
                 self._append_to_file(msg, logging.INFO)
 
+    def run_coefficient_test(self, test, f_flag):
+        base_test_case = BaseTestCase(
+            treatment_variable=next(self.scenario.variables[v] for v in test["mutations"]),
+            outcome_variable=next(self.scenario.variables[v] for v in test["expected_effect"]),
+            effect=test.get("effect", "direct"),
+        )
+        assert len(test["expected_effect"]) == 1, "Can only have one expected effect."
+        causal_test_case = CausalTestCase(
+            base_test_case=base_test_case,
+            expected_causal_effect=next(
+                self.effects[effect] for variable, effect in test["expected_effect"].items()
+            ),
+            estimate_type="coefficient",
+            effect_modifier_configuration={
+                self.scenario.variables[v] for v in test.get("effect_modifiers", [])
+            },
+        )
+        result = self._execute_test_case(causal_test_case=causal_test_case, test=test, f_flag=f_flag)
+        msg = (
+                f"Executing test: {test['name']} \n"
+                + f"  {causal_test_case} \n"
+                + "  "
+                + ("\n  ").join(str(result[1]).split("\n"))
+                + "==============\n"
+                + f"  Result: {'FAILED' if result[0] else 'Passed'}"
+        )
+        return msg
+
+    def run_ate_test(self, test, f_flag):
+        if "sample_size" in test:
+            sample_size = test["sample_size"]
+        else:
+            sample_size = 5
+        if "target_ks_score" in test:
+            target_ks_score = test["target_ks_score"]
+        else:
+            target_ks_score = 0.05
+        abstract_test = self._create_abstract_test_case(test, self.mutates, self.effects)
+        concrete_tests, _ = abstract_test.generate_concrete_tests(sample_size=sample_size, target_ks_score=target_ks_score)
+        failures, _ = self._execute_tests(concrete_tests, test, f_flag)
+
+        msg = (
+                f"Executing test: {test['name']} \n"
+                + "  abstract_test \n"
+                + f"  {abstract_test} \n"
+                + f"  {abstract_test.treatment_variable.name},"
+                + f"  {abstract_test.treatment_variable.distribution} \n"
+                + f"  Number of concrete tests for test case: {str(len(concrete_tests))} \n"
+                + f"  {failures}/{len(concrete_tests)} failed for {test['name']}"
+        )
+        return msg
+
     def _execute_tests(self, concrete_tests, test, f_flag):
         failures = 0
         details = []
@@ -222,7 +233,7 @@ def _populate_metas(self):
             meta.populate(self.data)
 
     def _execute_test_case(
-        self, causal_test_case: CausalTestCase, test: Iterable[Mapping], f_flag: bool
+            self, causal_test_case: CausalTestCase, test: Iterable[Mapping], f_flag: bool
     ) -> (bool, CausalTestResult):
         """Executes a singular test case, prints the results and returns the test case result
         :param causal_test_case: The concrete test case to be executed
@@ -262,7 +273,7 @@ def _execute_test_case(
         return failed, causal_test_result
 
     def _setup_test(
-        self, causal_test_case: CausalTestCase, test: Mapping, conditions: list[str] = None
+            self, causal_test_case: CausalTestCase, test: Mapping, conditions: list[str] = None
     ) -> tuple[CausalTestEngine, Estimator]:
         """Create the necessary inputs for a single test case
         :param causal_test_case: The concrete test case to be executed
@@ -347,7 +358,7 @@ def get_args(test_args=None) -> argparse.Namespace:
         parser.add_argument(
             "-w",
             help="Specify to overwrite any existing output files. This can lead to the loss of existing outputs if not "
-            "careful",
+                 "careful",
             action="store_true",
         )
         parser.add_argument(