diff --git a/README.md b/README.md index 0780788..ce91a99 100644 --- a/README.md +++ b/README.md @@ -41,7 +41,8 @@ controller = UnifiController( model_db_path=None, # Optional: Path to a custom model database file auth_retry_enabled=True, # Optional: Enable automatic retries on authentication failure auth_retry_count=3, # Optional: Number of authentication retries - auth_retry_delay=5 # Optional: Delay in seconds between authentication retries + auth_retry_delay=5, # Optional: Delay in seconds between authentication retries + request_timeout=30 # Optional: Timeout in seconds for controller HTTP requests ) # 2. Fetch Data (Example: Devices for the 'default' site) diff --git a/docs/api/client.rst b/docs/api/client.rst index 07b2bd6..3a382fd 100644 --- a/docs/api/client.rst +++ b/docs/api/client.rst @@ -30,7 +30,8 @@ Connecting to a Controller model_db_path=None, # Optional: Path to a custom model database file auth_retry_enabled=True, auth_retry_count=3, - auth_retry_delay=1 + auth_retry_delay=1, + request_timeout=30 ) Getting Devices at a Site diff --git a/docs/index.rst b/docs/index.rst index 0e72013..2840729 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -62,7 +62,8 @@ Basic usage: model_db_path=None, # Optional: Path to a custom model database file auth_retry_enabled=True, # Optional: Enable automatic retries on authentication failure auth_retry_count=3, # Optional: Number of authentication retries - auth_retry_delay=5 # Optional: Delay in seconds between authentication retries + auth_retry_delay=5, # Optional: Delay in seconds between authentication retries + request_timeout=30 # Optional: Timeout in seconds for controller HTTP requests ) # Example: Fetch devices for the 'default' site diff --git a/tests/test_request_timeout.py b/tests/test_request_timeout.py new file mode 100644 index 0000000..e768e13 --- /dev/null +++ b/tests/test_request_timeout.py @@ -0,0 +1,102 @@ +from typing import Optional + +from unifi_controller_api import UnifiController + + +class FakeResponse: + status_code = 200 + + def __init__(self, payload=None): + self._payload = payload or {"meta": {"rc": "ok"}} + + def json(self): + return self._payload + + def raise_for_status(self): + return None + + +class RecordingSession: + def __init__(self): + self.post_calls = [] + self.get_calls = [] + self.request_calls = [] + self.cookies = {} + + def post(self, url, **kwargs): + self.post_calls.append((url, kwargs)) + return FakeResponse() + + def get(self, url, **kwargs): + self.get_calls.append((url, kwargs)) + return FakeResponse({"meta": {"rc": "ok"}, "data": []}) + + def request(self, method, url, **kwargs): + self.request_calls.append((method, url, kwargs)) + return FakeResponse() + + +def make_controller(request_timeout: Optional[float] = None): + controller = UnifiController.__new__(UnifiController) + controller.controller_url = "https://controller.example" + controller.original_controller_url = "https://controller.example" + controller.is_udm_pro = False + controller.session = RecordingSession() + controller.verify_ssl = True + controller.auth_retry_enabled = False + controller.auth_retry_count = 1 + controller.auth_retry_delay = 0.1 + controller.request_timeout = request_timeout + return controller + + +def test_constructor_accepts_request_timeout_and_uses_it_for_authentication(monkeypatch): + session = RecordingSession() + monkeypatch.setattr("unifi_controller_api.api_client.requests.Session", lambda: session) + + UnifiController( + "https://controller.example", + "admin", + "secret", + request_timeout=12.5, + ) + + assert session.post_calls[0][1]["timeout"] == 12.5 + + +def test_get_requests_use_configured_request_timeout(): + controller = make_controller(request_timeout=7.0) + + controller.invoke_get_rest_api_call("https://controller.example/api/self/sites") + + assert controller.session.get_calls[0][1]["timeout"] == 7.0 + + +def test_mutating_requests_use_configured_request_timeout_by_default(): + controller = make_controller(request_timeout=3.0) + + controller._invoke_api_call("POST", "https://controller.example/api/s/default/cmd", json_payload={}) + + assert controller.session.request_calls[0][2]["timeout"] == 3.0 + + +def test_per_call_timeout_overrides_configured_request_timeout(): + controller = make_controller(request_timeout=3.0) + + controller._invoke_api_call( + "POST", + "https://controller.example/api/s/default/cmd", + json_payload={}, + timeout=9.0, + ) + + assert controller.session.request_calls[0][2]["timeout"] == 9.0 + + +def test_request_timeout_defaults_to_none_for_backward_compatibility(monkeypatch): + session = RecordingSession() + monkeypatch.setattr("unifi_controller_api.api_client.requests.Session", lambda: session) + + UnifiController("https://controller.example", "admin", "secret") + + assert session.post_calls[0][1]["timeout"] is None diff --git a/unifi_controller_api/api_client.py b/unifi_controller_api/api_client.py index 669d146..81f06c2 100644 --- a/unifi_controller_api/api_client.py +++ b/unifi_controller_api/api_client.py @@ -76,6 +76,7 @@ def __init__( auth_retry_enabled=True, auth_retry_count=3, auth_retry_delay=1, + request_timeout=None, ): """ Initialize the Unifi Controller client and authenticate. @@ -101,6 +102,9 @@ def __init__( Defaults to 3. auth_retry_delay: Delay in seconds between retry attempts (0.1-30). Defaults to 1. + request_timeout: Optional timeout in seconds to apply to controller + HTTP requests. Defaults to None to preserve the + requests library behavior of waiting indefinitely. """ if auth_retry_count < 1 or auth_retry_count > 10: raise ValueError("auth_retry_count must be between 1 and 10") @@ -120,6 +124,7 @@ def __init__( self.auth_retry_enabled = auth_retry_enabled self.auth_retry_count = auth_retry_count self.auth_retry_delay = auth_retry_delay + self.request_timeout = request_timeout if model_db_path is None: self.model_db_path = os.path.join( @@ -171,6 +176,7 @@ def authenticate(self, username, password): login_uri, json={"username": username, "password": password}, verify=self.verify_ssl, + timeout=self.request_timeout, ) response.raise_for_status() @@ -223,7 +229,7 @@ def _invoke_api_call( url: str, json_payload: Optional[Dict[str, Any]] = None, headers: Optional[Dict[str, str]] = None, - timeout: Optional[int] = None + timeout: Optional[float] = None ) -> requests.Response: """ Make an API request with the specified method, handling potential re-authentication. @@ -251,7 +257,7 @@ def _invoke_api_call( request_kwargs = { 'verify': self.verify_ssl, - 'timeout': timeout + 'timeout': timeout if timeout is not None else getattr(self, 'request_timeout', None) } if headers: @@ -389,9 +395,12 @@ def invoke_get_rest_api_call(self, url, headers=None): try: if headers: response = self.session.get( - url, headers=headers, verify=self.verify_ssl) + url, headers=headers, verify=self.verify_ssl, + timeout=getattr(self, 'request_timeout', None)) else: - response = self.session.get(url, verify=self.verify_ssl) + response = self.session.get( + url, verify=self.verify_ssl, + timeout=getattr(self, 'request_timeout', None)) if response.status_code == 401 and self.auth_retry_enabled: if hasattr(self, '_username') and hasattr(self, '_password'): @@ -411,10 +420,12 @@ def invoke_get_rest_api_call(self, url, headers=None): if headers: response = self.session.get( - url, headers=headers, verify=self.verify_ssl) + url, headers=headers, verify=self.verify_ssl, + timeout=getattr(self, 'request_timeout', None)) else: response = self.session.get( - url, verify=self.verify_ssl) + url, verify=self.verify_ssl, + timeout=getattr(self, 'request_timeout', None)) if response.status_code != 401: break