From 14fcfd4c21f63cf9d3bc1ba09e041411a71c8a7a Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Tue, 12 Nov 2024 09:01:50 -0600
Subject: [PATCH] Adjust client GET read benchmarks to include chunked and
 content-length (#9827)

---
 tests/test_benchmarks_client.py | 104 +++++++++++++++++++++++++++++---
 1 file changed, 97 insertions(+), 7 deletions(-)

diff --git a/tests/test_benchmarks_client.py b/tests/test_benchmarks_client.py
index 8c141581266..7daddcf3db2 100644
--- a/tests/test_benchmarks_client.py
+++ b/tests/test_benchmarks_client.py
@@ -4,7 +4,7 @@
 
 from pytest_codspeed import BenchmarkFixture
 
-from aiohttp import web
+from aiohttp import hdrs, web
 from aiohttp.pytest_plugin import AiohttpClient
 
 
@@ -33,7 +33,7 @@ def _run() -> None:
         loop.run_until_complete(run_client_benchmark())
 
 
-def test_one_hundred_get_requests_with_2048_payload(
+def test_one_hundred_get_requests_with_2048_chunked_payload(
     loop: asyncio.AbstractEventLoop,
     aiohttp_client: AiohttpClient,
     benchmark: BenchmarkFixture,
@@ -43,7 +43,9 @@ def test_one_hundred_get_requests_with_2048_payload(
     payload = b"a" * 2048
 
     async def handler(request: web.Request) -> web.Response:
-        return web.Response(body=payload)
+        resp = web.Response(body=payload)
+        resp.enable_chunked_encoding()
+        return resp
 
     app = web.Application()
     app.router.add_route("GET", "/", handler)
@@ -60,7 +62,7 @@ def _run() -> None:
         loop.run_until_complete(run_client_benchmark())
 
 
-def test_one_hundred_get_requests_with_32768_payload(
+def test_one_hundred_get_requests_with_32768_chunked_payload(
     loop: asyncio.AbstractEventLoop,
     aiohttp_client: AiohttpClient,
     benchmark: BenchmarkFixture,
@@ -70,7 +72,9 @@ def test_one_hundred_get_requests_with_32768_payload(
     payload = b"a" * 32768
 
     async def handler(request: web.Request) -> web.Response:
-        return web.Response(body=payload)
+        resp = web.Response(body=payload)
+        resp.enable_chunked_encoding()
+        return resp
 
     app = web.Application()
     app.router.add_route("GET", "/", handler)
@@ -87,7 +91,7 @@ def _run() -> None:
         loop.run_until_complete(run_client_benchmark())
 
 
-def test_one_hundred_get_requests_with_1mib_payload(
+def test_one_hundred_get_requests_with_1mib_chunked_payload(
     loop: asyncio.AbstractEventLoop,
     aiohttp_client: AiohttpClient,
     benchmark: BenchmarkFixture,
@@ -97,7 +101,93 @@ def test_one_hundred_get_requests_with_1mib_payload(
     payload = b"a" * 1024**2
 
     async def handler(request: web.Request) -> web.Response:
-        return web.Response(body=payload)
+        resp = web.Response(body=payload)
+        resp.enable_chunked_encoding()
+        return resp
+
+    app = web.Application()
+    app.router.add_route("GET", "/", handler)
+
+    async def run_client_benchmark() -> None:
+        client = await aiohttp_client(app)
+        for _ in range(message_count):
+            resp = await client.get("/")
+            await resp.read()
+        await client.close()
+
+    @benchmark
+    def _run() -> None:
+        loop.run_until_complete(run_client_benchmark())
+
+
+def test_one_hundred_get_requests_with_2048_content_length_payload(
+    loop: asyncio.AbstractEventLoop,
+    aiohttp_client: AiohttpClient,
+    benchmark: BenchmarkFixture,
+) -> None:
+    """Benchmark 100 GET requests with a small payload of 2048 bytes."""
+    message_count = 100
+    payload = b"a" * 2048
+    headers = {hdrs.CONTENT_LENGTH: str(len(payload))}
+
+    async def handler(request: web.Request) -> web.Response:
+        return web.Response(body=payload, headers=headers)
+
+    app = web.Application()
+    app.router.add_route("GET", "/", handler)
+
+    async def run_client_benchmark() -> None:
+        client = await aiohttp_client(app)
+        for _ in range(message_count):
+            resp = await client.get("/")
+            await resp.read()
+        await client.close()
+
+    @benchmark
+    def _run() -> None:
+        loop.run_until_complete(run_client_benchmark())
+
+
+def test_one_hundred_get_requests_with_32768_content_length_payload(
+    loop: asyncio.AbstractEventLoop,
+    aiohttp_client: AiohttpClient,
+    benchmark: BenchmarkFixture,
+) -> None:
+    """Benchmark 100 GET requests with a payload of 32768 bytes."""
+    message_count = 100
+    payload = b"a" * 32768
+    headers = {hdrs.CONTENT_LENGTH: str(len(payload))}
+
+    async def handler(request: web.Request) -> web.Response:
+        return web.Response(body=payload, headers=headers)
+
+    app = web.Application()
+    app.router.add_route("GET", "/", handler)
+
+    async def run_client_benchmark() -> None:
+        client = await aiohttp_client(app)
+        for _ in range(message_count):
+            resp = await client.get("/")
+            await resp.read()
+        await client.close()
+
+    @benchmark
+    def _run() -> None:
+        loop.run_until_complete(run_client_benchmark())
+
+
+def test_one_hundred_get_requests_with_1mib_content_length_payload(
+    loop: asyncio.AbstractEventLoop,
+    aiohttp_client: AiohttpClient,
+    benchmark: BenchmarkFixture,
+) -> None:
+    """Benchmark 100 GET requests with a payload of 1MiB bytes."""
+    message_count = 100
+    payload = b"a" * 1024**2
+    headers = {hdrs.CONTENT_LENGTH: str(len(payload))}
+
+    async def handler(request: web.Request) -> web.Response:
+        return web.Response(body=payload, headers=headers)
 
     app = web.Application()
     app.router.add_route("GET", "/", handler)