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_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_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