def test_execute(self, mock_name, mock_hook):
mock_name.return_value = TRANSFER_CONFIG_ID
mock_xcom = mock.MagicMock()
op = BigQueryCreateDataTransferOperator(
transfer_config=TRANSFER_CONFIG, project_id=PROJECT_ID, task_id="id"
)
op.xcom_push = mock_xcom
op.execute(None)
mock_hook.assert_called_once_with(
authorization_code=None,
metadata=None,
transfer_config=TRANSFER_CONFIG,
project_id=PROJECT_ID,
retry=None,
timeout=None,
)
)
# [END howto_bigquery_dts_create_args]
default_args = {"start_date": days_ago(1)}
with models.DAG(
"example_gcp_bigquery_dts",
default_args=default_args,
schedule_interval=None, # Override to match your needs
tags=['example'],
) as dag:
# [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', "