fix errors in integration tests

ihmeuw-msca · Dec 27, 2024 · 274b68d · 274b68d
1 parent d52f94a
commit 274b68d
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 40 deletions.
diff --git a/tests/integration/test_integration_stage_io.py b/tests/integration/test_integration_stage_io.py
@@ -125,16 +125,7 @@ def test_stage_model(stage_1, stage_2):
     stage_1_model_expected = {
         "name": "stage_1",
         "type": "DummyStage",
-        "config": {
-            "coef_bounds": None,
-            "holdout_columns": None,
-            "id_columns": None,
-            "model_type": None,
-            "observation_column": None,
-            "prediction_column": None,
-            "test_column": None,
-            "weights_column": None,
-        },
+        "config": {},
         "input_validation": {},
         "output_validation": {},
         "module": Path(__file__),
@@ -151,16 +142,7 @@ def test_stage_model(stage_1, stage_2):
     stage_2_model_expected = {
         "name": "stage_2",
         "type": "DummyStage",
-        "config": {
-            "coef_bounds": None,
-            "holdout_columns": None,
-            "id_columns": None,
-            "model_type": None,
-            "observation_column": None,
-            "prediction_column": None,
-            "test_column": None,
-            "weights_column": None,
-        },
+        "config": {},
         "input_validation": {},
         "output_validation": {},
         "module": Path(__file__),

diff --git a/tests/integration/test_integration_stage_io_validation.py b/tests/integration/test_integration_stage_io_validation.py
@@ -96,16 +96,7 @@ def stage_1_model_expected(test_base_dir):
         "name": "stage_1",
         "type": "DummyStage",
         "module": Path(__file__),
-        "config": {
-            "coef_bounds": None,
-            "holdout_columns": None,
-            "id_columns": None,
-            "model_type": None,
-            "observation_column": None,
-            "prediction_column": None,
-            "test_column": None,
-            "weights_column": None,
-        },
+        "config": {},
         "input_validation": {
             "covariates": {
                 "stage": "stage_0",
@@ -224,16 +215,7 @@ def stage_2_model_expected(test_base_dir):
     return {
         "name": "stage_2",
         "type": "DummyStage",
-        "config": {
-            "coef_bounds": None,
-            "holdout_columns": None,
-            "id_columns": None,
-            "model_type": None,
-            "observation_column": None,
-            "prediction_column": None,
-            "test_column": None,
-            "weights_column": None,
-        },
+        "config": {},
         "module": Path(__file__),
         "input_validation": {
             "data": {