tests/mock/mock_aiohttp.py

import json
from typing import Callable

from aiohttp.client import ClientSession

origin_request = ClientSession.request


class MockAioResponse:
    check_funcs: dict[tuple[str, str], Callable[[dict], str]] = {}
    rsp_cache: dict[str, str] = {}
    name = "aiohttp"
    status = 200

    def __init__(self, session, method, url, **kwargs) -> None:
        fn = self.check_funcs.get((method, url))
        _kwargs = {k: v for k, v in kwargs.items() if k != "proxy"}
        self.key = f"{self.name}-{method}-{url}-{fn(kwargs) if fn else json.dumps(_kwargs, sort_keys=True)}"
        self.mng = self.response = None
        if self.key not in self.rsp_cache:
            self.mng = origin_request(session, method, url, **kwargs)

    async def __aenter__(self):
        if self.response:
            await self.response.__aenter__()
            self.status = self.response.status
        elif self.mng:
            self.response = await self.mng.__aenter__()
        return self

    async def __aexit__(self, *args, **kwargs):
        if self.response:
            await self.response.__aexit__(*args, **kwargs)
            self.response = None
        elif self.mng:
            await self.mng.__aexit__(*args, **kwargs)
            self.mng = None

    async def json(self, *args, **kwargs):
        if self.key in self.rsp_cache:
            return self.rsp_cache[self.key]
        data = await self.response.json(*args, **kwargs)
        self.rsp_cache[self.key] = data
        return data

    @property
    def content(self):
        return self

    async def read(self):
        if self.key in self.rsp_cache:
            return eval(self.rsp_cache[self.key])
        data = await self.response.content.read()
        self.rsp_cache[self.key] = str(data)
        return data

    def raise_for_status(self):
        if self.response:
            self.response.raise_for_status()