From 2a70d455338192efcc77026e21c894ef5639b4a8 Mon Sep 17 00:00:00 2001
From: Dmitry Volodin <dv@nocproject.org>
Date: Wed, 28 Feb 2024 16:36:03 +0100
Subject: [PATCH] Timeouts

---
 src/async_client/client.rs    |  9 ++++++++-
 src/gufo/http/_fast.pyi       |  2 ++
 src/gufo/http/async_client.py | 10 ++++++++++
 tests/test_async_client.py    | 11 ++++++++++-
 tests/util.py                 |  1 +
 5 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/src/async_client/client.rs b/src/async_client/client.rs
index d3a5f6f..2dc7592 100644
--- a/src/async_client/client.rs
+++ b/src/async_client/client.rs
@@ -17,6 +17,7 @@ use reqwest::{
     Method,
 };
 use std::collections::HashMap;
+use std::time::Duration;
 
 #[pyclass(module = "gufo.http.async_client")]
 pub struct AsyncClient {
@@ -28,6 +29,8 @@ impl AsyncClient {
     #[new]
     fn new(
         validate_cert: bool,
+        connect_timeout: u64,
+        timeout: u64,
         max_redirect: Option<usize>,
         headers: Option<HashMap<&str, &[u8]>>,
         compression: Option<u8>,
@@ -61,10 +64,14 @@ impl AsyncClient {
                 builder = builder.brotli(true);
             }
         }
-        // Disable certificate validation
+        // Disable certificate validation when necessary
         if !validate_cert {
             builder = builder.danger_accept_invalid_certs(true);
         }
+        // Set timeouts
+        builder = builder
+            .connect_timeout(Duration::from_nanos(connect_timeout))
+            .timeout(Duration::from_nanos(timeout));
         //
         let client = builder
             .build()
diff --git a/src/gufo/http/_fast.pyi b/src/gufo/http/_fast.pyi
index de47cd3..306c168 100644
--- a/src/gufo/http/_fast.pyi
+++ b/src/gufo/http/_fast.pyi
@@ -40,6 +40,8 @@ class AsyncClient(object):
     def __init__(
         self: "AsyncClient",
         validate_cert: bool,
+        connect_timeout_ns: int,
+        timeout_ns: int,
         max_redirects: Optional[int],
         headers: Optional[Dict[str, bytes]],
         compression: Optional[int],
diff --git a/src/gufo/http/async_client.py b/src/gufo/http/async_client.py
index 3ea2e40..f16565b 100644
--- a/src/gufo/http/async_client.py
+++ b/src/gufo/http/async_client.py
@@ -16,6 +16,9 @@
 from .util import merge_dict
 
 MAX_REDIRECTS = 10
+DEFAULT_CONNECT_TIMEOUT = 30.0
+DEFAULT_TIMEOUT = 3600.0
+NS = 1_000_000_000.0
 
 
 class HttpClient(object):
@@ -34,19 +37,26 @@ class HttpClient(object):
             Set to `None` to disable compression support.
         validate_cert: Set to `False` to disable TLS certificate
             validation.
+        connect_timeout: Timeout to establish connection, in seconds.
+        timeout: Request timeout, in seconds.
     """
 
     headers: Optional[Dict[str, bytes]] = None
 
     def __init__(
         self: "HttpClient",
+        /,
         max_redirects: Optional[int] = MAX_REDIRECTS,
         headers: Optional[Dict[str, bytes]] = None,
         compression: Optional[int] = DEFLATE | GZIP | BROTLI,
         validate_cert: bool = True,
+        connect_timeout: float = DEFAULT_CONNECT_TIMEOUT,
+        timeout: float = DEFAULT_TIMEOUT,
     ) -> None:
         self._client = AsyncClient(
             validate_cert,
+            int(connect_timeout * NS),
+            int(timeout * NS),
             max_redirects,
             merge_dict(self.headers, headers),
             compression,
diff --git a/tests/test_async_client.py b/tests/test_async_client.py
index 170a91c..007d1b6 100644
--- a/tests/test_async_client.py
+++ b/tests/test_async_client.py
@@ -18,7 +18,7 @@
 from gufo.http.httpd import Httpd
 from gufo.http import GZIP
 
-from .util import URL_PREFIX
+from .util import URL_PREFIX, UNROUTABLE_URL
 
 
 def test_get(httpd: Httpd) -> None:
@@ -298,3 +298,12 @@ async def inner() -> None:
             assert b"</html>" in data
 
     asyncio.run(inner())
+
+
+def test_connect_timeout() -> None:
+    async def inner() -> None:
+        async with HttpClient(connect_timeout=1.0) as client:
+            resp = await client.get(UNROUTABLE_URL)
+            assert resp.status == 200
+
+    asyncio.run(inner())
diff --git a/tests/util.py b/tests/util.py
index 4aea85b..ebf4323 100644
--- a/tests/util.py
+++ b/tests/util.py
@@ -13,3 +13,4 @@
 HTTPD_ADDRESS = "127.0.0.1"
 HTTPD_PORT = random.randint(52000, 53999)
 URL_PREFIX = f"http://{HTTPD_HOST}:{HTTPD_PORT}"
+UNROUTABLE_URL = "http://192.0.2.1/"