From c5f2b8c381ffcdba66a894970edf6dcd02919f42 Mon Sep 17 00:00:00 2001
From: Felipe Selmo <fselmo2@gmail.com>
Date: Tue, 28 Jun 2022 16:56:35 -0600
Subject: [PATCH] Add async tests for stalecheck middleware

---
 tests/core/middleware/test_stalecheck.py | 97 +++++++++++++++++++++++-
 1 file changed, 96 insertions(+), 1 deletion(-)

diff --git a/tests/core/middleware/test_stalecheck.py b/tests/core/middleware/test_stalecheck.py
index 806b2cf654..f0ba75b2ed 100644
--- a/tests/core/middleware/test_stalecheck.py
+++ b/tests/core/middleware/test_stalecheck.py
@@ -1,10 +1,12 @@
-
 import pytest
+import sys
 from unittest.mock import (
     Mock,
     patch,
 )
 
+import pytest_asyncio
+
 from web3.datastructures import (
     AttributeDict,
 )
@@ -14,6 +16,7 @@
 from web3.middleware.stalecheck import (
     StaleBlockchain,
     _is_fresh,
+    async_make_stalecheck_middleware,
 )
 
 
@@ -120,3 +123,95 @@ def test_stalecheck_adds_block_to_cache(request_middleware, allowable_delay):
         request_middleware('', [])
         assert fresh_spy.call_count == 3
         assert fresh_spy.call_args == ((block, allowable_delay), )
+
+
+# -- async -- #
+
+
+min_version = pytest.mark.skipif(
+    sys.version_info < (3, 8),
+    reason="AsyncMock requires python3.8 or higher"
+)
+
+
+@pytest_asyncio.fixture
+async def request_async_middleware(allowable_delay):
+    from unittest.mock import AsyncMock
+    middleware = await async_make_stalecheck_middleware(allowable_delay)
+    make_request, web3 = AsyncMock(), AsyncMock()
+    initialized = await middleware(make_request, web3)
+    # for easier mocking, later:
+    initialized.web3 = web3
+    initialized.make_request = make_request
+    return initialized
+
+
+@pytest.mark.asyncio
+@min_version
+async def test_async_stalecheck_pass(request_async_middleware):
+    with patch("web3.middleware.stalecheck._is_fresh", return_value=True):
+        method, params = object(), object()
+        await request_async_middleware(method, params)
+        request_async_middleware.make_request.assert_called_once_with(method, params)
+
+
+@pytest.mark.asyncio
+@min_version
+async def test_async_stalecheck_fail(request_async_middleware, now):
+    with patch("web3.middleware.stalecheck._is_fresh", return_value=False):
+        request_async_middleware.web3.eth.get_block.return_value = stub_block(now)
+        with pytest.raises(StaleBlockchain):
+            await request_async_middleware("", [])
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("rpc_method", ["eth_getBlockByNumber"])
+@min_version
+async def test_async_stalecheck_ignores_get_by_block_methods(
+    request_async_middleware, rpc_method
+):
+    # This is especially critical for get_block("latest") which would cause
+    # infinite recursion
+    with patch("web3.middleware.stalecheck._is_fresh", side_effect=[False, True]):
+        await request_async_middleware(rpc_method, [])
+        assert not request_async_middleware.web3.eth.get_block.called
+
+
+@pytest.mark.asyncio
+@min_version
+async def test_async_stalecheck_calls_is_fresh_with_empty_cache(
+    request_async_middleware, allowable_delay
+):
+    with patch(
+        "web3.middleware.stalecheck._is_fresh", side_effect=[False, True]
+    ) as fresh_spy:
+        block = object()
+        request_async_middleware.web3.eth.get_block.return_value = block
+        await request_async_middleware("", [])
+        cache_call, live_call = fresh_spy.call_args_list
+        assert cache_call[0] == (None, allowable_delay)
+        assert live_call[0] == (block, allowable_delay)
+
+
+@pytest.mark.asyncio
+@min_version
+async def test_async_stalecheck_adds_block_to_cache(
+    request_async_middleware, allowable_delay
+):
+    with patch(
+        "web3.middleware.stalecheck._is_fresh", side_effect=[False, True, True]
+    ) as fresh_spy:
+        block = object()
+        request_async_middleware.web3.eth.get_block.return_value = block
+
+        # cache miss
+        await request_async_middleware("", [])
+        cache_call, live_call = fresh_spy.call_args_list
+        assert fresh_spy.call_count == 2
+        assert cache_call == ((None, allowable_delay), )
+        assert live_call == ((block, allowable_delay), )
+
+        # cache hit
+        await request_async_middleware("", [])
+        assert fresh_spy.call_count == 3
+        assert fresh_spy.call_args == ((block, allowable_delay), )