diff --git a/great_ai/external/__init__.py b/great_ai/external/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/great_ai/external/async_lru.py b/great_ai/external/async_lru.py new file mode 100644 index 0000000..c209d5a --- /dev/null +++ b/great_ai/external/async_lru.py @@ -0,0 +1,192 @@ +import asyncio +from collections import OrderedDict +from functools import _CacheInfo, _make_key, partial, wraps + +__version__ = "1.0.3" + +__all__ = ("alru_cache",) + + +def unpartial(fn): + while hasattr(fn, "func"): + fn = fn.func + + return fn + + +def _done_callback(fut, task): + if task.cancelled(): + fut.cancel() + return + + exc = task.exception() + if exc is not None: + fut.set_exception(exc) + return + + fut.set_result(task.result()) + + +def _cache_invalidate(wrapped, typed, *args, **kwargs): + key = _make_key(args, kwargs, typed) + + exists = key in wrapped._cache + + if exists: + wrapped._cache.pop(key) + + return exists + + +def _cache_clear(wrapped): + wrapped.hits = wrapped.misses = 0 + wrapped._cache = OrderedDict() + wrapped.tasks = set() + + +def _open(wrapped): + if not wrapped.closed: + raise RuntimeError("alru_cache is not closed") + + was_closed = ( + wrapped.hits == wrapped.misses == len(wrapped.tasks) == len(wrapped._cache) == 0 + ) + + if not was_closed: + raise RuntimeError("alru_cache was not closed correctly") + + wrapped.closed = False + + +def _close(wrapped, *, cancel=False, return_exceptions=True): + if wrapped.closed: + raise RuntimeError("alru_cache is closed") + + wrapped.closed = True + + if cancel: + for task in wrapped.tasks: + if not task.done(): # not sure is it possible + task.cancel() + + return _wait_closed(wrapped, return_exceptions=return_exceptions) + + +async def _wait_closed(wrapped, *, return_exceptions): + wait_closed = asyncio.gather(*wrapped.tasks, return_exceptions=return_exceptions) + + wait_closed.add_done_callback(partial(_close_waited, wrapped)) + + ret = await wait_closed + + # hack to get _close_waited callback to be executed + await asyncio.sleep(0) + + return ret + + +def _close_waited(wrapped, _): + wrapped.cache_clear() + + +def _cache_info(wrapped, maxsize): + return _CacheInfo( + wrapped.hits, + wrapped.misses, + maxsize, + len(wrapped._cache), + ) + + +def __cache_touch(wrapped, key): + try: + wrapped._cache.move_to_end(key) + except KeyError: # not sure is it possible + pass + + +def _cache_hit(wrapped, key): + wrapped.hits += 1 + __cache_touch(wrapped, key) + + +def _cache_miss(wrapped, key): + wrapped.misses += 1 + __cache_touch(wrapped, key) + + +def alru_cache( + fn=None, + maxsize=128, + typed=False, + *, + cache_exceptions=True, +): + def wrapper(fn): + _origin = unpartial(fn) + + if not asyncio.iscoroutinefunction(_origin): + raise RuntimeError("Coroutine function is required, got {}".format(fn)) + + # functools.partialmethod support + if hasattr(fn, "_make_unbound_method"): + fn = fn._make_unbound_method() + + @wraps(fn) + async def wrapped(*fn_args, **fn_kwargs): + if wrapped.closed: + raise RuntimeError("alru_cache is closed for {}".format(wrapped)) + + loop = asyncio.get_event_loop() + + key = _make_key(fn_args, fn_kwargs, typed) + + fut = wrapped._cache.get(key) + + if fut is not None: + if not fut.done(): + _cache_hit(wrapped, key) + return await asyncio.shield(fut) + + exc = fut._exception + + if exc is None or cache_exceptions: + _cache_hit(wrapped, key) + return fut.result() + + # exception here and cache_exceptions == False + wrapped._cache.pop(key) + + fut = loop.create_future() + task = loop.create_task(fn(*fn_args, **fn_kwargs)) + task.add_done_callback(partial(_done_callback, fut)) + + wrapped.tasks.add(task) + task.add_done_callback(wrapped.tasks.remove) + + wrapped._cache[key] = fut + + if maxsize is not None and len(wrapped._cache) > maxsize: + wrapped._cache.popitem(last=False) + + _cache_miss(wrapped, key) + return await asyncio.shield(fut) + + _cache_clear(wrapped) + wrapped._origin = _origin + wrapped.closed = False + wrapped.cache_info = partial(_cache_info, wrapped, maxsize) + wrapped.cache_clear = partial(_cache_clear, wrapped) + wrapped.invalidate = partial(_cache_invalidate, wrapped, typed) + wrapped.close = partial(_close, wrapped) + wrapped.open = partial(_open, wrapped) + + return wrapped + + if fn is None: + return wrapper + + if callable(fn) or hasattr(fn, "_make_unbound_method"): + return wrapper(fn) + + raise NotImplementedError("{} decorating is not supported".format(fn)) diff --git a/tests/external/__init__.py b/tests/external/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/external/conftest.py b/tests/external/conftest.py new file mode 100644 index 0000000..1fb0b17 --- /dev/null +++ b/tests/external/conftest.py @@ -0,0 +1,47 @@ +import asyncio +import gc +import os + +import pytest +from great_ai.external.async_lru import _CacheInfo + +asyncio.set_event_loop(None) + + +@pytest.fixture +def event_loop(request): + loop = asyncio.new_event_loop() + loop.set_debug(bool(os.environ.get("PYTHONASYNCIODEBUG"))) + + yield loop + + loop.call_soon(loop.stop) + loop.run_forever() + loop.close() + + gc.collect() + + +@pytest.fixture +def loop(event_loop, request): + asyncio.set_event_loop(None) + request.addfinalizer(lambda: asyncio.set_event_loop(None)) + + return event_loop + + +@pytest.fixture +def check_lru(request): + def _check_lru(wrapped, *, hits, misses, cache, tasks, maxsize=128): + assert wrapped.hits == hits + assert wrapped.misses == misses + assert len(wrapped._cache) == cache + assert len(wrapped.tasks) == tasks + assert wrapped.cache_info() == _CacheInfo( + hits=hits, + misses=misses, + maxsize=maxsize, + currsize=cache, + ) + + return _check_lru diff --git a/tests/external/test_basic.py b/tests/external/test_basic.py new file mode 100644 index 0000000..604662d --- /dev/null +++ b/tests/external/test_basic.py @@ -0,0 +1,187 @@ +import asyncio +from functools import partial + +import pytest +from great_ai.external.async_lru import alru_cache + +alru_cache_attrs = [ + "hits", + "misses", + "tasks", + "closed", + "cache_info", + "cache_clear", + "invalidate", + "close", + "open", +] + +alru_cache_calable_attrs = alru_cache_attrs.copy() +for attr in ["hits", "misses", "tasks", "closed"]: + alru_cache_calable_attrs.remove(attr) + + +def test_alru_cache_not_callable(loop): + with pytest.raises(NotImplementedError): + alru_cache("foo") + + +def test_alru_cache_not_coroutine(loop): + with pytest.raises(RuntimeError): + + @alru_cache + def not_coro(val): + return val + + +def test_alru_cache_deco(loop, check_lru): + asyncio.set_event_loop(loop) + + @alru_cache + async def coro(): + pass + + assert asyncio.iscoroutinefunction(coro) + + for attr in alru_cache_attrs: + assert hasattr(coro, attr) + for attr in alru_cache_calable_attrs: + assert callable(getattr(coro, attr)) + + assert isinstance(coro._cache, dict) + assert isinstance(coro.tasks, set) + check_lru(coro, hits=0, misses=0, cache=0, tasks=0) + + awaitable = coro() + assert asyncio.iscoroutine(awaitable) + loop.run_until_complete(awaitable) + + +def test_alru_cache_deco_called(check_lru, loop): + asyncio.set_event_loop(loop) + + @alru_cache() + async def coro(): + pass + + assert asyncio.iscoroutinefunction(coro) + + for attr in alru_cache_attrs: + assert hasattr(coro, attr) + for attr in alru_cache_calable_attrs: + assert callable(getattr(coro, attr)) + + assert isinstance(coro._cache, dict) + assert isinstance(coro.tasks, set) + check_lru(coro, hits=0, misses=0, cache=0, tasks=0) + + awaitable = coro() + assert asyncio.iscoroutine(awaitable) + loop.run_until_complete(awaitable) + + +def test_alru_cache_fn_called(check_lru, loop): + asyncio.set_event_loop(loop) + + async def coro(): + pass + + coro_wrapped = alru_cache(coro) + + assert asyncio.iscoroutinefunction(coro_wrapped) + + for attr in alru_cache_attrs: + assert hasattr(coro_wrapped, attr) + for attr in alru_cache_calable_attrs: + assert callable(getattr(coro_wrapped, attr)) + + assert isinstance(coro_wrapped._cache, dict) + assert isinstance(coro_wrapped.tasks, set) + check_lru(coro_wrapped, hits=0, misses=0, cache=0, tasks=0) + + awaitable = coro_wrapped() + assert asyncio.iscoroutine(awaitable) + loop.run_until_complete(awaitable) + + +def test_alru_cache_origin(loop): + asyncio.set_event_loop(loop) + + async def coro(): + pass + + coro_wrapped = alru_cache(coro) + + assert coro_wrapped._origin is coro + + coro_wrapped = alru_cache(partial(coro)) + + assert coro_wrapped._origin is coro + + +@pytest.mark.asyncio +async def test_alru_cache_await_same_result_async(check_lru, loop): + calls = 0 + val = object() + + @alru_cache() + async def coro(): + nonlocal calls + calls += 1 + + return val + + coros = [coro() for _ in range(100)] + ret = await asyncio.gather(*coros) + expected = [val] * 100 + assert ret == expected + check_lru(coro, hits=99, misses=1, cache=1, tasks=0) + + assert calls == 1 + assert await coro() is val + check_lru(coro, hits=100, misses=1, cache=1, tasks=0) + + +@pytest.mark.asyncio +async def test_alru_cache_await_same_result_coroutine(check_lru, loop): + calls = 0 + val = object() + + @alru_cache() + async def coro(): + nonlocal calls + calls += 1 + + return val + + coros = [coro() for _ in range(100)] + ret = await asyncio.gather(*coros) + expected = [val] * 100 + assert ret == expected + check_lru(coro, hits=99, misses=1, cache=1, tasks=0) + + assert calls == 1 + assert await coro() is val + check_lru(coro, hits=100, misses=1, cache=1, tasks=0) + + +@pytest.mark.asyncio +async def test_alru_cache_dict_not_shared(check_lru, loop): + async def coro(val): + return val + + coro1 = alru_cache()(coro) + coro2 = alru_cache()(coro) + + ret1 = await coro1(1) + check_lru(coro1, hits=0, misses=1, cache=1, tasks=0) + + ret2 = await coro2(1) + check_lru(coro2, hits=0, misses=1, cache=1, tasks=0) + + assert ret1 == ret2 + + assert coro1._cache[1].result() == coro2._cache[1].result() + assert coro1._cache != coro2._cache + assert coro1._cache.keys() == coro2._cache.keys() + assert coro1._cache is not coro2._cache diff --git a/tests/external/test_cache_clear.py b/tests/external/test_cache_clear.py new file mode 100644 index 0000000..4e59ea6 --- /dev/null +++ b/tests/external/test_cache_clear.py @@ -0,0 +1,22 @@ +import asyncio + +import pytest +from great_ai.external.async_lru import alru_cache + +pytestmark = pytest.mark.asyncio + + +async def test_cache_clear(check_lru, loop): + @alru_cache() + async def coro(val): + return val + + inputs = [1, 2, 3] + coros = [coro(v) for v in inputs] + ret = await asyncio.gather(*coros) + assert ret == inputs + check_lru(coro, hits=0, misses=3, cache=3, tasks=0) + + coro.cache_clear() + + check_lru(coro, hits=0, misses=0, cache=0, tasks=0) diff --git a/tests/external/test_cache_info.py b/tests/external/test_cache_info.py new file mode 100644 index 0000000..1a39b96 --- /dev/null +++ b/tests/external/test_cache_info.py @@ -0,0 +1,38 @@ +import asyncio + +import pytest +from great_ai.external.async_lru import alru_cache + +pytestmark = pytest.mark.asyncio + + +async def test_cache_info(check_lru, loop): + @alru_cache(maxsize=4) + async def coro(val): + return val + + inputs = [1, 2, 3] + coros = [coro(v) for v in inputs] + ret = await asyncio.gather(*coros) + assert ret == inputs + check_lru(coro, hits=0, misses=3, cache=3, tasks=0, maxsize=4) + + coro.cache_clear() + + check_lru(coro, hits=0, misses=0, cache=0, tasks=0, maxsize=4) + + inputs = [1, 1, 1] + coros = [coro(v) for v in inputs] + ret = await asyncio.gather(*coros) + assert ret == inputs + check_lru(coro, hits=2, misses=1, cache=1, tasks=0, maxsize=4) + + coro.cache_clear() + + check_lru(coro, hits=0, misses=0, cache=0, tasks=0, maxsize=4) + + inputs = [1, 2, 3, 4] * 2 + coros = [coro(v) for v in inputs] + ret = await asyncio.gather(*coros) + assert ret == inputs + check_lru(coro, hits=4, misses=4, cache=4, tasks=0, maxsize=4) diff --git a/tests/external/test_cache_invalidate.py b/tests/external/test_cache_invalidate.py new file mode 100644 index 0000000..48dfd8c --- /dev/null +++ b/tests/external/test_cache_invalidate.py @@ -0,0 +1,85 @@ +import asyncio + +import pytest +from great_ai.external.async_lru import alru_cache + +pytestmark = pytest.mark.asyncio + + +async def test_cache_invalidate(check_lru, loop): + @alru_cache() + async def coro(val): + return val + + inputs = [1, 2, 3] + + coro.invalidate(1) + coro.invalidate(2) + coro.invalidate(3) + + coros = [coro(v) for v in inputs] + + ret = await asyncio.gather(*coros) + + assert ret == inputs + check_lru(coro, hits=0, misses=3, cache=3, tasks=0) + + coro.invalidate(1) + check_lru(coro, hits=0, misses=3, cache=2, tasks=0) + coro.invalidate(2) + check_lru(coro, hits=0, misses=3, cache=1, tasks=0) + coro.invalidate(3) + check_lru(coro, hits=0, misses=3, cache=0, tasks=0) + + inputs = [1, 2, 3] + coros = [coro(v) for v in inputs] + + ret = await asyncio.gather(*coros) + + assert ret == inputs + check_lru(coro, hits=0, misses=6, cache=3, tasks=0) + + +async def test_cache_invalidate_multiple_args(check_lru, loop): + @alru_cache() + async def coro(*args): + return len(args) + + for i, size in enumerate(range(10)): + args = tuple(range(size)) + ret = await coro(*args) + assert ret == size + check_lru(coro, hits=0, misses=i + 1, cache=1, tasks=0) + coro.invalidate(*args) + check_lru(coro, hits=0, misses=i + 1, cache=0, tasks=0) + + for size in range(10): + args = tuple(range(size)) + ret = await coro(*args) + assert ret == size + check_lru(coro, hits=0, misses=20, cache=10, tasks=0) + + +async def test_cache_invalidate_multiple_args_different_order(check_lru, loop): + @alru_cache() + async def coro(*args): + return len(args) + + for i, size in enumerate(range(2, 10)): + args = tuple(range(size)) + rev_args = tuple(reversed(args)) + ret = await coro(*args) + assert ret == size + check_lru(coro, hits=0, misses=2 * i + 1, cache=i + 1, tasks=0) + ret = await coro(*rev_args) + # The reversed args should be a miss + check_lru(coro, hits=0, misses=2 * i + 2, cache=i + 2, tasks=0) + coro.invalidate(*rev_args) + # The reversed args should be invalidated + check_lru(coro, hits=0, misses=2 * i + 2, cache=i + 1, tasks=0) + + for i, size in enumerate(range(2, 10)): + args = tuple(range(size)) + ret = await coro(*args) + assert ret == size + check_lru(coro, hits=i + 1, misses=16, cache=8, tasks=0) diff --git a/tests/external/test_close.py b/tests/external/test_close.py new file mode 100644 index 0000000..3008899 --- /dev/null +++ b/tests/external/test_close.py @@ -0,0 +1,170 @@ +import asyncio + +import pytest +from great_ai.external.async_lru import alru_cache + +pytestmark = pytest.mark.asyncio + + +async def test_cache_close(check_lru, loop): + @alru_cache() + async def coro(val): + await asyncio.sleep(0.2) + + return val + + assert not coro.closed + + inputs = [1, 2, 3, 4, 5] + + coros = [coro(v) for v in inputs] + + gather = asyncio.gather(*coros) + + await asyncio.sleep(0.1) + + check_lru(coro, hits=0, misses=5, cache=5, tasks=5) + + close = coro.close() + + with pytest.raises(RuntimeError): + await coro(1) + + check_lru(coro, hits=0, misses=5, cache=5, tasks=5) + + ret_close = await close + + check_lru(coro, hits=0, misses=0, cache=0, tasks=0) + + ret_gather = await gather + + check_lru(coro, hits=0, misses=0, cache=0, tasks=0) + + assert set(ret_close) == set(ret_gather) == set(inputs) + + with pytest.raises(RuntimeError): + coro.close() + + +async def test_cache_close_cancel_return_exceptions(check_lru, loop): + @alru_cache() + async def coro(val): + await asyncio.sleep(0.2) + + return val + + inputs = [1, 2, 3, 4, 5] + + coros = [coro(v) for v in inputs] + + gather = asyncio.gather(*coros) + + await asyncio.sleep(0.1) + + close = coro.close(cancel=True) + + check_lru(coro, hits=0, misses=5, cache=5, tasks=5) + + ret_close = await close + + check_lru(coro, hits=0, misses=0, cache=0, tasks=0) + + for err in ret_close: + assert isinstance(err, asyncio.CancelledError) + + with pytest.raises(asyncio.CancelledError): + await gather + + +async def test_cache_close_cancel_not_return_exceptions(check_lru, loop): + @alru_cache() + async def coro(val): + await asyncio.sleep(0.2) + + return val + + inputs = [1, 2, 3, 4, 5] + + coros = [coro(v) for v in inputs] + + gather = asyncio.gather(*coros) + + await asyncio.sleep(0.1) + + close = coro.close(cancel=True, return_exceptions=False) + + check_lru(coro, hits=0, misses=5, cache=5, tasks=5) + + with pytest.raises(asyncio.CancelledError): + await close + + with pytest.raises(asyncio.CancelledError): + await gather + + check_lru(coro, hits=0, misses=0, cache=0, tasks=0) + + +async def test_cache_close_return_exceptions(check_lru, loop): + @alru_cache() + async def coro(val): + await asyncio.sleep(0.2) + + raise ZeroDivisionError + + inputs = [1, 2, 3, 4, 5] + + coros = [coro(v) for v in inputs] + + gather = asyncio.gather(*coros) + + await asyncio.sleep(0.1) + + close = coro.close() + + check_lru(coro, hits=0, misses=5, cache=5, tasks=5) + + with pytest.raises(ZeroDivisionError): + await gather + + check_lru(coro, hits=0, misses=5, cache=5, tasks=0) + + ret_close = await close + + check_lru(coro, hits=0, misses=0, cache=0, tasks=0) + + for err in ret_close: + assert isinstance(err, ZeroDivisionError) + + check_lru(coro, hits=0, misses=0, cache=0, tasks=0) + + +async def test_cache_close_not_return_exceptions(check_lru, loop): + @alru_cache() + async def coro(val): + await asyncio.sleep(0.2) + + raise ZeroDivisionError + + inputs = [1, 2, 3, 4, 5] + + coros = [coro(v) for v in inputs] + + gather = asyncio.gather(*coros, return_exceptions=True) + + await asyncio.sleep(0.1) + + close = coro.close(return_exceptions=False) + + check_lru(coro, hits=0, misses=5, cache=5, tasks=5) + + with pytest.raises(ZeroDivisionError): + await close + + check_lru(coro, hits=0, misses=0, cache=0, tasks=0) + + ret_gather = await gather + + for err in ret_gather: + assert isinstance(err, ZeroDivisionError) + + check_lru(coro, hits=0, misses=0, cache=0, tasks=0) diff --git a/tests/external/test_exception.py b/tests/external/test_exception.py new file mode 100644 index 0000000..3d4aa25 --- /dev/null +++ b/tests/external/test_exception.py @@ -0,0 +1,48 @@ +import asyncio + +import pytest +from great_ai.external.async_lru import alru_cache + +pytestmark = pytest.mark.asyncio + + +async def test_alru_cache_exception(check_lru, loop): + @alru_cache(cache_exceptions=True) + async def coro(val): + 1 / 0 + + inputs = [1, 1, 1] + coros = [coro(v) for v in inputs] + + ret = await asyncio.gather(*coros, return_exceptions=True) + + check_lru(coro, hits=2, misses=1, cache=1, tasks=0) + + for item in ret: + assert isinstance(item, ZeroDivisionError) + + with pytest.raises(ZeroDivisionError): + await coro(1) + + check_lru(coro, hits=3, misses=1, cache=1, tasks=0) + + +async def test_alru_not_cache_exception(check_lru, loop): + @alru_cache(cache_exceptions=False) + async def coro(val): + 1 / 0 + + inputs = [1, 1, 1] + coros = [coro(v) for v in inputs] + + ret = await asyncio.gather(*coros, return_exceptions=True) + + check_lru(coro, hits=2, misses=1, cache=1, tasks=0) + + for item in ret: + assert isinstance(item, ZeroDivisionError) + + with pytest.raises(ZeroDivisionError): + await coro(1) + + check_lru(coro, hits=2, misses=2, cache=1, tasks=0) diff --git a/tests/external/test_internals.py b/tests/external/test_internals.py new file mode 100644 index 0000000..1003bff --- /dev/null +++ b/tests/external/test_internals.py @@ -0,0 +1,366 @@ +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): + 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): + 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): + 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): + 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 == [] + assert mocked.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 == [] + assert mocked.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] + assert mocked.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] + assert mocked.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, + ) + assert mocked.called_once() + + +def test_close_waited(): + wrapped = Wrapped() + wrapped.cache_clear = partial(_cache_clear, wrapped) + + with mock.patch("great_ai.external.async_lru._cache_clear") as mocked: + _close_waited(wrapped, None) + + assert mocked.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) + + assert mocked.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) + + assert mocked.called_once() + + assert wrapped.misses == 2 + + _cache_miss(wrapped, 1) + + assert wrapped.misses == 3 diff --git a/tests/external/test_open.py b/tests/external/test_open.py new file mode 100644 index 0000000..b749df4 --- /dev/null +++ b/tests/external/test_open.py @@ -0,0 +1,39 @@ +import pytest +from great_ai.external.async_lru import alru_cache + +pytestmark = pytest.mark.asyncio + + +async def test_alru_cache_open(check_lru, loop): + @alru_cache() + async def coro(val): + return val + + await coro(1) + + check_lru(coro, hits=0, misses=1, cache=1, tasks=0) + + with pytest.raises(RuntimeError): + coro.open() + + close = coro.close() + + assert coro.closed + + with pytest.raises(RuntimeError): + await coro() + + with pytest.raises(RuntimeError): + coro.open() + + await close + + check_lru(coro, hits=0, misses=0, cache=0, tasks=0) + + coro.open() + + ret = await coro(1) + + assert ret == 1 + + check_lru(coro, hits=0, misses=1, cache=1, tasks=0) diff --git a/tests/external/test_partialmethod.py b/tests/external/test_partialmethod.py new file mode 100644 index 0000000..dfa888f --- /dev/null +++ b/tests/external/test_partialmethod.py @@ -0,0 +1,50 @@ +import asyncio +from functools import partial, partialmethod + +import pytest +from great_ai.external.async_lru import alru_cache + +pytestmark = pytest.mark.asyncio + + +async def test_partialmethod_basic(check_lru, loop): + class Obj: + async def _coro(self, val): + return val + + coro = alru_cache(partialmethod(_coro, 2)) + + obj = Obj() + + coros = [obj.coro() for _ in range(5)] + + check_lru(obj.coro, hits=0, misses=0, cache=0, tasks=0) + + ret = await asyncio.gather(*coros) + + check_lru(obj.coro, hits=4, misses=1, cache=1, tasks=0) + + assert ret == [2, 2, 2, 2, 2] + + +async def test_partialmethod_partial(check_lru, loop): + class Obj: + def __init__(self): + self.coro = alru_cache(partial(self._coro, 2)) + + async def __coro(self, val1, val2): + return val1 + val2 + + _coro = partialmethod(__coro, 1) + + obj = Obj() + + coros = [obj.coro() for _ in range(5)] + + check_lru(obj.coro, hits=0, misses=0, cache=0, tasks=0) + + ret = await asyncio.gather(*coros) + + check_lru(obj.coro, hits=4, misses=1, cache=1, tasks=0) + + assert ret == [3, 3, 3, 3, 3] diff --git a/tests/external/test_size.py b/tests/external/test_size.py new file mode 100644 index 0000000..c683c8e --- /dev/null +++ b/tests/external/test_size.py @@ -0,0 +1,58 @@ +import asyncio + +import pytest +from great_ai.external.async_lru import _make_key, alru_cache + +pytestmark = pytest.mark.asyncio + + +async def test_alru_cache_removing_lru_keys(check_lru, loop): + @alru_cache(maxsize=3) + async def coro(val): + return val + + key5 = _make_key((5,), {}, False) + key4 = _make_key((4,), {}, False) + key3 = _make_key((3,), {}, False) + key2 = _make_key((2,), {}, False) + key1 = _make_key((1,), {}, False) + + for i, v in enumerate([3, 4, 5]): + await coro(v) + check_lru(coro, hits=0, misses=i + 1, cache=i + 1, tasks=0, maxsize=3) + + check_lru(coro, hits=0, misses=3, cache=3, tasks=0, maxsize=3) + assert list(coro._cache) == [key3, key4, key5] + + for v in [3, 2, 1]: + await coro(v) + check_lru(coro, hits=1, misses=5, cache=3, tasks=0, maxsize=3) + assert list(coro._cache) == [key3, key2, key1] + + +async def test_alru_cache_none_max_size(check_lru, loop): + @alru_cache(maxsize=None) + async def coro(val): + return val + + inputs = [1, 2, 3, 4] * 2 + coros = [coro(v) for v in inputs] + + ret = await asyncio.gather(*coros) + + check_lru(coro, hits=4, misses=4, cache=4, tasks=0, maxsize=None) + assert ret == inputs + + +async def test_alru_cache_zero_max_size(check_lru, loop): + @alru_cache(maxsize=0) + async def coro(val): + return val + + inputs = [1, 2, 3, 4] * 2 + coros = [coro(v) for v in inputs] + + ret = await asyncio.gather(*coros) + + check_lru(coro, hits=0, misses=8, cache=0, tasks=0, maxsize=0) + assert ret == inputs diff --git a/tests/external/test_utils.py b/tests/external/test_utils.py new file mode 100644 index 0000000..aa52d5d --- /dev/null +++ b/tests/external/test_utils.py @@ -0,0 +1,19 @@ +from functools import partial + +from great_ai.external.async_lru import unpartial + + +def test_unpartial(): + def foo(): + pass + + assert unpartial(foo) is foo + + bar = partial(foo) + + assert unpartial(bar) is foo + + for _ in range(10): + bar = partial(bar) + + assert unpartial(bar) is foo