def test_execute(self, mock_hook):
op = BigQueryDataTransferServiceStartTransferRunsOperator(
transfer_config_id=TRANSFER_CONFIG_ID, task_id="id", project_id=PROJECT_ID
)
op.execute(None)
mock_hook.return_value.start_manual_transfer_runs.assert_called_once_with(
transfer_config_id=TRANSFER_CONFIG_ID,
project_id=PROJECT_ID,
requested_time_range=None,
requested_run_time=None,
metadata=None,
retry=None,
timeout=None,
)
def test_execute(self, mock_hook):
op = BigQueryDataTransferServiceStartTransferRunsOperator(
transfer_config_id=TRANSFER_CONFIG_ID, task_id="id", project_id=PROJECT_ID
)
ti = mock.MagicMock()
op.execute({'ti': ti})
mock_hook.return_value.start_manual_transfer_runs.assert_called_once_with(
transfer_config_id=TRANSFER_CONFIG_ID,
project_id=PROJECT_ID,
requested_time_range=None,
requested_run_time=None,
metadata=None,
retry=None,
timeout=None,
)
ti.xcom_push.assert_called_once_with(execution_date=None, key='run_id', value='123')
# [START howto_bigquery_create_data_transfer]
gcp_bigquery_create_transfer = BigQueryCreateDataTransferOperator(
transfer_config=TRANSFER_CONFIG,
project_id=GCP_PROJECT_ID,
task_id="gcp_bigquery_create_transfer",
)
transfer_config_id = (
"{{ task_instance.xcom_pull('gcp_bigquery_create_transfer', key='transfer_config_id') }}"
)
# [END howto_bigquery_create_data_transfer]
# [START howto_bigquery_start_transfer]
gcp_bigquery_start_transfer = BigQueryDataTransferServiceStartTransferRunsOperator(
task_id="gcp_bigquery_start_transfer",
transfer_config_id=transfer_config_id,
requested_run_time={"seconds": int(time.time() + 60)},
)
run_id = "{{ task_instance.xcom_pull('gcp_bigquery_start_transfer', key='run_id') }}"
# [END howto_bigquery_start_transfer]
# [START howto_bigquery_dts_sensor]
gcp_run_sensor = BigQueryDataTransferServiceTransferRunSensor(
task_id="gcp_run_sensor",
transfer_config_id=transfer_config_id,
run_id=run_id,
expected_statuses={"SUCCEEDED"},
)
# [END howto_bigquery_dts_sensor]
# [START howto_bigquery_delete_data_transfer]
