test: cover end-to-end retry on HTTP 504

Adds TestRetryOn504EndToEnd which drives a real request through
RetryableHTTPTransport / RetryableAsyncHTTPTransport and asserts the
underlying call fires max_retries times before the final 504 response is
returned. The existing tests only covered configuration/wiring, not the
actual retry loop behavior against a real httpx.Response.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: ionut-mihalache-uipath

tests/core/features/test_retry.py

@@ -3,6 +3,10 @@
 import logging
 from unittest.mock import MagicMock, patch
 
+import httpx
+import pytest
+
+from uipath.llm_client.httpx_client import UiPathHttpxAsyncClient, UiPathHttpxClient
 from uipath.llm_client.utils.exceptions import (
     UiPathBadGatewayError,
     UiPathGatewayTimeoutError,
@@ -19,6 +23,8 @@
     RetryConfig,
 )
 
+_NO_DELAY_CONFIG: RetryConfig = {"initial_delay": 0, "max_delay": 0, "jitter": 0}
+
 
 class TestDefaultRetryOnExceptions:
     """Pins the default retry set so HTTP 408/429/502/503/504/529 stay retryable."""
@@ -263,3 +269,79 @@ def test_with_logger_adds_before_sleep(self):
         result = _build_retryer(max_retries=2, retry_config=None, logger=logger)
         assert result is not None
         assert result.before_sleep is not None
+
+
+class TestRetryOn504EndToEnd:
+    """End-to-end coverage that a 504 response actually triggers the retry loop.
+
+    The other test classes verify the *configuration* (default exception set,
+    retryer wiring, wait math). These tests drive a real request through
+    ``RetryableHTTPTransport.handle_request`` / ``handle_async_request`` and
+    assert the underlying call fires ``max_retries`` times before the final 504
+    response is returned.
+    """
+
+    def test_sync_504_retries_max_retries_times_then_returns_response(self):
+        calls = 0
+
+        def fake_504(self: httpx.HTTPTransport, request: httpx.Request) -> httpx.Response:
+            nonlocal calls
+            calls += 1
+            return httpx.Response(504, content=b"{}", request=request)
+
+        client = UiPathHttpxClient(
+            base_url="https://example.com",
+            max_retries=3,
+            retry_config=_NO_DELAY_CONFIG,
+        )
+        try:
+            with patch.object(httpx.HTTPTransport, "handle_request", fake_504):
+                response = client.post("/anything", json={})
+        finally:
+            client.close()
+
+        assert calls == 3
+        assert response.status_code == 504
+
+    def test_sync_max_retries_zero_makes_exactly_one_call(self):
+        calls = 0
+
+        def fake_504(self: httpx.HTTPTransport, request: httpx.Request) -> httpx.Response:
+            nonlocal calls
+            calls += 1
+            return httpx.Response(504, content=b"{}", request=request)
+
+        client = UiPathHttpxClient(base_url="https://example.com", max_retries=0)
+        try:
+            with patch.object(httpx.HTTPTransport, "handle_request", fake_504):
+                response = client.post("/anything", json={})
+        finally:
+            client.close()
+
+        assert calls == 1
+        assert response.status_code == 504
+
+    @pytest.mark.asyncio
+    async def test_async_504_retries_max_retries_times_then_returns_response(self):
+        calls = 0
+
+        async def fake_504(
+            self: httpx.AsyncHTTPTransport, request: httpx.Request
+        ) -> httpx.Response:
+            nonlocal calls
+            calls += 1
+            return httpx.Response(504, content=b"{}", request=request)
+
+        client = UiPathHttpxAsyncClient(
+            base_url="https://example.com",
+            max_retries=3,
+            retry_config=_NO_DELAY_CONFIG,
+        )
+        try:
+            with patch.object(httpx.AsyncHTTPTransport, "handle_async_request", fake_504):
+                response = await client.post("/anything", json={})
+        finally:
+            await client.aclose()
+
+        assert calls == 3
+        assert response.status_code == 504