conftest.py

import asyncio
import os
import sys
import types

os.environ.setdefault('AIOHTTP_NO_WARN_LOGS', '1')

import aiohttp
import pytest

# This code adds the project root directory to the Python path, allowing imports to work correctly when running tests.
# Without this file, you might encounter ModuleNotFoundError when trying to import modules from your project, especially when running tests.
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__))))

from tests.helpers.embeddings import make_mock_embedder

__all__ = ['mock_embedder']

_original_client_session = aiohttp.ClientSession
_open_aiohttp_sessions: set[aiohttp.ClientSession] = set()


def _tracking_client_session(*args, **kwargs):
    session = _original_client_session(*args, **kwargs)
    _open_aiohttp_sessions.add(session)
    original_close = session.close

    async def _close_session(self, *close_args, **close_kwargs):  # type: ignore[unused-argument]
        if self in _open_aiohttp_sessions:
            _open_aiohttp_sessions.remove(self)
        await original_close(*close_args, **close_kwargs)

    session.close = types.MethodType(_close_session, session)  # type: ignore[assignment]
    return session


aiohttp.ClientSession = _tracking_client_session  # type: ignore[assignment]
aiohttp.client.ClientSession._warn_unclosed = lambda self, *args, **kwargs: None  # type: ignore[assignment]
aiohttp.connector.BaseConnector._warn_unclosed = lambda self, *args, **kwargs: None  # type: ignore[attr-defined, assignment]


@pytest.fixture
def mock_embedder():
    return make_mock_embedder()


@pytest.fixture(scope='session', autouse=True)
def _patch_aiohttp_client_session():
    yield

    aiohttp.ClientSession = _original_client_session  # type: ignore[assignment]

    async def _close_all():
        if not _open_aiohttp_sessions:
            return
        await asyncio.gather(
            *(session.close() for session in list(_open_aiohttp_sessions)),
            return_exceptions=True,
        )
        _open_aiohttp_sessions.clear()

    asyncio.run(_close_all())