test_azure_fileshare_to_gcs_2.py 1.3 KB

1234567891011121314151617181920212223242526272829303132333435
  1. def test_execute(self, gcs_mock_hook, azure_fileshare_mock_hook):
  2. """Test the execute function when the run is successful."""
  3. operator = AzureFileShareToGCSOperator(
  4. task_id=TASK_ID,
  5. share_name=AZURE_FILESHARE_SHARE,
  6. directory_name=AZURE_FILESHARE_DIRECTORY_NAME,
  7. azure_fileshare_conn_id=AZURE_FILESHARE_CONN_ID,
  8. gcp_conn_id=GCS_CONN_ID,
  9. dest_gcs=GCS_PATH_PREFIX,
  10. google_impersonation_chain=IMPERSONATION_CHAIN,
  11. )
  12. azure_fileshare_mock_hook.return_value.list_files.return_value = MOCK_FILES
  13. uploaded_files = operator.execute(None)
  14. gcs_mock_hook.return_value.upload.assert_has_calls(
  15. [
  16. mock.call('gcs-bucket', 'data/TEST1.csv', mock.ANY, gzip=False),
  17. mock.call('gcs-bucket', 'data/TEST3.csv', mock.ANY, gzip=False),
  18. mock.call('gcs-bucket', 'data/TEST2.csv', mock.ANY, gzip=False),
  19. ],
  20. any_order=True,
  21. )
  22. azure_fileshare_mock_hook.assert_called_once_with(AZURE_FILESHARE_CONN_ID)
  23. gcs_mock_hook.assert_called_once_with(
  24. gcp_conn_id=GCS_CONN_ID,
  25. delegate_to=None,
  26. impersonation_chain=IMPERSONATION_CHAIN,
  27. )
  28. assert sorted(MOCK_FILES) == sorted(uploaded_files)