def test_execute_timeout(self, mock_hook):
task = GCSObjectsWtihPrefixExistenceSensor(
task_id="task-id", bucket=TEST_BUCKET, prefix=TEST_PREFIX, poke_interval=0, timeout=1
)
mock_hook.return_value.list.return_value = []
with pytest.raises(AirflowSensorTimeout):
task.execute(mock.MagicMock)
mock_hook.return_value.list.assert_called_once_with(TEST_BUCKET, prefix=TEST_PREFIX)
def test_should_return_false_on_empty_list(self, mock_hook):
task = GCSObjectsWtihPrefixExistenceSensor(
task_id="task-id",
bucket=TEST_BUCKET,
prefix=TEST_PREFIX,
google_cloud_conn_id=TEST_GCP_CONN_ID,
delegate_to=TEST_DELEGATE_TO,
)
mock_hook.return_value.list.return_value = []
result = task.poke(mock.MagicMock)
self.assertEqual(False, result)
def test_should_pass_arguments_to_hook(self, mock_hook):
task = GCSObjectsWtihPrefixExistenceSensor(
task_id="task-id",
bucket=TEST_BUCKET,
prefix=TEST_PREFIX,
google_cloud_conn_id=TEST_GCP_CONN_ID,
delegate_to=TEST_DELEGATE_TO,
)
mock_hook.return_value.list.return_value = ["NOT_EMPTY_LIST"]
result = task.poke(mock.MagicMock)
mock_hook.assert_called_once_with(
delegate_to=TEST_DELEGATE_TO,
google_cloud_storage_conn_id=TEST_GCP_CONN_ID
)
mock_hook.return_value.list.assert_called_once_with(TEST_BUCKET, prefix=TEST_PREFIX)
self.assertEqual(True, result)
def create_watch_sensor(dag, task_id, url_prefix, **kwargs):
parsed_url = parse_gs_url(url_prefix)
return GCSObjectsWtihPrefixExistenceSensor(
task_id=task_id,
bucket=parsed_url['bucket'],
prefix=parsed_url['object'],
dag=dag,
**kwargs
)
def test_execute(self, mock_hook):
task = GCSObjectsWtihPrefixExistenceSensor(
task_id="task-id",
bucket=TEST_BUCKET,
prefix=TEST_PREFIX,
google_cloud_conn_id=TEST_GCP_CONN_ID,
delegate_to=TEST_DELEGATE_TO,
poke_interval=0)
generated_messages = ['test-prefix/obj%s' % i for i in range(5)]
mock_hook.return_value.list.return_value = generated_messages
response = task.execute(None)
mock_hook.assert_called_once_with(
delegate_to=TEST_DELEGATE_TO,
google_cloud_storage_conn_id=TEST_GCP_CONN_ID
)
mock_hook.return_value.list.assert_called_once_with(TEST_BUCKET, prefix=TEST_PREFIX)
self.assertEqual(response, generated_messages)
def test_should_pass_arguments_to_hook(self, mock_hook):
task = GCSObjectsWtihPrefixExistenceSensor(
task_id="task-id",
bucket=TEST_BUCKET,
prefix=TEST_PREFIX,
google_cloud_conn_id=TEST_GCP_CONN_ID,
delegate_to=TEST_DELEGATE_TO,
impersonation_chain=TEST_IMPERSONATION_CHAIN,
)
mock_hook.return_value.list.return_value = ["NOT_EMPTY_LIST"]
result = task.poke(mock.MagicMock)
mock_hook.assert_called_once_with(
delegate_to=TEST_DELEGATE_TO,
gcp_conn_id=TEST_GCP_CONN_ID,
impersonation_chain=TEST_IMPERSONATION_CHAIN,
)
mock_hook.return_value.list.assert_called_once_with(TEST_BUCKET,
prefix=TEST_PREFIX)
assert result is True
def test_execute(self, mock_hook):
task = GCSObjectsWtihPrefixExistenceSensor(
task_id="task-id",
bucket=TEST_BUCKET,
prefix=TEST_PREFIX,
google_cloud_conn_id=TEST_GCP_CONN_ID,
delegate_to=TEST_DELEGATE_TO,
impersonation_chain=TEST_IMPERSONATION_CHAIN,
poke_interval=0,
)
generated_messages = [f'test-prefix/obj{i}' for i in range(5)]
mock_hook.return_value.list.return_value = generated_messages
response = task.execute(None)
mock_hook.assert_called_once_with(
delegate_to=TEST_DELEGATE_TO,
google_cloud_storage_conn_id=TEST_GCP_CONN_ID,
impersonation_chain=TEST_IMPERSONATION_CHAIN,
)
mock_hook.return_value.list.assert_called_once_with(TEST_BUCKET,
prefix=TEST_PREFIX)
assert response == generated_messages
