Define an alternative can_handle logic by passing a callable. (#306)

Author: mindflayer

README.rst

@@ -225,6 +225,37 @@ It's very important that we test non-happy paths.
         with self.assertRaises(requests.exceptions.ConnectionError):
             requests.get(url)
 
+Example of how to mock a call with a custom `can_handle` function
+=================================================================
+.. code-block:: python
+
+    import json
+
+    from mocket import mocketize
+    from mocket.mocks.mockhttp import Entry
+    import requests
+
+    @mocketize
+    def test_can_handle():
+        Entry.single_register(
+            Entry.GET,
+            url,
+            body=json.dumps({"message": "Nope... not this time!"}),
+            headers={"content-type": "application/json"},
+            can_handle_fun=lambda path, qs_dict: path == "/ip" and qs_dict,
+        )
+        Entry.single_register(
+            Entry.GET,
+            url,
+            body=json.dumps({"message": "There you go!"}),
+            headers={"content-type": "application/json"},
+            can_handle_fun=lambda path, qs_dict: path == "/ip" and not qs_dict,
+        )
+        resp = requests.get("https://httpbin.org/ip")
+        assert resp.status_code == 200
+        assert resp.json() == {"message": "There you go!"}
+
+
 Example of how to record real socket traffic
 ============================================
 
@@ -251,10 +282,12 @@ You probably know what *VCRpy* is capable of, that's the *mocket*'s way of achie
 
 HTTPretty compatibility layer
 =============================
-Mocket HTTP mock can work as *HTTPretty* replacement for many different use cases. Two main features are missing:
+Mocket HTTP mock can work as *HTTPretty* replacement for many different use cases. Two main features are missing, or better said, are implemented differently:
+
+- URL entries containing regular expressions, *Mocket* implements `can_handle_fun` which is way simpler to use and more powerful;
+- response body from functions (used mostly to fake errors, *Mocket* accepts an `exception` instead).
 
-- URL entries containing regular expressions;
-- response body from functions (used mostly to fake errors, *mocket* doesn't need to do it this way).
+Both features are documented above.
 
 Two features which are against the Zen of Python, at least imho (*mindflayer*), but of course I am open to call it into question.
 

mocket/__init__.py

@@ -31,4 +31,4 @@
     "FakeSSLContext",
 )
 
-__version__ = "3.13.10"
+__version__ = "3.13.11"

mocket/mocks/mockhttp.py

@@ -2,6 +2,7 @@
 import time
 from functools import cached_property
 from http.server import BaseHTTPRequestHandler
+from typing import Callable, Optional
 from urllib.parse import parse_qs, unquote, urlsplit
 
 from h11 import SERVER, Connection, Data
@@ -82,9 +83,7 @@ def __init__(self, body="", status=200, headers=None):
         self.status = status
 
         self.set_base_headers()
-
-        if headers is not None:
-            self.set_extra_headers(headers)
+        self.set_extra_headers(headers)
 
         self.data = self.get_protocol_data() + self.body
 
@@ -142,9 +141,19 @@ class Entry(MocketEntry):
     request_cls = Request
     response_cls = Response
 
-    default_config = {"match_querystring": True}
+    default_config = {"match_querystring": True, "can_handle_fun": None}
+    _can_handle_fun: Optional[Callable] = None
+
+    def __init__(
+        self,
+        uri,
+        method,
+        responses,
+        match_querystring: bool = True,
+        can_handle_fun: Optional[Callable] = None,
+    ):
+        self._can_handle_fun = can_handle_fun if can_handle_fun else self._can_handle
 
-    def __init__(self, uri, method, responses, match_querystring: bool = True):
         uri = urlsplit(uri)
 
         port = uri.port
@@ -177,6 +186,18 @@ def collect(self, data):
 
         return consume_response
 
+    def _can_handle(self, path: str, qs_dict: dict) -> bool:
+        """
+        The default can_handle function, which checks if the path match,
+        and if match_querystring is True, also checks if the querystring matches.
+        """
+        can_handle = path == self.path
+        if self._match_querystring:
+            can_handle = can_handle and qs_dict == parse_qs(
+                self.query, keep_blank_values=True
+            )
+        return can_handle
+
     def can_handle(self, data):
         r"""
         >>> e = Entry('http://www.github.com/?bar=foo&foobar', Entry.GET, (Response(b'<html/>'),))
@@ -192,13 +213,12 @@ def can_handle(self, data):
         except ValueError:
             return self is getattr(Mocket, "_last_entry", None)
 
-        uri = urlsplit(path)
-        can_handle = uri.path == self.path and method == self.method
-        if self._match_querystring:
-            kw = dict(keep_blank_values=True)
-            can_handle = can_handle and parse_qs(uri.query, **kw) == parse_qs(
-                self.query, **kw
-            )
+        _request = urlsplit(path)
+
+        can_handle = method == self.method and self._can_handle_fun(
+            _request.path, parse_qs(_request.query, keep_blank_values=True)
+        )
+
         if can_handle:
             Mocket._last_entry = self
         return can_handle
@@ -249,8 +269,27 @@ def single_register(
         headers=None,
         exception=None,
         match_querystring=True,
+        can_handle_fun=None,
         **config,
     ):
+        """
+        A helper method to register a single Response for a given URI and method.
+        Instead of passing a list of Response objects, you can just pass the response
+        parameters directly.
+
+        Args:
+            method (str): The HTTP method (e.g., 'GET', 'POST').
+            uri (str): The URI to register the response for.
+            body (str, optional): The body of the response. Defaults to an empty string.
+            status (int, optional): The HTTP status code. Defaults to 200.
+            headers (dict, optional): A dictionary of headers to include in the response. Defaults to None.
+            exception (Exception, optional): An exception to raise instead of returning a response. Defaults to None.
+            match_querystring (bool, optional): Whether to match the querystring in the URI. Defaults to True.
+            can_handle_fun (Callable, optional): A custom function to determine if the Entry can handle a request.
+                Defaults to None. If None, the default matching logic is used. The function should accept two parameters:
+                path (str), and querystring params (dict), and return a boolean. Method is matched before the function call.
+            **config: Additional configuration options.
+        """
         response = (
             exception
             if exception
@@ -262,5 +301,6 @@ def single_register(
             uri,
             response,
             match_querystring=match_querystring,
+            can_handle_fun=can_handle_fun,
             **config,
         )

mocket/plugins/httpretty/__init__.py

@@ -139,6 +139,4 @@ def __getattr__(self, name):
     "HEAD",
     "PATCH",
     "register_uri",
-    "str",
-    "bytes",
 )

tests/test_http.py

@@ -455,3 +455,30 @@ def test_mocket_with_no_path(self):
         response = urlopen("http://httpbin.local/")
         self.assertEqual(response.code, 202)
         self.assertEqual(Mocket._entries[("httpbin.local", 80)][0].path, "/")
+
+    @mocketize
+    def test_can_handle(self):
+        Entry.single_register(
+            Entry.POST,
+            "http://testme.org/foobar",
+            body=json.dumps({"message": "Spooky!"}),
+            match_querystring=False,
+        )
+        Entry.single_register(
+            Entry.GET,
+            "http://testme.org/",
+            body=json.dumps({"message": "Gotcha!"}),
+            can_handle_fun=lambda p, q: p.endswith("/foobar") and "a" in q,
+        )
+        Entry.single_register(
+            Entry.GET,
+            "http://testme.org/foobar",
+            body=json.dumps({"message": "Missed!"}),
+            match_querystring=False,
+        )
+        response = requests.get("http://testme.org/foobar?a=1")
+        self.assertEqual(response.status_code, 200)
+        self.assertEqual(response.json(), {"message": "Gotcha!"})
+        response = requests.get("http://testme.org/foobar?b=2")
+        self.assertEqual(response.status_code, 200)
+        self.assertEqual(response.json(), {"message": "Missed!"})

tests/test_https.py

@@ -91,3 +91,24 @@ def test_raise_exception_from_single_register():
     Entry.single_register(Entry.GET, url, exception=OSError())
     with pytest.raises(requests.exceptions.ConnectionError):
         requests.get(url)
+
+
+@mocketize
+def test_can_handle():
+    Entry.single_register(
+        Entry.GET,
+        "https://httpbin.org",
+        body=json.dumps({"message": "Nope... not this time!"}),
+        headers={"content-type": "application/json"},
+        can_handle_fun=lambda path, qs_dict: path == "/ip" and qs_dict,
+    )
+    Entry.single_register(
+        Entry.GET,
+        "https://httpbin.org",
+        body=json.dumps({"message": "There you go!"}),
+        headers={"content-type": "application/json"},
+        can_handle_fun=lambda path, qs_dict: path == "/ip" and not qs_dict,
+    )
+    resp = requests.get("https://httpbin.org/ip")
+    assert resp.status_code == 200
+    assert resp.json() == {"message": "There you go!"}

tests/test_httpx.py

@@ -194,3 +194,22 @@ async def test_httpx_fixture(httpx_client):
         response = await client.get(url)
 
         assert response.json() == data
+
+
+@pytest.mark.asyncio
+async def test_httpx_fixture_with_can_handle_fun(httpx_client):
+    url = "https://foo.bar/barfoo"
+    data = {"message": "Gotcha!"}
+
+    Entry.single_register(
+        Entry.GET,
+        "https://foo.bar",
+        body=json.dumps(data),
+        headers={"content-type": "application/json"},
+        can_handle_fun=lambda p, q: p.endswith("foo"),
+    )
+
+    async with httpx_client as client:
+        response = await client.get(url)
+
+        assert response.json() == data