Optimize

Author: SlavaSkvortsov

pytest_httpx_blockage/blockage.py

@@ -1,36 +1,62 @@
 from contextlib import contextmanager
-from typing import Any, Generator
+from typing import Any, Dict, Generator, Tuple
 from unittest.mock import patch
 
+from httpcore import AsyncByteStream, SyncByteStream
 from httpcore._async.connection import AsyncHTTPConnection
 from httpcore._sync.connection import SyncHTTPConnection
-from httpx import URL
+from httpcore._types import URL, Headers
 
+from pytest_httpx_blockage.contextvar import is_blockage_enabled
 from pytest_httpx_blockage.exceptions import RequestBlockageException
 
+base_request_sync = SyncHTTPConnection.request
+base_request_async = AsyncHTTPConnection.arequest
+
 
 def side_effect(
+    self: SyncHTTPConnection,
     method: bytes,
     url: URL,
     *args: Any,
     **kwargs: Any,
-) -> None:
-    raise RequestBlockageException(f'Unmocked "{method.decode()}" request to host="{url}"')
+) -> Tuple[int, Headers, SyncByteStream, Dict[Any, Any]]:
+    if is_blockage_enabled.get():
+        raise RequestBlockageException(f'Unmocked "{method.decode()}" request to host="{url}"')
+    else:
+        return base_request_sync(
+            self,
+            method,
+            url,
+            *args,
+            **kwargs,
+        )
 
 
 async def async_side_effect(
+    self: AsyncHTTPConnection,
     method: bytes,
     url: URL,
     *args: Any,
     **kwargs: Any,
-) -> None:
-    side_effect(method=method, url=url)
+) -> Tuple[int, Headers, AsyncByteStream, Dict[Any, Any]]:
+    if is_blockage_enabled.get():
+        raise RequestBlockageException(f'Unmocked "{method.decode()}" request to host="{url}"')
+    else:
+        return await base_request_async(
+            self,
+            method,
+            url,
+            *args,
+            **kwargs,
+        )
 
 
 @contextmanager
 def blockage() -> Generator[None, None, None]:
-    patch_sync = patch.object(SyncHTTPConnection, 'request')
-    patch_async = patch.object(AsyncHTTPConnection, 'arequest')
+    patch_sync = patch.object(SyncHTTPConnection, 'request', autospec=True)
+    patch_async = patch.object(AsyncHTTPConnection, 'arequest', autospec=True)
+
     with patch_sync as mocked_sync, patch_async as mocked_async:
         mocked_sync.side_effect = side_effect
         mocked_async.side_effect = async_side_effect

pytest_httpx_blockage/contextvar.py

@@ -0,0 +1,3 @@
+from contextvars import ContextVar
+
+is_blockage_enabled: ContextVar[bool] = ContextVar('is_blockage_enabled', default=True)

pytest_httpx_blockage/plugin.py

@@ -5,6 +5,7 @@
 from _pytest.fixtures import SubRequest
 
 from pytest_httpx_blockage.blockage import blockage
+from pytest_httpx_blockage.contextvar import is_blockage_enabled
 
 
 def pytest_addoption(parser: Parser) -> None:
@@ -34,7 +35,7 @@ def pytest_addoption(parser: Parser) -> None:
     )
 
 
-@pytest.fixture(autouse=True)
+@pytest.fixture(scope='session', autouse=True)
 def _blockage(request: SubRequest) -> Generator[None, None, None]:
     config = request.config
     blockage_enabled = config.getini('blockage-httpx') or config.getoption('--blockage-httpx')
@@ -43,11 +44,19 @@ def _blockage(request: SubRequest) -> Generator[None, None, None]:
         yield
         return
 
+    with blockage():
+        yield
+
+
+@pytest.fixture(autouse=True)
+def _blockage_disable_mark(request: SubRequest) -> Generator[None, None, None]:
+    config = request.config
     disable_mark_name = config.getini('disable-blockage-mark') or config.getoption('--disable-blockage-mark')
     mark = request.node.get_closest_marker(disable_mark_name or 'integration')
     if mark is not None:
-        yield
-        return
+        is_blockage_enabled.set(False)
 
-    with blockage():
+    try:
         yield
+    finally:
+        is_blockage_enabled.set(True)