From c95d26e3bd9e250e776d84b7bcc0b235ea5c138a Mon Sep 17 00:00:00 2001 From: TianKai Ma Date: Mon, 10 Aug 2026 00:05:37 +0800 Subject: [PATCH] fix: clean up failed login sessions --- src/utils/auth.py | 54 +++++++++++++++++------------ tests/test_auth.py | 85 +++++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 117 insertions(+), 22 deletions(-) diff --git a/src/utils/auth.py b/src/utils/auth.py index b77163ab66..4aa2ced689 100644 --- a/src/utils/auth.py +++ b/src/utils/auth.py @@ -398,37 +398,49 @@ def __init__(self, headless: bool = True, *, after_login_services: bool = True): async def __aenter__(self) -> RequestSession: """Initialize browser, perform login and return RequestSession""" - self.playwright = await async_playwright().__aenter__() - self.browser = await self.playwright.chromium.launch(headless=self.headless) - self.context = await self.browser.new_context(locale="zh-CN") - self.page = await self.context.new_page() - - await self._login() - if self.after_login_services: - await self._after_login() - - client = await self._create_http_client() - timeout_ms = self.timeout_ms if self.timeout_ms > 0 else 10 * 60 * 1000 - self.request_session = RequestSession( - client=client, - page=self.page, - timeout_ms=timeout_ms, - ) - return self.request_session + try: + self.playwright = await async_playwright().__aenter__() + self.browser = await self.playwright.chromium.launch(headless=self.headless) + self.context = await self.browser.new_context(locale="zh-CN") + self.page = await self.context.new_page() + + await self._login() + if self.after_login_services: + await self._after_login() + + client = await self._create_http_client() + timeout_ms = self.timeout_ms if self.timeout_ms > 0 else 10 * 60 * 1000 + self.request_session = RequestSession( + client=client, + page=self.page, + timeout_ms=timeout_ms, + ) + return self.request_session + except BaseException: + await self._cleanup() + raise async def __aexit__(self, exc_type, exc, tb): """Cleanup browser resources""" + await self._cleanup() + + async def _cleanup(self) -> None: + """Close every resource that was successfully initialized.""" with suppress(Exception): if self.request_session: await self.request_session.close() with suppress(Exception): - await self.page.close() + if page := getattr(self, "page", None): + await page.close() with suppress(Exception): - await self.context.close() + if context := getattr(self, "context", None): + await context.close() with suppress(Exception): - await self.browser.close() + if browser := getattr(self, "browser", None): + await browser.close() with suppress(Exception): - await self.playwright.stop() + if playwright := getattr(self, "playwright", None): + await playwright.stop() async def _login(self): """Perform USTC login sequence""" diff --git a/tests/test_auth.py b/tests/test_auth.py index 2970d90f23..fa1b696a51 100644 --- a/tests/test_auth.py +++ b/tests/test_auth.py @@ -1,7 +1,7 @@ import json import os import unittest -from unittest.mock import patch +from unittest.mock import AsyncMock, patch import httpx @@ -142,5 +142,88 @@ def test_after_login_services_can_be_disabled(self) -> None: self.assertFalse(session.after_login_services) +class USTCSessionCleanupTest(unittest.IsolatedAsyncioTestCase): + def create_session(self) -> USTCSession: + with patch.dict( + os.environ, + { + "USTC_PASSPORT_USERNAME": "user", + "USTC_PASSPORT_PASSWORD": "password", + "USTC_PASSPORT_TOTP_URL": "", + }, + ): + return USTCSession(after_login_services=False) + + async def test_launch_failure_stops_playwright(self) -> None: + playwright = _FakePlaywright(launch_error=RuntimeError("launch failed")) + session = self.create_session() + + with ( + patch( + "src.utils.auth.async_playwright", return_value=_FakeManager(playwright) + ), + self.assertRaisesRegex(RuntimeError, "launch failed"), + ): + await session.__aenter__() + + playwright.stop.assert_awaited_once() + + async def test_login_failure_closes_all_browser_resources(self) -> None: + playwright = _FakePlaywright() + session = self.create_session() + + with ( + patch( + "src.utils.auth.async_playwright", return_value=_FakeManager(playwright) + ), + patch.object( + session, "_login", AsyncMock(side_effect=RuntimeError("login failed")) + ), + self.assertRaisesRegex(RuntimeError, "login failed"), + ): + await session.__aenter__() + + playwright.page.close.assert_awaited_once() + playwright.context.close.assert_awaited_once() + playwright.browser.close.assert_awaited_once() + playwright.stop.assert_awaited_once() + + +class _FakeManager: + def __init__(self, playwright) -> None: + self.playwright = playwright + + async def __aenter__(self): + return self.playwright + + +class _FakeChromium: + def __init__(self, browser, launch_error=None) -> None: + self.browser = browser + self.launch_error = launch_error + + async def launch(self, *, headless): + if self.launch_error: + raise self.launch_error + return self.browser + + +class _FakePlaywright: + def __init__(self, launch_error=None) -> None: + self.page = type("Page", (), {"close": AsyncMock()})() + self.context = type( + "Context", + (), + {"new_page": AsyncMock(return_value=self.page), "close": AsyncMock()}, + )() + self.browser = type( + "Browser", + (), + {"new_context": AsyncMock(return_value=self.context), "close": AsyncMock()}, + )() + self.chromium = _FakeChromium(self.browser, launch_error) + self.stop = AsyncMock() + + if __name__ == "__main__": unittest.main()