def test_execute(self, mock_upload_data, mock_xcom, mock_sheet_hook, mock_gcs_hook):
context = {}
data = ["data1", "data2"]
mock_sheet_hook.return_value.get_sheet_titles.return_value = RANGES
mock_upload_data.side_effect = [PATH, PATH]
op = GoogleSheetsToGCSOperator(
task_id="test_task",
spreadsheet_id=SPREADSHEET_ID,
destination_bucket=BUCKET,
sheet_filter=FILTER,
destination_path=PATH,
gcp_conn_id=GCP_CONN_ID,
impersonation_chain=IMPERSONATION_CHAIN,
)
op.execute(context)
mock_sheet_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID, delegate_to=None, impersonation_chain=IMPERSONATION_CHAIN,
)
mock_gcs_hook.assert_called_once_with(
gcp_conn_id=GCP_CONN_ID, delegate_to=None, impersonation_chain=IMPERSONATION_CHAIN,
)
mock_sheet_hook.return_value.get_sheet_titles.assert_called_once_with(
spreadsheet_id=SPREADSHEET_ID, sheet_filter=FILTER
)
calls = [mock.call(spreadsheet_id=SPREADSHEET_ID, range_=r) for r in RANGES]
mock_sheet_hook.return_value.get_values.has_calls(calls)
calls = [mock.call(mock_gcs_hook, mock_sheet_hook, r, v) for r, v in zip(RANGES, data)]
mock_upload_data.has_calls(calls)
mock_xcom.assert_called_once_with(context, "destination_objects", [PATH, PATH])
def test_upload_data(self, mock_tempfile, mock_writer):
filename = "file://97g23r"
file_handle = mock.MagicMock()
mock_tempfile.return_value.__enter__.return_value = file_handle
mock_tempfile.return_value.__enter__.return_value.name = filename
mock_sheet_hook = mock.MagicMock()
mock_sheet_hook.get_spreadsheet.return_value = {
"properties": {
"title": SHEET_TITLE
}
}
expected_dest_file = f"{PATH}/{SHEET_TITLE}_{RANGE}.csv"
mock_gcs_hook = mock.MagicMock()
op = GoogleSheetsToGCSOperator(
task_id="test_task",
spreadsheet_id=SPREADSHEET_ID,
destination_bucket=BUCKET,
sheet_filter=FILTER,
destination_path=PATH,
)
result = op._upload_data(
gcs_hook=mock_gcs_hook,
hook=mock_sheet_hook,
sheet_range=RANGE,
sheet_values=VALUES,
)
# Test writing to file
mock_sheet_hook.get_spreadsheet.assert_called_once_with(SPREADSHEET_ID)
mock_writer.assert_called_once_with(file_handle)
mock_writer.return_value.writerows.assert_called_once_with(VALUES)
file_handle.flush.assert_called_once_with()
# Test upload
mock_gcs_hook.upload.assert_called_once_with(
bucket_name=BUCKET,
object_name=expected_dest_file,
filename=filename)
# Assert path to file is returned
assert result == expected_dest_file
from airflow.providers.google.suite.transfers.gcs_to_sheets import GCSToGoogleSheetsOperator
BUCKET = os.environ.get("GCP_GCS_BUCKET", "example-test-bucket3")
SPREADSHEET_ID = os.environ.get("SPREADSHEET_ID", "example-spreadsheetID")
NEW_SPREADSHEET_ID = os.environ.get("NEW_SPREADSHEET_ID", "1234567890qwerty")
with models.DAG(
"example_gcs_to_sheets",
start_date=datetime(2021, 1, 1),
schedule_interval='@once', # Override to match your needs
catchup=False,
tags=["example"],
) as dag:
upload_sheet_to_gcs = GoogleSheetsToGCSOperator(
task_id="upload_sheet_to_gcs",
destination_bucket=BUCKET,
spreadsheet_id=SPREADSHEET_ID,
)
# [START upload_gcs_to_sheets]
upload_gcs_to_sheet = GCSToGoogleSheetsOperator(
task_id="upload_gcs_to_sheet",
bucket_name=BUCKET,
object_name="{{ task_instance.xcom_pull('upload_sheet_to_gcs')[0] }}",
spreadsheet_id=NEW_SPREADSHEET_ID,
)
# [END upload_gcs_to_sheets]
upload_sheet_to_gcs >> upload_gcs_to_sheet
