Skip to content

Commit

Permalink
fix: test
Browse files Browse the repository at this point in the history
  • Loading branch information
@deryrahman
deryrahman committed Jul 11, 2024
1 parent cf4f246 commit 0c44beb
Showing 1 changed file with 16 additions and 4 deletions.
20 changes: 16 additions & 4 deletions task/bq2bq/executor/tests/test_transformation.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,7 @@ def test_partition_transform_execute(self, BigqueryServiceMock):
bigquery_service.transform_load.assert_called_with(query=final_query,
write_disposition=WriteDisposition.WRITE_TRUNCATE,
destination_table="bq_project.playground_dev.abcd$20190101",
dry_run=False,
allow_field_addition=False)

@mock.patch("bumblebee.bigquery_service.BigqueryService")
Expand Down Expand Up @@ -90,6 +91,7 @@ def test_table_transform(self, BigqueryServiceMock):
bigquery_service.transform_load.assert_called_with(query=final_query,
write_disposition=WriteDisposition.WRITE_TRUNCATE,
destination_table="bq_project.playground_dev.abcd",
dry_run=False,
allow_field_addition=False)

@mock.patch("bumblebee.bigquery_service.BigqueryService")
Expand All @@ -113,6 +115,7 @@ def test_single_partition_transform_1d_window_0_offset_without_spillover(self, B
bigquery_service.transform_load.assert_called_with(query=final_query,
write_disposition=WriteDisposition.WRITE_TRUNCATE,
destination_table="bq_project.playground_dev.abcd$20190101",
dry_run=False,
allow_field_addition=False)

@mock.patch("bumblebee.bigquery_service.BigqueryService")
Expand All @@ -135,6 +138,7 @@ def test_single_partition_transform_2d_window_24h_offset_without_spillover(self,
bigquery_service.transform_load.assert_called_with(query=final_query,
write_disposition=WriteDisposition.WRITE_TRUNCATE,
destination_table="bq_project.playground_dev.abcd$20190104",
dry_run=False,
allow_field_addition=False)

@mock.patch("bumblebee.bigquery_service.BigqueryService")
Expand Down Expand Up @@ -174,6 +178,7 @@ def test_single_partition_transform_7d_window_without_spillover(self, BigquerySe
bigquery_service.transform_load.assert_called_with(query=final_query,
write_disposition=WriteDisposition.WRITE_TRUNCATE,
destination_table="bq_project.playground_dev.abcd$20190103",
dry_run=False,
allow_field_addition=False)

@mock.patch("bumblebee.bigquery_service.BigqueryService")
Expand Down Expand Up @@ -220,7 +225,7 @@ def test_dml_transform(self, BigqueryServiceMock):
task.execute()

final_query = """select count(1) from table where date >= '2019-01-02' and date < '2019-01-03'"""
bigquery_service.execute_query.assert_called_with(final_query)
bigquery_service.execute_query.assert_called_with(final_query,dry_run=False)

@mock.patch("bumblebee.bigquery_service.BigqueryService")
def test_execute_dry_run(self, BigqueryServiceMock):
Expand All @@ -237,7 +242,11 @@ def test_execute_dry_run(self, BigqueryServiceMock):
task = TableTransformation(bigquery_service, task_config, query, localized_start_time,
localized_end_time, dry_run, localized_execution_time)
task.transform()
bigquery_service.transform_load.assert_not_called()
bigquery_service.transform_load.assert_called_with(query=query,
write_disposition=WriteDisposition.WRITE_TRUNCATE,
destination_table="bq_project.playground_dev.abcd$20190101",
dry_run=True,
allow_field_addition=False)


@mock.patch("bumblebee.bigquery_service.BigqueryService")
Expand All @@ -262,6 +271,7 @@ def test_allow_field_addition(self, BigqueryServiceMock):
bigquery_service.transform_load.assert_called_with(query=final_query,
write_disposition=WriteDisposition.WRITE_TRUNCATE,
destination_table="bq_project.playground_dev.abcd",
dry_run=False,
allow_field_addition=True)


Expand Down Expand Up @@ -289,7 +299,7 @@ def test_should_run_dml_merge_statements(self, BigqueryServiceMock):
transformation.transform()

final_query = """select count(1) from table where date >= '2019-02-01' and date < '2019-02-02'"""
bigquery_service.execute_query.assert_called_with(final_query)
bigquery_service.execute_query.assert_called_with(final_query,dry_run=False)

@mock.patch("bumblebee.bigquery_service.BigqueryService")
def test_should_run_table_task(self, BigqueryServiceMock):
Expand Down Expand Up @@ -325,6 +335,7 @@ def get_table_mock(table_name):
bigquery_service.transform_load.assert_called_with(query=final_query,
write_disposition=WriteDisposition.WRITE_APPEND,
destination_table="bq_project.playground_dev.abcd",
dry_run=False,
allow_field_addition=False)

@mock.patch("bumblebee.bigquery_service.BigqueryService")
Expand Down Expand Up @@ -362,6 +373,7 @@ def get_table_mock(table_name):
bigquery_service.transform_load.assert_called_with(query=final_query,
write_disposition=WriteDisposition.WRITE_TRUNCATE,
destination_table="bq_project.playground_dev.abcd",
dry_run=False,
allow_field_addition=False)

@mock.patch("bumblebee.bigquery_service.BigqueryService")
Expand Down Expand Up @@ -401,7 +413,7 @@ def get_table_mock(table_name):
transformation.transform()

final_query = """-- Optimus generated\nDECLARE partitions ARRAY<DATE>;\n\n\n\nCREATE TEMP TABLE `opt__partitions` AS (\n select count(1) from table where date >= '__dstart__' and date < '__dend__'\n);\n\nSET (partitions) = (\n SELECT AS STRUCT\n array_agg(DISTINCT DATE(`event_timestamp`))\n FROM opt__partitions\n);\n\nMERGE INTO\n `bq_project.playground_dev.abcd` AS target\nUSING\n (\n Select * from `opt__partitions`\n ) AS source\nON FALSE\nWHEN NOT MATCHED BY SOURCE AND DATE(`event_timestamp`) IN UNNEST(partitions)\nTHEN DELETE\nWHEN NOT MATCHED THEN INSERT\n (\n \n )\nVALUES\n (\n \n );\n"""
bigquery_service.execute_query.assert_called_with(final_query)
bigquery_service.execute_query.assert_called_with(final_query,dry_run=False)

@mock.patch("bumblebee.bigquery_service.BigqueryService")
def test_should_fail_if_partition_task_for_ingestion_time_without_filter_in_REPLACE_MERGE(self, BigqueryServiceMock):
Expand Down

0 comments on commit 0c44beb

Please sign in to comment.