From 4404b5da6eb2bfb512c6c583f7bbcf92c3fa32f2 Mon Sep 17 00:00:00 2001
From: Iurii Pliner <yury.pliner@gmail.com>
Date: Tue, 22 Oct 2024 17:06:57 +0100
Subject: [PATCH] Do not retry low timeout response (#276)

---
 aio_request/pipeline.py           | 11 +++++++++--
 tests/conftest.py                 |  2 +-
 tests/test_parallel_strategy.py   | 19 ++++++++++++++++++-
 tests/test_sequential_strategy.py | 19 ++++++++++++++++++-
 4 files changed, 46 insertions(+), 5 deletions(-)

diff --git a/aio_request/pipeline.py b/aio_request/pipeline.py
index 45ce749..2c0bf0b 100644
--- a/aio_request/pipeline.py
+++ b/aio_request/pipeline.py
@@ -47,11 +47,18 @@ async def execute(
 
 
 class LowTimeoutModule(RequestModule):
-    __slots__ = ("__low_timeout_threshold",)
+    __slots__ = ("__low_timeout_threshold", "__timeout_response")
 
     def __init__(self, low_timeout_threshold: float):
         self.__low_timeout_threshold = low_timeout_threshold
 
+        headers = multidict.CIMultiDict[str]()
+        headers[Header.X_DO_NOT_RETRY] = "1"
+        self.__timeout_response = EmptyResponse(
+            status=408,
+            headers=multidict.CIMultiDictProxy[str](headers),
+        )
+
     async def execute(
         self,
         next: NextModuleFunc,
@@ -62,7 +69,7 @@ async def execute(
         priority: Priority,
     ) -> ClosableResponse:
         if deadline.expired or deadline.timeout < self.__low_timeout_threshold:
-            return EmptyResponse(status=408)
+            return self.__timeout_response
 
         return await next(endpoint, request, deadline, priority)
 
diff --git a/tests/conftest.py b/tests/conftest.py
index a8edd11..07656d9 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -17,7 +17,7 @@
 @dataclasses.dataclass(frozen=True)
 class FakeResponseConfiguration:
     status: int
-    delay_seconds: float
+    delay_seconds: float = 0
 
 
 class FakeTransport(aio_request.Transport):
diff --git a/tests/test_parallel_strategy.py b/tests/test_parallel_strategy.py
index baa836b..a4be083 100644
--- a/tests/test_parallel_strategy.py
+++ b/tests/test_parallel_strategy.py
@@ -3,7 +3,24 @@
 from .conftest import FakeResponseConfiguration, FakeTransport
 
 
-async def test_timeout_because_of_expiration():
+async def test_timeout_due_to_low_timeout():
+    client = aio_request.setup(
+        transport=FakeTransport([FakeResponseConfiguration(status=200)]),
+        endpoint="http://service.com",
+    )
+    deadline = aio_request.Deadline.from_timeout(0.004)
+    response_ctx = client.request(
+        aio_request.get("hello"),
+        deadline=deadline,
+        strategy=aio_request.parallel_strategy(),
+    )
+    async with response_ctx as response:
+        assert response.status == 408
+        assert aio_request.Header.X_DO_NOT_RETRY in response.headers
+        assert not deadline.expired
+
+
+async def test_timeout_due_to_expiration():
     client = aio_request.setup(
         transport=FakeTransport(
             [
diff --git a/tests/test_sequential_strategy.py b/tests/test_sequential_strategy.py
index da26b56..2d3a85a 100644
--- a/tests/test_sequential_strategy.py
+++ b/tests/test_sequential_strategy.py
@@ -6,7 +6,24 @@
 from .conftest import FakeResponseConfiguration, FakeTransport
 
 
-async def test_timeout_because_of_expiration():
+async def test_timeout_due_to_low_timeout():
+    client = aio_request.setup(
+        transport=FakeTransport([FakeResponseConfiguration(status=200)]),
+        endpoint="http://service.com",
+    )
+    deadline = aio_request.Deadline.from_timeout(0.004)
+    response_ctx = client.request(
+        aio_request.get("hello"),
+        deadline=deadline,
+        strategy=aio_request.sequential_strategy(attempts_count=3, delays_provider=aio_request.linear_backoff_delays()),
+    )
+    async with response_ctx as response:
+        assert response.status == 408
+        assert aio_request.Header.X_DO_NOT_RETRY in response.headers
+        assert not deadline.expired
+
+
+async def test_timeout_due_to_expiration():
     client = aio_request.setup(
         transport=FakeTransport([FakeResponseConfiguration(status=200, delay_seconds=5)]),
         endpoint="http://service.com",