feat: Implement asynchronous timeout context manager (#1569)

* feat: implement async timeout guard

* add docstring

* clean whitespace

* update import file name

* add missing return statement

* update test cases

* update test cases

* include underlying timeout exception in trace

* avoid the cost of actual time

google/auth/aio/transport/aiohttp.py

@@ -14,3 +14,54 @@
 
 """Transport adapter for Asynchronous HTTP Requests.
 """
+
+
+from google.auth.exceptions import TimeoutError
+
+import asyncio
+import time
+from contextlib import asynccontextmanager
+
+
+@asynccontextmanager
+async def timeout_guard(timeout):
+    """
+    timeout_guard is an asynchronous context manager to apply a timeout to an asynchronous block of code.
+
+    Args:
+        timeout (float): The time in seconds before the context manager times out.
+
+    Raises:
+        google.auth.exceptions.TimeoutError: If the code within the context exceeds the provided timeout.
+
+    Usage:
+        async with timeout_guard(10) as with_timeout:
+            await with_timeout(async_function())
+    """
+    start = time.monotonic()
+    total_timeout = timeout
+
+    def _remaining_time():
+        elapsed = time.monotonic() - start
+        remaining = total_timeout - elapsed
+        if remaining <= 0:
+            raise TimeoutError(
+                f"Context manager exceeded the configured timeout of {total_timeout}s."
+            )
+        return remaining
+
+    async def with_timeout(coro):
+        try:
+            remaining = _remaining_time()
+            response = await asyncio.wait_for(coro, remaining)
+            return response
+        except (asyncio.TimeoutError, TimeoutError) as e:
+            raise TimeoutError(
+                f"The operation {coro} exceeded the configured timeout of {total_timeout}s."
+            ) from e
+
+    try:
+        yield with_timeout
+
+    finally:
+        _remaining_time()

google/auth/exceptions.py

@@ -98,3 +98,7 @@ class InvalidType(DefaultCredentialsError, TypeError):
 
 class OSError(DefaultCredentialsError, EnvironmentError):
     """Used to wrap EnvironmentError(OSError after python3.3)."""
+
+
+class TimeoutError(GoogleAuthError):
+    """Used to indicate a timeout error occurred during an HTTP request."""

tests/transport/aio/test_aiohttp.py

@@ -13,4 +13,71 @@
 # limitations under the License.
 
 import google.auth.aio.transport.aiohttp as auth_aiohttp
-import pytest # type: ignore
+import pytest  # type: ignore
+import asyncio
+from google.auth.exceptions import TimeoutError
+from unittest.mock import patch
+
+
+@pytest.fixture
+async def simple_async_task():
+    return True
+
+
+class TestTimeoutGuard(object):
+    default_timeout = 1
+
+    def make_timeout_guard(self, timeout):
+        return auth_aiohttp.timeout_guard(timeout)
+
+    @pytest.mark.asyncio
+    async def test_timeout_with_simple_async_task_within_bounds(
+        self, simple_async_task
+    ):
+        task = False
+        with patch("time.monotonic", side_effect=[0, 0.25, 0.75]):
+            with patch("asyncio.wait_for", lambda coro, timeout: coro):
+                async with self.make_timeout_guard(
+                    timeout=self.default_timeout
+                ) as with_timeout:
+                    task = await with_timeout(simple_async_task)
+
+        # Task succeeds.
+        assert task is True
+
+    @pytest.mark.asyncio
+    async def test_timeout_with_simple_async_task_out_of_bounds(
+        self, simple_async_task
+    ):
+        task = False
+        with patch("time.monotonic", side_effect=[0, 1, 1]):
+            with patch("asyncio.wait_for", lambda coro, timeout: coro):
+                with pytest.raises(TimeoutError) as exc:
+                    async with self.make_timeout_guard(
+                        timeout=self.default_timeout
+                    ) as with_timeout:
+                        task = await with_timeout(simple_async_task)
+
+        # Task does not succeed and the context manager times out i.e. no remaining time left.
+        assert task is False
+        assert exc.match(
+            f"Context manager exceeded the configured timeout of {self.default_timeout}s."
+        )
+
+    @pytest.mark.asyncio
+    async def test_timeout_with_async_task_timing_out_before_context(
+        self, simple_async_task
+    ):
+        task = False
+        with pytest.raises(TimeoutError) as exc:
+            async with self.make_timeout_guard(
+                timeout=self.default_timeout
+            ) as with_timeout:
+                with patch("asyncio.wait_for", side_effect=asyncio.TimeoutError):
+                    task = await with_timeout(simple_async_task)
+
+        # Task does not complete i.e. the operation times out.
+        assert task is False
+        assert exc.match(
+            f"The operation {simple_async_task} exceeded the configured timeout of {self.default_timeout}s."
+        )