Clean up failed login sessions

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"""

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()