From 659b1dd739a0ebfa29b10fb406381eacf89ae94d Mon Sep 17 00:00:00 2001
From: Jacob Tomlinson <jacobtomlinson@users.noreply.github.com>
Date: Fri, 14 Jun 2024 16:37:03 +0100
Subject: [PATCH] Allow setting timeouts on httpx client (#409)

---
 docs/client.md         | 30 ++++++++++++++++++++++++++++++
 kr8s/_api.py           | 14 +++++++++++++-
 kr8s/tests/test_api.py | 20 ++++++++++++++++++++
 3 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/docs/client.md b/docs/client.md
index ee6c7588..be389e02 100644
--- a/docs/client.md
+++ b/docs/client.md
@@ -122,6 +122,36 @@ async with api.call_api("GET", base="/version", version="") as r:
 print(version)
 ```
 
+## Timeouts
+
+All API calls are made by `httpx` under the hood. There may be cases where you want to manually control the timeout of these requests, especially when interacting with clusters under a heavy load or are some distance away.
+
+To set the timeout you can set the `.timeout` attribute on the API object. This value can be set to anything that the `timeout` keyword argument in `httpx` accepts.
+
+`````{tab-set}
+
+````{tab-item} Sync
+:sync: sync
+```python
+import kr8s
+
+api = kr8s.api()
+api.timeout = 10  # Set the default timeout for all calls to 10 seconds
+```
+````
+
+````{tab-item} Async
+:sync: async
+```python
+import kr8s
+
+api = await kr8s.asyncio.api()
+api.timeout = 10  # Set the default timeout for all calls to 10 seconds
+```
+````
+
+`````
+
 (client-caching)=
 
 ## Client caching
diff --git a/kr8s/_api.py b/kr8s/_api.py
index f89d29dc..b8323283 100644
--- a/kr8s/_api.py
+++ b/kr8s/_api.py
@@ -50,6 +50,7 @@ def __init__(self, **kwargs) -> None:
         self._kubeconfig = kwargs.get("kubeconfig")
         self._serviceaccount = kwargs.get("serviceaccount")
         self._session = None
+        self._timeout = None
         self.auth = KubeAuth(
             url=self._url,
             kubeconfig=self._kubeconfig,
@@ -75,6 +76,16 @@ async def f():
 
         return f().__await__()
 
+    @property
+    def timeout(self):
+        return self._timeout
+
+    @timeout.setter
+    def timeout(self, value):
+        self._timeout = value
+        if self._session:
+            self._session.timeout = value
+
     async def _create_session(self) -> None:
         headers = {"User-Agent": self.__version__, "content-type": "application/json"}
         if self.auth.token:
@@ -87,6 +98,7 @@ async def _create_session(self) -> None:
             base_url=self.auth.server,
             headers=headers,
             verify=await self.auth.ssl_context(),
+            timeout=self._timeout,
         )
 
     def _construct_url(
@@ -173,7 +185,7 @@ async def call_api(
                     continue
                 else:
                     raise
-            except httpx.ReadTimeout as e:
+            except httpx.TimeoutException as e:
                 raise APITimeoutError(
                     "Timeout while waiting for the Kubernetes API server"
                 ) from e
diff --git a/kr8s/tests/test_api.py b/kr8s/tests/test_api.py
index e3ecef33..f98cc87d 100644
--- a/kr8s/tests/test_api.py
+++ b/kr8s/tests/test_api.py
@@ -8,6 +8,7 @@
 
 import kr8s
 import kr8s.asyncio
+from kr8s._exceptions import APITimeoutError
 from kr8s.asyncio.objects import Pod, Table
 
 
@@ -291,3 +292,22 @@ async def test_api_resources_cache(caplog):
     assert caplog.text.count('/apis/ "HTTP/1.1 200 OK"') == 1
     await api.api_resources()
     assert caplog.text.count('/apis/ "HTTP/1.1 200 OK"') == 1
+
+
+async def test_api_timeout():
+    from httpx import Timeout
+
+    api = await kr8s.asyncio.api()
+    api.timeout = 10
+    await api.version()
+    assert api._session.timeout.read == 10
+    api.timeout = 20
+    await api.version()
+    assert api._session.timeout.read == 20
+    api.timeout = Timeout(30)
+    await api.version()
+    assert api._session.timeout.read == 30
+
+    api.timeout = 0.00001
+    with pytest.raises(APITimeoutError):
+        await api.version()