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