Add tests for top-level writing to BQ helper

Author: mashalifshin

jobs/ads-incrementality-dap-collector/ads_incrementality_dap_collector/constants.py

@@ -7,7 +7,7 @@
 PROCESS_TIMEOUT = 1200  # 20 mins
 
 CONFIG_FILE_NAME = "config.json"     # See example_config.json for the contents and structure of the job config file.
-LOG_FILE_NAME = f"ads-incrementality-dap-collector-{datetime.now()}.log"
+LOG_FILE_NAME = f"{datetime.now()}-ads-incrementality-dap-collector.log"
 
 COLLECTOR_RESULTS_SCHEMA = [
     bigquery.SchemaField("collection_start", "DATE", mode="REQUIRED", description="Start date of the collected time window, inclusive."),

jobs/ads-incrementality-dap-collector/tests/test_helpers.py

@@ -4,7 +4,7 @@
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../ads_incrementality_dap_collector')))
 
 import pytest
-
+import re
 from unittest import TestCase
 from unittest.mock import patch
 
@@ -13,9 +13,14 @@
     mock_nimbus_experiment, mock_control_row, mock_treatment_a_row, mock_treatment_b_row,
     mock_task_id, mock_nimbus_unparseable_experiment,
     mock_tasks_to_collect, mock_dap_config, mock_experiment_config,
-    mock_dap_subprocess_success, mock_dap_subprocess_fail, mock_dap_subprocess_raise
+    mock_dap_subprocess_success, mock_dap_subprocess_fail, mock_dap_subprocess_raise,
+    mock_collected_tasks, mock_bq_config,
+    mock_create_dataset_success, mock_create_table_success, mock_insert_rows_json_success,
+    mock_create_dataset_fail, mock_create_table_fail, mock_insert_rows_json_fail
+)
+from ads_incrementality_dap_collector.helpers import (
+    get_experiment, prepare_results_rows, collect_dap_results, write_results_to_bq
 )
-from ads_incrementality_dap_collector.helpers import get_experiment, prepare_results_rows, collect_dap_results
 
 class TestHelpers(TestCase):
     @patch("requests.get", side_effect=mock_nimbus_success)
@@ -69,3 +74,43 @@ def test_collect_dap_results_raise(self, mock_dap_subprocess_raise):
         with pytest.raises(Exception, match=f'Collection failed for {task_id}, 1, stderr: Uh-oh'):
             collect_dap_results(tasks_to_collect, mock_dap_config(), mock_experiment_config())
             self.assertEqual(1, mock_dap_subprocess_success.call_count)
+
+    @patch("google.cloud.bigquery.Client.create_dataset", side_effect=mock_create_dataset_success)
+    @patch("google.cloud.bigquery.Client.create_table", side_effect=mock_create_table_success)
+    @patch("google.cloud.bigquery.Client.insert_rows_json", side_effect=mock_insert_rows_json_success)
+    def test_write_results_to_bq_success(self, mock_insert_rows_json_success, mock_create_table_success, mock_create_dataset_success):
+        collected_tasks = mock_collected_tasks()
+        write_results_to_bq(collected_tasks, mock_bq_config())
+        self.assertEqual(1, mock_create_dataset_success.call_count)
+        self.assertEqual(1, mock_create_table_success.call_count)
+        self.assertEqual(len(collected_tasks["mubArkO3So8Co1X98CBo62-lSCM4tB-NZPOUGJ83N1o"]), mock_insert_rows_json_success.call_count)
+
+    @patch("google.cloud.bigquery.Client.create_dataset", side_effect=mock_create_dataset_fail)
+    @patch("google.cloud.bigquery.Client.create_table", side_effect=mock_create_table_success)
+    @patch("google.cloud.bigquery.Client.insert_rows_json", side_effect=mock_insert_rows_json_success)
+    def test_write_results_to_bq_create_dataset_fail(self, mock_insert_rows_json_success, mock_create_table_success, mock_create_dataset_fail):
+        with pytest.raises(Exception, match='BQ create dataset Uh-oh'):
+            write_results_to_bq(mock_collected_tasks(), mock_bq_config())
+            self.assertEqual(1, mock_create_dataset_fail.call_count)
+            self.assertEqual(0, mock_create_table_success.call_count)
+            self.assertEqual(0, mock_insert_rows_json_success.call_count)
+
+    @patch("google.cloud.bigquery.Client.create_dataset", side_effect=mock_create_dataset_success)
+    @patch("google.cloud.bigquery.Client.create_table", side_effect=mock_create_table_fail)
+    @patch("google.cloud.bigquery.Client.insert_rows_json", side_effect=mock_insert_rows_json_success)
+    def test_write_results_to_bq_create_table_fail(self, mock_insert_rows_json_success, mock_create_table_fail, mock_create_dataset_success):
+        with pytest.raises(Exception, match='Failed to create BQ table: some-gcp-project-id.ads_dap.incrementality'):
+            write_results_to_bq(mock_collected_tasks(), mock_bq_config())
+            self.assertEqual(1, mock_create_dataset_success.call_count)
+            self.assertEqual(1, mock_create_dataset_fail.call_count)
+            self.assertEqual(0, mock_insert_rows_json_success.call_count)
+
+    @patch("google.cloud.bigquery.Client.create_dataset", side_effect=mock_create_dataset_success)
+    @patch("google.cloud.bigquery.Client.create_table", side_effect=mock_create_table_success)
+    @patch("google.cloud.bigquery.Client.insert_rows_json", side_effect=mock_insert_rows_json_fail)
+    def test_write_results_to_bq_insert_rows_fail(self, mock_insert_rows_json_fail, mock_create_table_success, mock_create_dataset_success):
+        with pytest.raises(Exception, match=re.escape("Error inserting rows into some-gcp-project-id.ads_dap.incrementality: [{'key': 0, 'errors': 'Problem writing bucket 1 results'}, {'key': 1, 'errors': 'Problem writing bucket 2 results'}, {'key': 2, 'errors': 'Problem writing bucket 3 results'}]")):
+            write_results_to_bq(mock_collected_tasks(), mock_bq_config())
+            self.assertEqual(1, mock_create_dataset_success.call_count)
+            self.assertEqual(1, mock_create_table_success.call_count)
+            self.assertEqual(1, mock_insert_rows_json_fail.call_count)

jobs/ads-incrementality-dap-collector/tests/test_mocks.py

@@ -1,5 +1,7 @@
+from google.cloud import bigquery
+from collections.abc import Mapping, Sequence
 from subprocess import CompletedProcess
-from models import IncrementalityBranchResultsRow, NimbusExperiment, DAPConfig, ExperimentConfig
+from models import IncrementalityBranchResultsRow, NimbusExperiment, BQConfig, DAPConfig, ExperimentConfig
 from tests.test_mock_responses import NIMBUS_SUCCESS, NIMBUS_NOT_AN_INCREMENTALITY_EXPERIMENT
 
 class MockResponse:
@@ -87,6 +89,20 @@ def mock_tasks_to_collect() -> dict[str, dict[int, IncrementalityBranchResultsRo
         }
     }
 
+def mock_collected_tasks() -> dict[str, dict[int, IncrementalityBranchResultsRow]]:
+    experiment = mock_nimbus_experiment()
+    tasks_to_collect = {
+        "mubArkO3So8Co1X98CBo62-lSCM4tB-NZPOUGJ83N1o": {
+            1: mock_control_row(experiment),
+            2: mock_treatment_b_row(experiment),
+            3: mock_treatment_a_row(experiment)
+        }
+    }
+    tasks_to_collect["mubArkO3So8Co1X98CBo62-lSCM4tB-NZPOUGJ83N1o"][1].value_count = 13645
+    tasks_to_collect["mubArkO3So8Co1X98CBo62-lSCM4tB-NZPOUGJ83N1o"][2].value_count = 18645
+    tasks_to_collect["mubArkO3So8Co1X98CBo62-lSCM4tB-NZPOUGJ83N1o"][3].value_count = 9645
+    return tasks_to_collect
+
 def mock_dap_config() -> DAPConfig:
     return DAPConfig(
         hpke_config="AQAgAAEAAQAgpdceoGiuWvIiogA8SPCdprkhWMNtLq_y0GSePI7EhXE",
@@ -101,6 +117,13 @@ def mock_experiment_config() -> ExperimentConfig:
         batch_duration=604800
     )
 
+def mock_bq_config() -> BQConfig:
+    return BQConfig(
+        project="some-gcp-project-id",
+        namespace="ads_dap",
+        table="incrementality"
+    )
+
 def mock_dap_subprocess_success(args: list[str], capture_output: bool, text: bool, check: bool, timeout: int) -> CompletedProcess:
     return CompletedProcess(args=[
         './collect',
@@ -137,3 +160,25 @@ def mock_dap_subprocess_fail(args: list[str], capture_output: bool, text: bool,
 
 def mock_dap_subprocess_raise(args: list[str], capture_output: bool, text: bool, check: bool, timeout: int) -> CompletedProcess:
     raise Exception("Collection failed for mubArkO3So8Co1X98CBo62-lSCM4tB-NZPOUGJ83N1o, 1, stderr: Uh-oh") from None
+
+def mock_create_dataset_success(data_set: str, exists_ok: bool):
+    pass
+
+def mock_create_dataset_fail(data_set: str, exists_ok: bool):
+    raise Exception("BQ create dataset Uh-oh")
+
+def mock_create_table_success(table: bigquery.Table, exists_ok: bool):
+    pass
+
+def mock_create_table_fail(table: bigquery.Table, exists_ok: bool):
+    raise Exception("BQ create table Uh-oh")
+
+def mock_insert_rows_json_success(table: str, json_rows: dict) -> Sequence[Mapping]:
+    return []
+
+def mock_insert_rows_json_fail(table: str, json_rows: dict) -> Sequence[Mapping]:
+    return [
+        {"key": 0, "errors": "Problem writing bucket 1 results" },
+        {"key": 1, "errors": "Problem writing bucket 2 results"},
+        {"key": 2, "errors": "Problem writing bucket 3 results"}
+    ]