def test_execute(self, mock_xcom, mock_hook):
context = {}
spreadsheet = mock.MagicMock()
mock_hook.return_value.create_spreadsheet.return_value = {
"spreadsheetId": SPREADSHEET_ID,
"spreadsheetUrl": SPREADSHEET_URL,
}
op = GoogleSheetsCreateSpreadsheetOperator(
task_id="test_task", spreadsheet=spreadsheet, gcp_conn_id=GCP_CONN_ID
)
op.execute(context)
mock_hook.return_value.create_spreadsheet.assert_called_once_with(spreadsheet=spreadsheet)
calls = [
mock.call(context, "spreadsheet_id", SPREADSHEET_ID),
mock.call(context, "spreadsheet_url", SPREADSHEET_URL),
]
mock_xcom.has_calls(calls)
with models.DAG(
"example_sheets_gcs",
default_args=default_args,
schedule_interval=None, # Override to match your needs
tags=["example"],
) as dag:
# [START upload_sheet_to_gcs]
upload_sheet_to_gcs = GoogleSheetsToGCSOperator(
task_id="upload_sheet_to_gcs",
destination_bucket=GCS_BUCKET,
spreadsheet_id=SPREADSHEET_ID,
)
# [END upload_sheet_to_gcs]
# [START create_spreadsheet]
create_spreadsheet = GoogleSheetsCreateSpreadsheetOperator(
task_id="create_spreadsheet", spreadsheet=SPREADSHEET)
# [END create_spreadsheet]
# [START print_spreadsheet_url]
print_spreadsheet_url = BashOperator(
task_id="print_spreadsheet_url",
bash_command=
"echo {{ task_instance.xcom_pull('create_spreadsheet', key='spreadsheet_url') }}",
)
# [END print_spreadsheet_url]
# [START upload_gcs_to_sheet]
upload_gcs_to_sheet = GCSToGoogleSheetsOperator(
task_id="upload_gcs_to_sheet",
bucket_name=GCS_BUCKET,
object_name="{{ task_instance.xcom_pull('upload_sheet_to_gcs')[0] }}",