diff --git a/src/ahriman/core/status/web_client.py b/src/ahriman/core/status/web_client.py index 3dc2f00f..cabdc0c1 100644 --- a/src/ahriman/core/status/web_client.py +++ b/src/ahriman/core/status/web_client.py @@ -102,15 +102,22 @@ class WebClient(Client, LazyLogging): return address, False @contextlib.contextmanager - def __get_session(self) -> Generator[requests.Session, None, None]: + def __get_session(self, session: requests.Session | None = None) -> Generator[requests.Session, None, None]: """ execute request and handle exceptions + Args: + session(requests.Session | None, optional): session to be used or stored instance property otherwise + (Default value = None) + Yields: requests.Session: session for requests """ try: - yield self.__session + if session is not None: + yield session # use session from arguments + else: + yield self.__session # use instance generated session except requests.RequestException as e: if self.suppress_errors: return @@ -136,13 +143,16 @@ class WebClient(Client, LazyLogging): return session session = requests.Session() - self._login() + self._login(session) return session - def _login(self) -> None: + def _login(self, session: requests.Session) -> None: """ process login to the service + + Args: + session(requests.Session): request session to login """ if self.user is None: return # no auth configured @@ -152,7 +162,7 @@ class WebClient(Client, LazyLogging): "password": self.user.password } - with self.__get_session() as session: + with self.__get_session(session): response = session.post(self._login_url, json=payload) response.raise_for_status() diff --git a/tests/ahriman/core/status/test_web_client.py b/tests/ahriman/core/status/test_web_client.py index a259172d..b4756be9 100644 --- a/tests/ahriman/core/status/test_web_client.py +++ b/tests/ahriman/core/status/test_web_client.py @@ -55,7 +55,7 @@ def test_create_session(web_client: WebClient, mocker: MockerFixture) -> None: session = web_client._create_session(use_unix_socket=False) assert isinstance(session, requests.Session) assert not isinstance(session, requests_unixsocket.Session) - login_mock.assert_called_once_with() + login_mock.assert_called_once_with(pytest.helpers.anyvar(int)) def test_create_session_unix_socket(web_client: WebClient, mocker: MockerFixture) -> None: @@ -80,7 +80,7 @@ def test_login(web_client: WebClient, user: User, mocker: MockerFixture) -> None "password": user.password } - web_client._login() + web_client._login(requests.Session()) requests_mock.assert_called_once_with(pytest.helpers.anyvar(str, True), json=payload) @@ -90,7 +90,7 @@ def test_login_failed(web_client: WebClient, user: User, mocker: MockerFixture) """ web_client.user = user mocker.patch("requests.Session.post", side_effect=Exception()) - web_client._login() + web_client._login(requests.Session()) def test_login_failed_http_error(web_client: WebClient, user: User, mocker: MockerFixture) -> None: @@ -99,7 +99,7 @@ def test_login_failed_http_error(web_client: WebClient, user: User, mocker: Mock """ web_client.user = user mocker.patch("requests.Session.post", side_effect=requests.exceptions.HTTPError()) - web_client._login() + web_client._login(requests.Session()) def test_login_skip(web_client: WebClient, mocker: MockerFixture) -> None: @@ -107,7 +107,7 @@ def test_login_skip(web_client: WebClient, mocker: MockerFixture) -> None: must skip login if no user set """ requests_mock = mocker.patch("requests.Session.post") - web_client._login() + web_client._login(requests.Session()) requests_mock.assert_not_called()