def test_execute(self, gcs_mock_hook, adls_one_mock_hook, adls_two_mock_hook): """Test the execute function when the run is successful.""" operator = AdlsToGoogleCloudStorageOperator( task_id=TASK_ID, src_adls=ADLS_PATH_1, dest_gcs=GCS_PATH, replace=False, azure_data_lake_conn_id=AZURE_CONN_ID, google_cloud_storage_conn_id=GCS_CONN_ID) adls_one_mock_hook.return_value.list.return_value = MOCK_FILES adls_two_mock_hook.return_value.list.return_value = MOCK_FILES # gcs_mock_hook.return_value.upload.side_effect = _assert_upload uploaded_files = operator.execute(None) gcs_mock_hook.return_value.upload.assert_has_calls([ mock.call(bucket_name='test', filename=mock.ANY, object_name='test/path/PARQUET.parquet', gzip=False), mock.call(bucket_name='test', filename=mock.ANY, object_name='test/path/TEST3.csv', gzip=False), mock.call(bucket_name='test', filename=mock.ANY, object_name='test/path/PIC.png', gzip=False), mock.call(bucket_name='test', filename=mock.ANY, object_name='test/TEST1.csv', gzip=False), mock.call(bucket_name='test', filename=mock.ANY, object_name='test/TEST2.csv', gzip=False) ], any_order=True) adls_one_mock_hook.assert_called_once_with( azure_data_lake_conn_id=AZURE_CONN_ID) adls_two_mock_hook.assert_called_once_with( azure_data_lake_conn_id=AZURE_CONN_ID) gcs_mock_hook.assert_called_once_with( google_cloud_storage_conn_id=GCS_CONN_ID, delegate_to=None) # we expect MOCK_FILES to be uploaded self.assertEqual(sorted(MOCK_FILES), sorted(uploaded_files))
def test_init(self): """Test AdlsToGoogleCloudStorageOperator instance is properly initialized.""" operator = AdlsToGoogleCloudStorageOperator( task_id=TASK_ID, src_adls=ADLS_PATH_1, dest_gcs=GCS_PATH, replace=False, azure_data_lake_conn_id=AZURE_CONN_ID, gcp_conn_id=GCS_CONN_ID) self.assertEqual(operator.task_id, TASK_ID) self.assertEqual(operator.src_adls, ADLS_PATH_1) self.assertEqual(operator.dest_gcs, GCS_PATH) self.assertEqual(operator.replace, False) self.assertEqual(operator.gcp_conn_id, GCS_CONN_ID) self.assertEqual(operator.azure_data_lake_conn_id, AZURE_CONN_ID)