Add proxies and timeout parameters

pybatfish/client/restv2helper.py

@@ -441,6 +441,8 @@ def _delete(session: Session, url_tail: str, params: dict[str, Any] | None = Non
         headers=_get_headers(session),
         params=params,
         verify=session.verify_ssl_certs,
+        timeout=session.timeout,
+        proxies=session.proxies,
     )
     _check_response_status(response)
 
@@ -461,6 +463,8 @@ def _get(
         params=params,
         stream=stream,
         verify=session.verify_ssl_certs,
+        timeout=session.timeout,
+        proxies=session.proxies,
     )
     _check_response_status(response)
     return response
@@ -521,6 +525,8 @@ def _post(
         headers=headers,
         params=params,
         verify=session.verify_ssl_certs,
+        timeout=session.timeout,
+        proxies=session.proxies,
     )
     _check_response_status(response)
     return None
@@ -544,6 +550,8 @@ def _put(session, url_tail, params=None, json=None, stream=None):
         headers=headers,
         verify=session.verify_ssl_certs,
         params=params,
+        timeout=session.timeout,
+        proxies=session.proxies,
     )
     _check_response_status(response)
     return None

pybatfish/client/session.py

@@ -319,6 +319,8 @@ class Session:
     :ivar port_v2: Legacy alias for port
     :ivar ssl: Whether to use SSL when connecting to Batfish (False by default)
     :ivar api_key: Your API key
+    :ivar proxies: A dictionary of proxies to use for the session. See `requests` documentation for details.
+    :ivar timeout: Timeout for requests in seconds (default is 30 seconds)
     """
 
     def __init__(
@@ -333,6 +335,8 @@ def __init__(
         verify_ssl_certs: bool = Options.verify_ssl_certs,
         api_key: str = CoordConsts.DEFAULT_API_KEY,
         load_questions: bool = True,
+        proxies: dict[str, str] | None = None,
+        timeout: int = 30,
     ):
         # Coordinator args
         self.host: str = host
@@ -360,6 +364,8 @@ def __init__(
 
         self.elapsed_delay: int = 5
         self.stale_timeout: int = 5
+        self.proxies = proxies
+        self.timeout = timeout
 
         # Auto-load question templates
         if load_questions:

tests/client/test_restv2helper.py

@@ -56,6 +56,8 @@ def session() -> Session:
     s.get_base_url2.return_value = BASE_URL
     s.api_key = "0000"
     s.verify_ssl_certs = True
+    s.timeout = 30
+    s.proxies = None
     return s
 
 
@@ -92,6 +94,8 @@ def test_delete(session: Session, request_session: Mock) -> None:
         headers=_get_headers(session),
         params=None,
         verify=session.verify_ssl_certs,
+        proxies=session.proxies,
+        timeout=session.timeout,
     )
 
 
@@ -111,6 +115,8 @@ def test_get(session: Session, request_session: Mock) -> None:
         params=None,
         stream=False,
         verify=session.verify_ssl_certs,
+        proxies=session.proxies,
+        timeout=session.timeout,
     )
 
     # Fast-failing session
@@ -124,6 +130,8 @@ def test_get(session: Session, request_session: Mock) -> None:
         params=None,
         stream=False,
         verify=session.verify_ssl_certs,
+        proxies=session.proxies,
+        timeout=session.timeout,
     )
 
 
@@ -144,6 +152,8 @@ def test_post_json(session, request_session):
         headers=_get_headers(session),
         params=None,
         verify=session.verify_ssl_certs,
+        proxies=session.proxies,
+        timeout=session.timeout,
     )
 
 
@@ -165,6 +175,8 @@ def test_post_stream(session, request_session):
             headers=expected_headers,
             params=None,
             verify=session.verify_ssl_certs,
+            proxies=session.proxies,
+            timeout=session.timeout,
         )
 
 
@@ -184,6 +196,8 @@ def test_put(session, request_session):
         headers=_get_headers(session),
         verify=session.verify_ssl_certs,
         params=None,
+        proxies=session.proxies,
+        timeout=session.timeout,
     )