diff --git a/pybatfish/client/restv2helper.py b/pybatfish/client/restv2helper.py index 01ab2b32..85d65fd0 100644 --- a/pybatfish/client/restv2helper.py +++ b/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 diff --git a/pybatfish/client/session.py b/pybatfish/client/session.py index b9167b36..d8bec964 100644 --- a/pybatfish/client/session.py +++ b/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: diff --git a/tests/client/test_restv2helper.py b/tests/client/test_restv2helper.py index f334be31..7a1f91b7 100644 --- a/tests/client/test_restv2helper.py +++ b/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, )