mirror of
https://github.com/home-assistant/core
synced 2024-08-15 18:25:44 +02:00
334 lines
9.2 KiB
Python
334 lines
9.2 KiB
Python
"""Aiohttp test utils."""
|
|
import asyncio
|
|
from contextlib import contextmanager
|
|
from http import HTTPStatus
|
|
import re
|
|
from unittest import mock
|
|
from urllib.parse import parse_qs
|
|
|
|
from aiohttp import ClientSession
|
|
from aiohttp.client_exceptions import ClientError, ClientResponseError
|
|
from aiohttp.streams import StreamReader
|
|
from multidict import CIMultiDict
|
|
from yarl import URL
|
|
|
|
from homeassistant.const import EVENT_HOMEASSISTANT_CLOSE
|
|
from homeassistant.helpers.json import json_dumps, json_loads
|
|
|
|
RETYPE = type(re.compile(""))
|
|
|
|
|
|
def mock_stream(data):
|
|
"""Mock a stream with data."""
|
|
protocol = mock.Mock(_reading_paused=False)
|
|
stream = StreamReader(protocol, limit=2**16)
|
|
stream.feed_data(data)
|
|
stream.feed_eof()
|
|
return stream
|
|
|
|
|
|
class AiohttpClientMocker:
|
|
"""Mock Aiohttp client requests."""
|
|
|
|
def __init__(self):
|
|
"""Initialize the request mocker."""
|
|
self._mocks = []
|
|
self._cookies = {}
|
|
self.mock_calls = []
|
|
|
|
def request(
|
|
self,
|
|
method,
|
|
url,
|
|
*,
|
|
auth=None,
|
|
status=HTTPStatus.OK,
|
|
text=None,
|
|
data=None,
|
|
content=None,
|
|
json=None,
|
|
params=None,
|
|
headers={},
|
|
exc=None,
|
|
cookies=None,
|
|
side_effect=None,
|
|
):
|
|
"""Mock a request."""
|
|
if not isinstance(url, RETYPE):
|
|
url = URL(url)
|
|
if params:
|
|
url = url.with_query(params)
|
|
|
|
self._mocks.append(
|
|
AiohttpClientMockResponse(
|
|
method=method,
|
|
url=url,
|
|
status=status,
|
|
response=content,
|
|
json=json,
|
|
text=text,
|
|
cookies=cookies,
|
|
exc=exc,
|
|
headers=headers,
|
|
side_effect=side_effect,
|
|
)
|
|
)
|
|
|
|
def get(self, *args, **kwargs):
|
|
"""Register a mock get request."""
|
|
self.request("get", *args, **kwargs)
|
|
|
|
def put(self, *args, **kwargs):
|
|
"""Register a mock put request."""
|
|
self.request("put", *args, **kwargs)
|
|
|
|
def post(self, *args, **kwargs):
|
|
"""Register a mock post request."""
|
|
self.request("post", *args, **kwargs)
|
|
|
|
def delete(self, *args, **kwargs):
|
|
"""Register a mock delete request."""
|
|
self.request("delete", *args, **kwargs)
|
|
|
|
def options(self, *args, **kwargs):
|
|
"""Register a mock options request."""
|
|
self.request("options", *args, **kwargs)
|
|
|
|
def patch(self, *args, **kwargs):
|
|
"""Register a mock patch request."""
|
|
self.request("patch", *args, **kwargs)
|
|
|
|
@property
|
|
def call_count(self):
|
|
"""Return the number of requests made."""
|
|
return len(self.mock_calls)
|
|
|
|
def clear_requests(self):
|
|
"""Reset mock calls."""
|
|
self._mocks.clear()
|
|
self._cookies.clear()
|
|
self.mock_calls.clear()
|
|
|
|
def create_session(self, loop):
|
|
"""Create a ClientSession that is bound to this mocker."""
|
|
session = ClientSession(loop=loop, json_serialize=json_dumps)
|
|
# Setting directly on `session` will raise deprecation warning
|
|
object.__setattr__(session, "_request", self.match_request)
|
|
return session
|
|
|
|
async def match_request(
|
|
self,
|
|
method,
|
|
url,
|
|
*,
|
|
data=None,
|
|
auth=None,
|
|
params=None,
|
|
headers=None,
|
|
allow_redirects=None,
|
|
timeout=None,
|
|
json=None,
|
|
cookies=None,
|
|
**kwargs,
|
|
):
|
|
"""Match a request against pre-registered requests."""
|
|
data = data or json
|
|
url = URL(url)
|
|
if params:
|
|
url = url.with_query(params)
|
|
|
|
for response in self._mocks:
|
|
if response.match_request(method, url, params):
|
|
self.mock_calls.append((method, url, data, headers))
|
|
if response.side_effect:
|
|
response = await response.side_effect(method, url, data)
|
|
if response.exc:
|
|
raise response.exc
|
|
return response
|
|
|
|
assert False, "No mock registered for {} {} {}".format(
|
|
method.upper(), url, params
|
|
)
|
|
|
|
|
|
class AiohttpClientMockResponse:
|
|
"""Mock Aiohttp client response."""
|
|
|
|
def __init__(
|
|
self,
|
|
method,
|
|
url,
|
|
status=HTTPStatus.OK,
|
|
response=None,
|
|
json=None,
|
|
text=None,
|
|
cookies=None,
|
|
exc=None,
|
|
headers=None,
|
|
side_effect=None,
|
|
):
|
|
"""Initialize a fake response."""
|
|
if json is not None:
|
|
text = json_dumps(json)
|
|
if text is not None:
|
|
response = text.encode("utf-8")
|
|
if response is None:
|
|
response = b""
|
|
|
|
self.charset = "utf-8"
|
|
self.method = method
|
|
self._url = url
|
|
self.status = status
|
|
self.response = response
|
|
self.exc = exc
|
|
self.side_effect = side_effect
|
|
self._headers = CIMultiDict(headers or {})
|
|
self._cookies = {}
|
|
|
|
if cookies:
|
|
for name, data in cookies.items():
|
|
cookie = mock.MagicMock()
|
|
cookie.value = data
|
|
self._cookies[name] = cookie
|
|
|
|
def match_request(self, method, url, params=None):
|
|
"""Test if response answers request."""
|
|
if method.lower() != self.method.lower():
|
|
return False
|
|
|
|
# regular expression matching
|
|
if isinstance(self._url, RETYPE):
|
|
return self._url.search(str(url)) is not None
|
|
|
|
if (
|
|
self._url.scheme != url.scheme
|
|
or self._url.host != url.host
|
|
or self._url.path != url.path
|
|
):
|
|
return False
|
|
|
|
# Ensure all query components in matcher are present in the request
|
|
request_qs = parse_qs(url.query_string)
|
|
matcher_qs = parse_qs(self._url.query_string)
|
|
for key, vals in matcher_qs.items():
|
|
for val in vals:
|
|
try:
|
|
request_qs.get(key, []).remove(val)
|
|
except ValueError:
|
|
return False
|
|
|
|
return True
|
|
|
|
@property
|
|
def headers(self):
|
|
"""Return content_type."""
|
|
return self._headers
|
|
|
|
@property
|
|
def cookies(self):
|
|
"""Return dict of cookies."""
|
|
return self._cookies
|
|
|
|
@property
|
|
def url(self):
|
|
"""Return yarl of URL."""
|
|
return self._url
|
|
|
|
@property
|
|
def content_type(self):
|
|
"""Return yarl of URL."""
|
|
return self._headers.get("content-type")
|
|
|
|
@property
|
|
def content(self):
|
|
"""Return content."""
|
|
return mock_stream(self.response)
|
|
|
|
async def read(self):
|
|
"""Return mock response."""
|
|
return self.response
|
|
|
|
async def text(self, encoding="utf-8", errors="strict"):
|
|
"""Return mock response as a string."""
|
|
return self.response.decode(encoding, errors=errors)
|
|
|
|
async def json(self, encoding="utf-8", content_type=None, loads=json_loads):
|
|
"""Return mock response as a json."""
|
|
return loads(self.response.decode(encoding))
|
|
|
|
def release(self):
|
|
"""Mock release."""
|
|
|
|
def raise_for_status(self):
|
|
"""Raise error if status is 400 or higher."""
|
|
if self.status >= 400:
|
|
request_info = mock.Mock(real_url="http://example.com")
|
|
raise ClientResponseError(
|
|
request_info=request_info,
|
|
history=None,
|
|
code=self.status,
|
|
headers=self.headers,
|
|
)
|
|
|
|
def close(self):
|
|
"""Mock close."""
|
|
|
|
|
|
@contextmanager
|
|
def mock_aiohttp_client():
|
|
"""Context manager to mock aiohttp client."""
|
|
mocker = AiohttpClientMocker()
|
|
|
|
def create_session(hass, *args, **kwargs):
|
|
session = mocker.create_session(hass.loop)
|
|
|
|
async def close_session(event):
|
|
"""Close session."""
|
|
await session.close()
|
|
|
|
hass.bus.async_listen_once(EVENT_HOMEASSISTANT_CLOSE, close_session)
|
|
|
|
return session
|
|
|
|
with mock.patch(
|
|
"homeassistant.helpers.aiohttp_client._async_create_clientsession",
|
|
side_effect=create_session,
|
|
):
|
|
yield mocker
|
|
|
|
|
|
class MockLongPollSideEffect:
|
|
"""Imitate a long_poll request.
|
|
|
|
It should be created and used as a side effect for a GET/PUT/etc. request.
|
|
Once created, actual responses are queued with queue_response
|
|
If queue is empty, will await until done.
|
|
"""
|
|
|
|
def __init__(self):
|
|
"""Initialize the queue."""
|
|
self.semaphore = asyncio.Semaphore(0)
|
|
self.response_list = []
|
|
self.stopping = False
|
|
|
|
async def __call__(self, method, url, data):
|
|
"""Fetch the next response from the queue or wait until the queue has items."""
|
|
if self.stopping:
|
|
raise ClientError()
|
|
await self.semaphore.acquire()
|
|
kwargs = self.response_list.pop(0)
|
|
return AiohttpClientMockResponse(method=method, url=url, **kwargs)
|
|
|
|
def queue_response(self, **kwargs):
|
|
"""Add a response to the long_poll queue."""
|
|
self.response_list.append(kwargs)
|
|
self.semaphore.release()
|
|
|
|
def stop(self):
|
|
"""Stop the current request and future ones.
|
|
|
|
This avoids an exception if there is someone waiting when exiting test.
|
|
"""
|
|
self.stopping = True
|
|
self.queue_response(exc=ClientError())
|