diff --git a/api/conf/asgi_handler.py b/api/conf/asgi_handler.py index ed68f94e3f6..5e26de95985 100644 --- a/api/conf/asgi_handler.py +++ b/api/conf/asgi_handler.py @@ -64,7 +64,8 @@ async def __call__(self, scope, receive, send): async def shutdown(self): live_handlers = 0 - for handler_ref in self._on_shutdown: + while self._on_shutdown: + handler_ref = self._on_shutdown.pop() if not (handler := handler_ref()): self.logger.debug("Reference dead, skipping handler") continue diff --git a/api/test/unit/utils/test_aiohttp.py b/api/test/unit/utils/test_aiohttp.py new file mode 100644 index 00000000000..221d52026df --- /dev/null +++ b/api/test/unit/utils/test_aiohttp.py @@ -0,0 +1,64 @@ +import asyncio + +import pytest +from asgiref.sync import async_to_sync + +from api.utils.aiohttp import get_aiohttp_session +from conf.asgi import application + + +@pytest.fixture(autouse=True) +def get_new_loop(): + loops: list[asyncio.AbstractEventLoop] = [] + + def _get_new_loop(): + loop = asyncio.new_event_loop() + loops.append(loop) + return loop + + yield _get_new_loop + + async_to_sync(application.shutdown)() + for loop in loops: + loop.close() + + +def test_reuses_session_within_same_loop(get_new_loop): + loop = get_new_loop() + + session_1 = loop.run_until_complete(get_aiohttp_session()) + session_2 = loop.run_until_complete(get_aiohttp_session()) + + assert session_1 is session_2 + + +def test_creates_new_session_for_separate_loops(get_new_loop): + loop_1 = get_new_loop() + loop_2 = get_new_loop() + + loop_1_session = loop_1.run_until_complete(get_aiohttp_session()) + loop_2_session = loop_2.run_until_complete(get_aiohttp_session()) + + assert loop_1_session is not loop_2_session + + +def test_multiple_loops_reuse_separate_sessions(get_new_loop): + loop_1 = get_new_loop() + loop_2 = get_new_loop() + + loop_1_session_1 = loop_1.run_until_complete(get_aiohttp_session()) + loop_1_session_2 = loop_1.run_until_complete(get_aiohttp_session()) + loop_2_session_1 = loop_2.run_until_complete(get_aiohttp_session()) + loop_2_session_2 = loop_2.run_until_complete(get_aiohttp_session()) + + assert loop_1_session_1 is loop_1_session_2 + assert loop_2_session_1 is loop_2_session_2 + + +def test_registers_shutdown_cleanup(get_new_loop): + assert len(application._on_shutdown) == 0 + + loop = get_new_loop() + loop.run_until_complete(get_aiohttp_session()) + + assert len(application._on_shutdown) == 1