374 lines
8.2 KiB
Python
374 lines
8.2 KiB
Python
import asyncio
|
|
from collections import OrderedDict
|
|
from functools import _make_key, partial
|
|
from unittest import mock
|
|
|
|
import pytest
|
|
from great_ai.external.async_lru import (
|
|
__cache_touch,
|
|
_cache_clear,
|
|
_cache_hit,
|
|
_cache_info,
|
|
_cache_invalidate,
|
|
_cache_miss,
|
|
_close,
|
|
_close_waited,
|
|
_done_callback,
|
|
_open,
|
|
_wait_closed,
|
|
)
|
|
|
|
|
|
class Wrapped:
|
|
pass
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_done_callback_cancelled():
|
|
loop = asyncio.get_running_loop()
|
|
task = loop.create_future()
|
|
fut = loop.create_future()
|
|
|
|
task.add_done_callback(partial(_done_callback, fut))
|
|
|
|
task.cancel()
|
|
|
|
await asyncio.sleep(0)
|
|
|
|
assert fut.cancelled()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_done_callback_exception():
|
|
loop = asyncio.get_running_loop()
|
|
task = loop.create_future()
|
|
fut = loop.create_future()
|
|
|
|
task.add_done_callback(partial(_done_callback, fut))
|
|
|
|
exc = ZeroDivisionError()
|
|
|
|
task.set_exception(exc)
|
|
|
|
await asyncio.sleep(0)
|
|
|
|
with pytest.raises(ZeroDivisionError):
|
|
await fut
|
|
|
|
with pytest.raises(ZeroDivisionError):
|
|
fut.result()
|
|
|
|
assert fut.exception() is exc
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_done_callback():
|
|
loop = asyncio.get_running_loop()
|
|
task = loop.create_future()
|
|
fut = loop.create_future()
|
|
|
|
task.add_done_callback(partial(_done_callback, fut))
|
|
|
|
task.set_result(1)
|
|
|
|
await asyncio.sleep(0)
|
|
|
|
assert fut.result() == 1
|
|
|
|
|
|
def test_cache_invalidate_typed():
|
|
wrapped = Wrapped()
|
|
wrapped._cache = {}
|
|
|
|
args = (1,)
|
|
kwargs = {"1": 1}
|
|
|
|
from_cache = _cache_invalidate(wrapped, True, *args, **kwargs)
|
|
|
|
assert not from_cache
|
|
|
|
key = _make_key(args, kwargs, True)
|
|
|
|
wrapped._cache[key] = 0
|
|
|
|
from_cache = _cache_invalidate(wrapped, True, *args, **kwargs)
|
|
|
|
assert from_cache
|
|
|
|
assert len(wrapped._cache) == 0
|
|
|
|
wrapped._cache[key] = 0
|
|
|
|
args = (1.0,)
|
|
|
|
from_cache = _cache_invalidate(wrapped, True, *args, **kwargs)
|
|
|
|
assert not from_cache
|
|
|
|
wrapped._cache[key] = 1
|
|
|
|
|
|
def test_cache_invalidate_not_typed():
|
|
wrapped = Wrapped()
|
|
wrapped._cache = {}
|
|
|
|
args = (1,)
|
|
kwargs = {"1": 1}
|
|
|
|
from_cache = _cache_invalidate(wrapped, False, *args, **kwargs)
|
|
|
|
assert not from_cache
|
|
|
|
key = _make_key(args, kwargs, False)
|
|
|
|
wrapped._cache[key] = 0
|
|
|
|
from_cache = _cache_invalidate(wrapped, False, *args, **kwargs)
|
|
|
|
assert from_cache
|
|
|
|
assert len(wrapped._cache) == 0
|
|
|
|
wrapped._cache[key] = 0
|
|
|
|
args = (1.0,)
|
|
|
|
from_cache = _cache_invalidate(wrapped, False, *args, **kwargs)
|
|
|
|
assert from_cache
|
|
|
|
assert len(wrapped._cache) == 0
|
|
|
|
|
|
def test_cache_clear():
|
|
wrapped = Wrapped()
|
|
|
|
attrs = ["hits", "_cache", "tasks"]
|
|
for attr in attrs:
|
|
assert not hasattr(wrapped, attr)
|
|
|
|
_cache_clear(wrapped)
|
|
|
|
for attr in attrs:
|
|
assert hasattr(wrapped, attr)
|
|
|
|
assert wrapped.hits == wrapped.misses == 0
|
|
assert isinstance(wrapped._cache, dict)
|
|
assert len(wrapped._cache) == 0
|
|
assert isinstance(wrapped.tasks, set)
|
|
assert len(wrapped.tasks) == 0
|
|
|
|
_cache = wrapped._cache
|
|
tasks = wrapped.tasks
|
|
|
|
_cache_clear(wrapped)
|
|
|
|
assert wrapped._cache is not _cache
|
|
assert wrapped.tasks is not tasks
|
|
|
|
|
|
def test_open():
|
|
wrapped = Wrapped()
|
|
wrapped.hits = wrapped.misses = 1
|
|
wrapped._cache = {}
|
|
wrapped.tasks = set()
|
|
wrapped.closed = True
|
|
|
|
with pytest.raises(RuntimeError):
|
|
_open(wrapped)
|
|
|
|
wrapped.hits = wrapped.misses = 0
|
|
|
|
_open(wrapped)
|
|
|
|
assert not wrapped.closed
|
|
|
|
with pytest.raises(RuntimeError):
|
|
_open(wrapped)
|
|
|
|
|
|
def test_close(loop):
|
|
wrapped = Wrapped()
|
|
wrapped.closed = False
|
|
wrapped.tasks = set()
|
|
|
|
awaitable = _close(wrapped, cancel=False, return_exceptions=True)
|
|
loop.run_until_complete(awaitable)
|
|
|
|
assert wrapped.closed
|
|
|
|
with pytest.raises(RuntimeError):
|
|
_close(wrapped, cancel=False, return_exceptions=True)
|
|
|
|
fut = loop.create_future()
|
|
wrapped.closed = False
|
|
wrapped.tasks = {fut}
|
|
|
|
awaitable = _close(wrapped, cancel=True, return_exceptions=True)
|
|
loop.run_until_complete(awaitable)
|
|
|
|
assert fut.cancelled()
|
|
|
|
fut = loop.create_future()
|
|
fut.set_result(None)
|
|
wrapped.closed = False
|
|
wrapped.tasks = {fut}
|
|
|
|
awaitable = _close(wrapped, cancel=True, return_exceptions=True)
|
|
loop.run_until_complete(awaitable)
|
|
|
|
assert not fut.cancelled()
|
|
|
|
fut = loop.create_future()
|
|
fut.set_exception(ZeroDivisionError)
|
|
wrapped.closed = False
|
|
wrapped.tasks = {fut}
|
|
|
|
awaitable = _close(wrapped, cancel=True, return_exceptions=True)
|
|
loop.run_until_complete(awaitable)
|
|
|
|
assert not fut.cancelled()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_wait_closed():
|
|
loop = asyncio.get_running_loop()
|
|
wrapped = Wrapped()
|
|
wrapped.tasks = set()
|
|
|
|
with mock.patch("great_ai.external.async_lru._close_waited") as mocked:
|
|
ret = await _wait_closed(
|
|
wrapped,
|
|
return_exceptions=True,
|
|
)
|
|
assert ret == []
|
|
mocked.assert_called_once()
|
|
|
|
asyncio.set_event_loop(loop)
|
|
with mock.patch("great_ai.external.async_lru._close_waited") as mocked:
|
|
ret = await _wait_closed(
|
|
wrapped,
|
|
return_exceptions=True,
|
|
)
|
|
assert ret == []
|
|
mocked.assert_called_once()
|
|
asyncio.set_event_loop(None)
|
|
|
|
fut = loop.create_future()
|
|
fut.set_result(None)
|
|
wrapped.tasks = {fut}
|
|
with mock.patch("great_ai.external.async_lru._close_waited") as mocked:
|
|
ret = await _wait_closed(
|
|
wrapped,
|
|
return_exceptions=True,
|
|
)
|
|
assert ret == [None]
|
|
mocked.assert_called_once()
|
|
|
|
exc = ZeroDivisionError()
|
|
fut = loop.create_future()
|
|
fut.set_exception(exc)
|
|
wrapped.tasks = {fut}
|
|
with mock.patch("great_ai.external.async_lru._close_waited") as mocked:
|
|
ret = await _wait_closed(
|
|
wrapped,
|
|
return_exceptions=True,
|
|
)
|
|
assert ret == [exc]
|
|
mocked.assert_called_once()
|
|
|
|
fut = loop.create_future()
|
|
fut.set_exception(ZeroDivisionError)
|
|
wrapped.tasks = {fut}
|
|
with mock.patch("great_ai.external.async_lru._close_waited") as mocked:
|
|
with pytest.raises(ZeroDivisionError):
|
|
await _wait_closed(
|
|
wrapped,
|
|
return_exceptions=False,
|
|
)
|
|
# the awaited result raised before _wait_closed could flush its done
|
|
# callback, so yield once to let the scheduled _close_waited run
|
|
await asyncio.sleep(0)
|
|
mocked.assert_called_once()
|
|
|
|
|
|
def test_close_waited():
|
|
wrapped = Wrapped()
|
|
# _close_waited calls wrapped.cache_clear(); mock that directly (patching the
|
|
# module-level _cache_clear would not affect the already-bound partial).
|
|
wrapped.cache_clear = mock.Mock()
|
|
|
|
_close_waited(wrapped, None)
|
|
|
|
wrapped.cache_clear.assert_called_once()
|
|
|
|
|
|
def test_cache_info():
|
|
wrapped = Wrapped()
|
|
wrapped._cache = {}
|
|
wrapped.hits = wrapped.misses = 0
|
|
|
|
assert (0, 0, 3, 0) == _cache_info(wrapped, 3)
|
|
|
|
wrapped._cache[1] = 1
|
|
|
|
assert (0, 0, 1, 1) == _cache_info(wrapped, 1)
|
|
|
|
wrapped.hits = 2
|
|
wrapped.misses = 3
|
|
wrapped._cache[2] = 2
|
|
|
|
assert (2, 3, 5, 2) == _cache_info(wrapped, 5)
|
|
|
|
|
|
def test__cache_touch():
|
|
wrapped = Wrapped()
|
|
|
|
wrapped._cache = OrderedDict()
|
|
wrapped._cache[1] = 1
|
|
wrapped._cache[2] = 2
|
|
|
|
__cache_touch(wrapped, 1)
|
|
assert list(wrapped._cache) == [2, 1]
|
|
|
|
__cache_touch(wrapped, 2)
|
|
assert list(wrapped._cache) == [1, 2]
|
|
|
|
# test KeyError
|
|
__cache_touch(wrapped, 100)
|
|
|
|
|
|
def test_cache_hit():
|
|
wrapped = Wrapped()
|
|
wrapped.hits = 1
|
|
wrapped._cache = OrderedDict()
|
|
wrapped._cache[1] = 1
|
|
|
|
with mock.patch("great_ai.external.async_lru.__cache_touch") as mocked:
|
|
_cache_hit(wrapped, 1)
|
|
|
|
mocked.assert_called_once()
|
|
|
|
assert wrapped.hits == 2
|
|
|
|
_cache_hit(wrapped, 1)
|
|
|
|
assert wrapped.hits == 3
|
|
|
|
|
|
def test_cache_miss():
|
|
wrapped = Wrapped()
|
|
wrapped.misses = 1
|
|
wrapped._cache = OrderedDict()
|
|
wrapped._cache[1] = 1
|
|
|
|
with mock.patch("great_ai.external.async_lru.__cache_touch") as mocked:
|
|
_cache_miss(wrapped, 1)
|
|
|
|
mocked.assert_called_once()
|
|
|
|
assert wrapped.misses == 2
|
|
|
|
_cache_miss(wrapped, 1)
|
|
|
|
assert wrapped.misses == 3
|