def test_should_support_db_backend(self, mock_session):
with _prepare_app():
mock_backend = DatabaseBackend(app=celery_executor.app,
url="sqlite3://")
with mock.patch.object(celery_executor.app, 'backend',
mock_backend):
mock_session = mock_backend.ResultSession.return_value # pylint: disable=no-member
mock_session.query.return_value.filter.return_value.all.return_value = [
mock.MagicMock(
**{
"to_dict.return_value": {
"status": "SUCCESS",
"task_id": "123"
}
})
]
fetcher = BulkStateFetcher()
result = fetcher.get_many([
mock.MagicMock(task_id="123"),
mock.MagicMock(task_id="456"),
])
self.assertEqual(result, {
'123': ('SUCCESS', None),
'456': ("PENDING", None)
})
def test_should_support_db_backend(self, mock_session):
with _prepare_app():
mock_backend = DatabaseBackend(app=celery_executor.app,
url="sqlite3://")
with mock.patch.object(
celery_executor.app, 'backend',
mock_backend), self.assertLogs(
"airflow.executors.celery_executor.BulkStateFetcher",
level="DEBUG") as cm:
mock_session = mock_backend.ResultSession.return_value # pylint: disable=no-member
mock_session.query.return_value.filter.return_value.all.return_value = [
mock.MagicMock(
**{
"to_dict.return_value": {
"status": "SUCCESS",
"task_id": "123"
}
})
]
fetcher = BulkStateFetcher()
result = fetcher.get_many([
mock.MagicMock(task_id="123"),
mock.MagicMock(task_id="456"),
])
assert result == {'123': ('SUCCESS', None), '456': ("PENDING", None)}
assert [
'DEBUG:airflow.executors.celery_executor.BulkStateFetcher:Fetched 2 state(s) for 2 task(s)'
] == cm.output
def test_should_support_kv_backend(self, mock_mget):
with _prepare_app():
mock_backend = BaseKeyValueStoreBackend(app=celery_executor.app)
with mock.patch.object(
celery_executor.app, 'backend',
mock_backend), self.assertLogs(
"airflow.executors.celery_executor.BulkStateFetcher",
level="DEBUG") as cm:
fetcher = BulkStateFetcher()
result = fetcher.get_many([
mock.MagicMock(task_id="123"),
mock.MagicMock(task_id="456"),
])
# Assert called - ignore order
mget_args, _ = mock_mget.call_args
assert set(mget_args[0]) == {
b'celery-task-meta-456', b'celery-task-meta-123'
}
mock_mget.assert_called_once_with(mock.ANY)
assert result == {'123': ('SUCCESS', None), '456': ("PENDING", None)}
assert [
'DEBUG:airflow.executors.celery_executor.BulkStateFetcher:Fetched 2 state(s) for 2 task(s)'
] == cm.output
def test_should_support_base_backend(self):
with _prepare_app():
mock_backend = mock.MagicMock(autospec=BaseBackend)
with mock.patch.object(celery_executor.app, 'backend', mock_backend):
fetcher = BulkStateFetcher(1)
result = fetcher.get_many([
ClassWithCustomAttributes(task_id="123", state='SUCCESS'),
ClassWithCustomAttributes(task_id="456", state="PENDING"),
])
self.assertEqual(result, {'123': 'SUCCESS', '456': "PENDING"})
def test_should_support_kv_backend(self, mock_mget):
with _prepare_app():
mock_backend = BaseKeyValueStoreBackend(app=celery_executor.app)
with mock.patch.object(celery_executor.app, 'backend', mock_backend):
fetcher = BulkStateFetcher()
result = fetcher.get_many([
mock.MagicMock(task_id="123"),
mock.MagicMock(task_id="456"),
])
# Assert called - ignore order
mget_args, _ = mock_mget.call_args
self.assertEqual(set(mget_args[0]), {b'celery-task-meta-456', b'celery-task-meta-123'})
mock_mget.assert_called_once_with(mock.ANY)
self.assertEqual(result, {'123': 'SUCCESS', '456': "PENDING"})
def test_should_support_base_backend(self):
with _prepare_app():
mock_backend = mock.MagicMock(autospec=BaseBackend)
with mock.patch.object(
celery_executor.app, 'backend',
mock_backend), self.assertLogs(
"airflow.executors.celery_executor.BulkStateFetcher",
level="DEBUG") as cm:
fetcher = BulkStateFetcher(1)
result = fetcher.get_many([
ClassWithCustomAttributes(task_id="123", state='SUCCESS'),
ClassWithCustomAttributes(task_id="456", state="PENDING"),
])
assert result == {'123': ('SUCCESS', None), '456': ("PENDING", None)}
assert [
'DEBUG:airflow.executors.celery_executor.BulkStateFetcher:Fetched 2 state(s) for 2 task(s)'
] == cm.output
