def test_passing_arguments_to_hook(self, mock_hook):
task = BigQueryTablePartitionExistenceSensor(
task_id='task-id',
project_id=TEST_PROJECT_ID,
dataset_id=TEST_DATASET_ID,
table_id=TEST_TABLE_ID,
partition_id=TEST_PARTITION_ID,
bigquery_conn_id=TEST_GCP_CONN_ID,
delegate_to=TEST_DELEGATE_TO,
impersonation_chain=TEST_IMPERSONATION_CHAIN,
)
mock_hook.return_value.table_partition_exists.return_value = True
results = task.poke(mock.MagicMock())
self.assertEqual(True, results)
mock_hook.assert_called_once_with(
bigquery_conn_id=TEST_GCP_CONN_ID,
delegate_to=TEST_DELEGATE_TO,
impersonation_chain=TEST_IMPERSONATION_CHAIN,
)
mock_hook.return_value.table_partition_exists.assert_called_once_with(
project_id=TEST_PROJECT_ID,
dataset_id=TEST_DATASET_ID,
table_id=TEST_TABLE_ID,
partition_id=TEST_PARTITION_ID,
)
task_id="check_table_exists",
project_id=PROJECT_ID,
dataset_id=DATASET_NAME,
table_id=TABLE_NAME)
# [END howto_sensor_bigquery_table]
execute_insert_query = BigQueryExecuteQueryOperator(
task_id="execute_insert_query",
sql=INSERT_ROWS_QUERY,
use_legacy_sql=False)
# [START howto_sensor_bigquery_table_partition]
check_table_partition_exists = BigQueryTablePartitionExistenceSensor(
task_id="check_table_partition_exists",
project_id=PROJECT_ID,
dataset_id=DATASET_NAME,
table_id=TABLE_NAME,
partition_id=PARTITION_NAME,
)
# [END howto_sensor_bigquery_table_partition]
delete_dataset = BigQueryDeleteDatasetOperator(task_id="delete_dataset",
dataset_id=DATASET_NAME,
delete_contents=True)
create_dataset >> create_table
create_table >> check_table_exists
create_table >> execute_insert_query
execute_insert_query >> check_table_partition_exists
check_table_exists >> delete_dataset
check_table_partition_exists >> delete_dataset