From 293d884731df97fc96d9efbbbf13ababde8d3db4 Mon Sep 17 00:00:00 2001
From: Adam Hopkins <adam@amhopkins.com>
Date: Thu, 22 Sep 2022 10:33:52 +0300
Subject: [PATCH] Add request counter (#129)

---
 sanic_ext/__init__.py               |  2 ++
 sanic_ext/extras/request.py         | 46 +++++++++++++++++++++++++++++
 setup.cfg                           |  2 +-
 tests/conftest.py                   |  1 +
 tests/extra/test_request_counted.py | 42 ++++++++++++++++++++++++++
 5 files changed, 92 insertions(+), 1 deletion(-)
 create mode 100644 sanic_ext/extras/request.py
 create mode 100644 tests/extra/test_request_counted.py

diff --git a/sanic_ext/__init__.py b/sanic_ext/__init__.py
index 71eee0e..b9ba15d 100644
--- a/sanic_ext/__init__.py
+++ b/sanic_ext/__init__.py
@@ -6,6 +6,7 @@
 from sanic_ext.extensions.http.cors import cors
 from sanic_ext.extensions.openapi import openapi
 from sanic_ext.extensions.templating.render import render
+from sanic_ext.extras.request import CountedRequest
 from sanic_ext.extras.serializer.decorator import serializer
 from sanic_ext.extras.validation.decorator import validate
 
@@ -13,6 +14,7 @@
 
 __all__ = [
     "Config",
+    "CountedRequest",
     "Extend",
     "Extension",
     "cors",
diff --git a/sanic_ext/extras/request.py b/sanic_ext/extras/request.py
new file mode 100644
index 0000000..f60daf2
--- /dev/null
+++ b/sanic_ext/extras/request.py
@@ -0,0 +1,46 @@
+from itertools import count
+
+from sanic import Request, Sanic
+from sanic.compat import Header
+from sanic.models.protocol_types import TransportProtocol
+
+
+class CountedRequest(Request):
+    __slots__ = ()
+
+    _counter = count()
+    count = next(_counter)
+
+    def __init__(
+        self,
+        url_bytes: bytes,
+        headers: Header,
+        version: str,
+        method: str,
+        transport: TransportProtocol,
+        app: Sanic,
+        head: bytes = b"",
+        stream_id: int = 0,
+    ):
+        super().__init__(
+            url_bytes,
+            headers,
+            version,
+            method,
+            transport,
+            app,
+            head,
+            stream_id,
+        )
+        self.__class__._increment()
+        if hasattr(self.app, "multiplexer"):
+            self.app.multiplexer.state["request_count"] = self.__class__.count
+
+    @classmethod
+    def _increment(cls):
+        cls.count = next(cls._counter)
+
+    @classmethod
+    def reset_count(cls):
+        cls._counter = count()
+        cls.count = next(cls._counter)
diff --git a/setup.cfg b/setup.cfg
index 4f5ca27..3b2eb3e 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -28,7 +28,7 @@ sanic_ext =
 
 [options.extras_require]
 test =
-    sanic_testing==22.9.0b1
+    sanic_testing>=22.9.0b1
     coverage
     pytest
     pytest-cov
diff --git a/tests/conftest.py b/tests/conftest.py
index f4a1f0e..1f022c6 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -22,6 +22,7 @@ def reset_globals():
     OperationStore.reset()
     Extend.reset()
     Extension.reset()
+    Sanic._app_registry.clear()
 
 
 @pytest.fixture(autouse=True)
diff --git a/tests/extra/test_request_counted.py b/tests/extra/test_request_counted.py
new file mode 100644
index 0000000..5f26e04
--- /dev/null
+++ b/tests/extra/test_request_counted.py
@@ -0,0 +1,42 @@
+from unittest.mock import Mock
+
+import pytest
+from sanic import Sanic
+from sanic.compat import Header
+from sanic.response import json
+from sanic_testing.reusable import ReusableClient
+
+from sanic_ext import CountedRequest
+
+
+@pytest.fixture(autouse=True)
+def reset_counter():
+    yield
+    CountedRequest.reset_count()
+
+
+def test_counter_increments(app: Sanic):
+    app.request_class = CountedRequest
+
+    @app.get("/")
+    async def handler(request: CountedRequest):
+        return json({"count": request.count})
+
+    @app.get("/info")
+    async def info(request: CountedRequest):
+        return json({"state": request.app.m.state})
+
+    with ReusableClient(app) as client:
+        for i in range(1, 10):
+            _, response = client.get("/")
+            assert response.json["count"] == i
+
+
+def test_counter_increment_on_state(app: Sanic):
+    mock = Mock()
+    mock.state = {}
+    app.multiplexer = mock
+
+    for i in range(1, 10):
+        CountedRequest(b"/", Header({}), "", "", Mock(), app)
+        assert CountedRequest.count == i