Move new JSON config test to BQ system tests file

GoogleCloudPlatform · Feb 7, 2024 · e4e2f4a · e4e2f4a
1 parent 1828691
commit e4e2f4a
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 54 deletions.
diff --git a/tests/system/data_sources/test_bigquery.py b/tests/system/data_sources/test_bigquery.py
@@ -403,6 +403,34 @@
 LIMIT 10
 """.strip()
 
+TEST_JSON_VALIDATION_CONFIG = {
+    consts.CONFIG_TYPE: "Column",
+    consts.CONFIG_SOURCE_CONN_NAME: "mock-conn",
+    consts.CONFIG_TARGET_CONN_NAME: "mock-conn",
+    consts.CONFIG_TABLE_NAME: "dvt_core_types",
+    consts.CONFIG_SCHEMA_NAME: "pso_data_validator",
+    consts.CONFIG_TARGET_SCHEMA_NAME: "pso_data_validator",
+    consts.CONFIG_TARGET_TABLE_NAME: "dvt_core_types",
+    consts.CONFIG_LABELS: [],
+    consts.CONFIG_THRESHOLD: 0.0,
+    consts.CONFIG_FORMAT: "table",
+    consts.CONFIG_RESULT_HANDLER: None,
+    consts.CONFIG_FILTERS: [],
+    consts.CONFIG_USE_RANDOM_ROWS: False,
+    consts.CONFIG_RANDOM_ROW_BATCH_SIZE: None,
+    consts.CONFIG_FILTER_STATUS: ["fail"],
+    consts.CONFIG_AGGREGATES: [
+        {
+            consts.CONFIG_SOURCE_COLUMN: None,
+            consts.CONFIG_TARGET_COLUMN: None,
+            consts.CONFIG_FIELD_ALIAS: "count",
+            consts.CONFIG_TYPE: "count",
+        },
+    ],
+    consts.CONFIG_SOURCE_CONN: BQ_CONN,
+    consts.CONFIG_TARGET_CONN: BQ_CONN,
+}
+
 
 def test_count_validator():
     validator = data_validation.DataValidation(CONFIG_COUNT_VALID, verbose=True)
@@ -1271,3 +1299,28 @@ def test_custom_query_validation_core_types(mock_conn):
     df = validator.execute()
     # With filter on failures the data frame should be empty
     assert len(df) == 0
+
+
+@mock.patch(
+    "data_validation.state_manager.StateManager.get_connection_config",
+    return_value=BQ_CONN,
+)
+def test_column_validation_convert_config_to_json(mock_conn):
+    parser = cli_tools.configure_arg_parser()
+    args = parser.parse_args(
+        [
+            "validate",
+            "column",
+            "-sc=mock-conn",
+            "-tc=mock-conn",
+            "-tbls=pso_data_validator.dvt_core_types",
+            "--filter-status=fail",
+            "--config-file-json=bq-column-validation.json",
+        ]
+    )
+    config_managers = main.build_config_managers_from_args(args)
+    assert len(config_managers) == 1
+
+    json_config = main.convert_config_to_json(config_managers)
+    # assert structure
+    assert json_config == TEST_JSON_VALIDATION_CONFIG
diff --git a/tests/unit/test__main.py b/tests/unit/test__main.py
@@ -21,7 +21,6 @@
 from data_validation import __main__ as main
 from tests.system.data_sources.test_bigquery import BQ_CONN
 
-PROJECT_ID = os.getenv("PROJECT_ID")
 
 TEST_CONN = '{"source_type":"Example"}'
 CLI_ARGS = {
@@ -89,34 +88,6 @@
     "config_dir": "gs://pso-kokoro-resources/resources/test/unit/test__main/4partitions",
 }
 
-TEST_JSON_VALIDATION_CONFIG = {
-    consts.CONFIG_TYPE: "Column",
-    consts.CONFIG_SOURCE_CONN_NAME: "mock-conn",
-    consts.CONFIG_TARGET_CONN_NAME: "mock-conn",
-    consts.CONFIG_TABLE_NAME: "dvt_core_types",
-    consts.CONFIG_SCHEMA_NAME: "pso_data_validator",
-    consts.CONFIG_TARGET_SCHEMA_NAME: "pso_data_validator",
-    consts.CONFIG_TARGET_TABLE_NAME: "dvt_core_types",
-    consts.CONFIG_LABELS: [],
-    consts.CONFIG_THRESHOLD: 0.0,
-    consts.CONFIG_FORMAT: "table",
-    consts.CONFIG_RESULT_HANDLER: None,
-    consts.CONFIG_FILTERS: [],
-    consts.CONFIG_USE_RANDOM_ROWS: False,
-    consts.CONFIG_RANDOM_ROW_BATCH_SIZE: None,
-    consts.CONFIG_FILTER_STATUS: ["fail"],
-    consts.CONFIG_AGGREGATES: [
-        {
-            consts.CONFIG_SOURCE_COLUMN: None,
-            consts.CONFIG_TARGET_COLUMN: None,
-            consts.CONFIG_FIELD_ALIAS: "count",
-            consts.CONFIG_TYPE: "count",
-        },
-    ],
-    consts.CONFIG_SOURCE_CONN: BQ_CONN,
-    consts.CONFIG_TARGET_CONN: BQ_CONN,
-}
-
 
 @mock.patch(
     "argparse.ArgumentParser.parse_args",
@@ -224,28 +195,3 @@ def test_config_runner_3(mock_args, mock_build, mock_run, caplog):
     assert mock_run.call_args.args[0].config_dir is None
     assert os.path.basename(mock_run.call_args.args[0].config_file) == "0002.yaml"
     assert len(mock_run.call_args.args[1]) == 1
-
-
-@mock.patch(
-    "data_validation.state_manager.StateManager.get_connection_config",
-    return_value=BQ_CONN,
-)
-def test_column_validation_convert_config_to_json(mock_conn):
-    parser = cli_tools.configure_arg_parser()
-    args = parser.parse_args(
-        [
-            "validate",
-            "column",
-            "-sc=mock-conn",
-            "-tc=mock-conn",
-            "-tbls=pso_data_validator.dvt_core_types",
-            "--filter-status=fail",
-            "--config-file-json=bq-column-validation.json",
-        ]
-    )
-    config_managers = main.build_config_managers_from_args(args)
-    assert len(config_managers) == 1
-
-    json_config = main.convert_config_to_json(config_managers)
-    # assert structure
-    assert json_config == TEST_JSON_VALIDATION_CONFIG