diff --git a/README.md b/README.md index cac6bee..8cea223 100644 --- a/README.md +++ b/README.md @@ -36,7 +36,7 @@ For `session` a valid `aiohttp.ClientSession` needs to be provided. If ssl verif `device_token` should be added when using a two-step authentication account, otherwise DSM will ask to login with a One Time Password (OTP) and requests will fail (see the login section for more details). -Default `timeout` is 10 seconds. +Default `timeout` is 10 seconds. `SynologyDSM` also takes a `aiohttp.ClientTimeout` as `timeout`. ## Login diff --git a/src/synology_dsm/synology_dsm.py b/src/synology_dsm/synology_dsm.py index f776b51..8b48c57 100644 --- a/src/synology_dsm/synology_dsm.py +++ b/src/synology_dsm/synology_dsm.py @@ -10,7 +10,7 @@ from typing import Any, Coroutine, TypedDict from urllib.parse import quote, urlencode -import aiohttp +from aiohttp import ClientError, ClientSession, ClientTimeout from yarl import URL from .api import SynoBaseApi @@ -61,20 +61,23 @@ class SynologyDSM: def __init__( self, - session: aiohttp.ClientSession, + session: ClientSession, dsm_ip: str, dsm_port: int, username: str, password: str, use_https: bool = False, - timeout: int = 10, + timeout: int | ClientTimeout = 10, device_token: str | None = None, debugmode: bool = False, ): """Constructor method.""" self.username = username self._password = password - self._aiohttp_timeout = aiohttp.ClientTimeout(total=timeout) + if isinstance(timeout, ClientTimeout): + self._aiohttp_timeout = timeout + else: + self._aiohttp_timeout = ClientTimeout(total=timeout) self._debugmode = debugmode # Session @@ -387,9 +390,9 @@ async def _execute_request( return await response.text() # We got a 400, 401 or 404 ... - raise aiohttp.ClientError(response) + raise ClientError(response) - except (aiohttp.ClientError, asyncio.TimeoutError, JSONDecodeError) as exp: + except (ClientError, asyncio.TimeoutError, JSONDecodeError) as exp: raise SynologyDSMRequestException(exp) from exp async def update( diff --git a/tests/test_synology_dsm.py b/tests/test_synology_dsm.py index c37de80..8eadc6b 100644 --- a/tests/test_synology_dsm.py +++ b/tests/test_synology_dsm.py @@ -2,6 +2,7 @@ # pylint: disable=protected-access import pytest +from aiohttp import ClientTimeout from synology_dsm.api.core.external_usb import SynoCoreExternalUSB from synology_dsm.api.core.security import SynoCoreSecurity @@ -278,7 +279,16 @@ async def test_login_failed(self, version): assert not dsm._session_id @pytest.mark.parametrize("version", [5, 6, 7]) - def test_request_timeout(self, version): + @pytest.mark.parametrize( + "timeout,expected_result", + [ + (2, (2, None)), + (15, (15, None)), + (ClientTimeout(total=5), (5, None)), + (ClientTimeout(total=60, connect=15), (60, 15)), + ], + ) + def test_request_timeout(self, version, timeout, expected_result): """Test request timeout.""" dsm = SynologyDSMMock( None, @@ -287,10 +297,11 @@ def test_request_timeout(self, version): VALID_USER, VALID_PASSWORD, VALID_HTTPS, - timeout=2, + timeout=timeout, ) dsm.dsm_version = version - assert dsm._aiohttp_timeout.total == 2 + assert dsm._aiohttp_timeout.total == expected_result[0] + assert dsm._aiohttp_timeout.connect == expected_result[1] @pytest.mark.asyncio async def test_request_get(self, dsm):