diff --git a/src/runpod_flash/cli/commands/login.py b/src/runpod_flash/cli/commands/login.py index 2483566b..df80e68f 100644 --- a/src/runpod_flash/cli/commands/login.py +++ b/src/runpod_flash/cli/commands/login.py @@ -1,6 +1,4 @@ import asyncio -import datetime as dt -from typing import Optional import typer from rich.console import Console @@ -15,20 +13,8 @@ console = Console() -POLL_INTERVAL_SECONDS = 2.0 -DEFAULT_TIMEOUT_SECONDS = 600.0 - -def _parse_expires_at(value: Optional[str]) -> Optional[dt.datetime]: - if not value: - return None - try: - return dt.datetime.fromisoformat(value.replace("Z", "+00:00")) - except ValueError: - return None - - -async def _login(open_browser: bool, timeout_seconds: float) -> None: +async def _login(open_browser: bool) -> None: async with RunpodGraphQLClient(require_api_key=False) as client: request = await client.create_flash_auth_request() request_id = request.get("id") @@ -45,46 +31,40 @@ async def _login(open_browser: bool, timeout_seconds: float) -> None: if open_browser: typer.launch(auth_url) - expires_at = _parse_expires_at(request.get("expiresAt")) - deadline = dt.datetime.now(dt.timezone.utc) + dt.timedelta( - seconds=timeout_seconds - ) - if expires_at and expires_at < deadline: - deadline = expires_at + api_key = console.input( + "Paste the API key shown after authorization: ", password=True + ).strip() - with console.status("[dim]Waiting for authorization...[/dim]"): - while True: - status_payload = await client.get_flash_auth_request_status(request_id) - status = status_payload.get("status") - api_key = status_payload.get("apiKey") - - if api_key and status in {"APPROVED", "CONSUMED"}: - check_and_migrate_legacy_credentials() - path = save_api_key(api_key) - console.print( - f"[green]Logged in.[/green] Credentials saved to [dim]{path}[/dim]" - ) - console.print() - return - - if status in {"DENIED", "EXPIRED", "CONSUMED"}: - raise RuntimeError(f"login failed: {status.lower()}") + if not api_key: + raise RuntimeError( + "no api key provided. copy the key shown in the browser after approval " + "and paste it here." + ) + if not api_key.startswith("rpa_"): + raise RuntimeError("invalid api key format. Runpod API keys start with rpa_.") - if dt.datetime.now(dt.timezone.utc) >= deadline: - raise RuntimeError("login timed out") + async with RunpodGraphQLClient(api_key=api_key) as client: + if not await client.verify_api_key(): + raise RuntimeError( + "api key could not be verified. Check the pasted key and try again." + ) - await asyncio.sleep(POLL_INTERVAL_SECONDS) + check_and_migrate_legacy_credentials() + path = save_api_key(api_key) + console.print(f"[green]Logged in.[/green] Credentials saved to [dim]{path}[/dim]") + console.print() def login_command( no_open: bool = typer.Option(False, "--no-open", help="do not open the browser"), - timeout: float = typer.Option( - DEFAULT_TIMEOUT_SECONDS, "--timeout", help="max wait time in seconds" - ), ): """Authenticate and save a Runpod API key for flash.""" try: - asyncio.run(_login(open_browser=not no_open, timeout_seconds=timeout)) + asyncio.run(_login(open_browser=not no_open)) except RuntimeError as exc: print_error(console, str(exc)) raise typer.Exit(code=1) + except KeyboardInterrupt: + console.print() + print_error(console, "login cancelled") + raise typer.Exit(code=130) diff --git a/src/runpod_flash/core/api/runpod.py b/src/runpod_flash/core/api/runpod.py index 42c23d87..b5c27c80 100644 --- a/src/runpod_flash/core/api/runpod.py +++ b/src/runpod_flash/core/api/runpod.py @@ -860,19 +860,26 @@ async def create_flash_auth_request(self) -> Dict[str, Any]: result = await self._execute_graphql(mutation) return result.get("createFlashAuthRequest", {}) - async def get_flash_auth_request_status(self, request_id: str) -> Dict[str, Any]: + async def verify_api_key(self) -> bool: query = """ - query flashAuthRequestStatus($flashAuthRequestId: String!) { - flashAuthRequestStatus(flashAuthRequestId: $flashAuthRequestId) { - id - status - expiresAt - apiKey + query verifyApiKey { + myself { + __typename } } """ - result = await self._execute_graphql(query, {"flashAuthRequestId": request_id}) - return result.get("flashAuthRequestStatus", {}) + try: + result = await self._execute_graphql(query) + except ( + _GraphQLHTTPStatusError, + _GraphQLNetworkError, + _GraphQLErrorResponse, + asyncio.TimeoutError, + ) as exc: + log.debug("Failed to verify Runpod API key: %s", exc) + return False + + return isinstance(result.get("myself"), dict) async def close(self): """Close the HTTP session.""" diff --git a/tests/unit/test_login.py b/tests/unit/test_login.py index 9d061a27..8a7c2cde 100644 --- a/tests/unit/test_login.py +++ b/tests/unit/test_login.py @@ -5,24 +5,6 @@ import pytest -from runpod_flash.cli.commands.login import _parse_expires_at - - -class TestParseExpiresAt: - def test_iso_format(self): - result = _parse_expires_at("2026-03-01T12:00:00Z") - assert result is not None - assert result.year == 2026 - - def test_none_input(self): - assert _parse_expires_at(None) is None - - def test_empty_string(self): - assert _parse_expires_at("") is None - - def test_invalid_string(self): - assert _parse_expires_at("not-a-date") is None - class TestGraphQLClientNoKeyForLogin: """Login mutations must not send stored credentials.""" @@ -61,16 +43,14 @@ def test_require_api_key_true_loads_key(self): assert client.api_key == "loaded-key" -def _make_mock_client(**status_return): +def _make_mock_client(): """Build an AsyncMock that works as an async context manager.""" client = AsyncMock() client.create_flash_auth_request.return_value = { "id": "req-123", "expiresAt": None, } - client.get_flash_auth_request_status.return_value = status_return - # _login uses `async with RunpodGraphQLClient(...) as client:`, - # so __aenter__ must return the same mock instance + client.verify_api_key.return_value = True client.__aenter__.return_value = client return client @@ -85,42 +65,73 @@ def _get_login_fn(): class TestLoginFlow: - async def test_login_denied(self): - mock_client = _make_mock_client(status="DENIED", apiKey=None) + async def test_login_saves_pasted_key(self, isolate_credentials_file): + mock_client = _make_mock_client() _login = _get_login_fn() - with patch( - "runpod_flash.cli.commands.login.RunpodGraphQLClient", - return_value=mock_client, + with ( + patch( + "runpod_flash.cli.commands.login.RunpodGraphQLClient", + return_value=mock_client, + ), + patch("runpod_flash.cli.commands.login.console") as mock_console, ): - with pytest.raises(RuntimeError, match="login failed: denied"): - await _login(open_browser=False, timeout_seconds=5) + mock_console.input.return_value = "rpa_pasted-api-key" + await _login(open_browser=False) + assert isolate_credentials_file.exists() + assert "rpa_pasted-api-key" in isolate_credentials_file.read_text() + mock_console.input.assert_called_once_with( + "Paste the API key shown after authorization: ", password=True + ) - async def test_login_approved_saves_key(self, isolate_credentials_file): - mock_client = _make_mock_client(status="APPROVED", apiKey="fresh-api-key") + async def test_login_empty_key_raises(self): + mock_client = _make_mock_client() _login = _get_login_fn() - with patch( - "runpod_flash.cli.commands.login.RunpodGraphQLClient", - return_value=mock_client, + with ( + patch( + "runpod_flash.cli.commands.login.RunpodGraphQLClient", + return_value=mock_client, + ), + patch("runpod_flash.cli.commands.login.console") as mock_console, ): - await _login(open_browser=False, timeout_seconds=5) - assert isolate_credentials_file.exists() - assert "fresh-api-key" in isolate_credentials_file.read_text() + mock_console.input.return_value = " " + with pytest.raises(RuntimeError, match="copy the key shown"): + await _login(open_browser=False) - async def test_login_expired(self): - mock_client = _make_mock_client(status="EXPIRED", apiKey=None) + async def test_login_invalid_key_format_raises(self): + mock_client = _make_mock_client() _login = _get_login_fn() - with patch( - "runpod_flash.cli.commands.login.RunpodGraphQLClient", - return_value=mock_client, + with ( + patch( + "runpod_flash.cli.commands.login.RunpodGraphQLClient", + return_value=mock_client, + ), + patch("runpod_flash.cli.commands.login.console") as mock_console, + ): + mock_console.input.return_value = "not-a-runpod-key" + with pytest.raises(RuntimeError, match="Runpod API keys start with rpa_"): + await _login(open_browser=False) + + async def test_login_unverified_key_raises(self): + mock_client = _make_mock_client() + mock_client.verify_api_key.return_value = False + _login = _get_login_fn() + + with ( + patch( + "runpod_flash.cli.commands.login.RunpodGraphQLClient", + return_value=mock_client, + ), + patch("runpod_flash.cli.commands.login.console") as mock_console, ): - with pytest.raises(RuntimeError, match="login failed: expired"): - await _login(open_browser=False, timeout_seconds=5) + mock_console.input.return_value = "rpa_bad-key" + with pytest.raises(RuntimeError, match="api key could not be verified"): + await _login(open_browser=False) async def test_no_request_id_raises(self): - mock_client = _make_mock_client(status="APPROVED", apiKey="key") + mock_client = _make_mock_client() mock_client.create_flash_auth_request.return_value = {} _login = _get_login_fn() @@ -129,4 +140,4 @@ async def test_no_request_id_raises(self): return_value=mock_client, ): with pytest.raises(RuntimeError, match="auth request failed"): - await _login(open_browser=False, timeout_seconds=5) + await _login(open_browser=False) diff --git a/tests/unit/test_login_extended.py b/tests/unit/test_login_extended.py index ff0edf40..a371012b 100644 --- a/tests/unit/test_login_extended.py +++ b/tests/unit/test_login_extended.py @@ -1,16 +1,12 @@ """Extended tests for flash login command and auth GraphQL methods. -Covers gaps in test_login.py: +Covers: - open_browser=True path (typer.launch) -- CONSUMED status with and without apiKey -- expiresAt deadline capping -- Timeout branch - login_command CLI wrapper -- create_flash_auth_request and get_flash_auth_request_status direct tests +- create_flash_auth_request direct test - GraphQL session without API key (Authorization header omission) """ -import datetime as dt import os from unittest.mock import AsyncMock, patch @@ -25,19 +21,19 @@ def _fresh_login_module(): return importlib.import_module("runpod_flash.cli.commands.login") -def _make_mock_client(**status_return): +def _make_mock_client(): """Build an AsyncMock that works as an async context manager.""" client = AsyncMock() client.create_flash_auth_request.return_value = { "id": "req-123", "expiresAt": None, } - client.get_flash_auth_request_status.return_value = status_return + client.verify_api_key.return_value = True client.__aenter__.return_value = client return client -# ── _login flow gaps ───────────────────────────────────────────────────── +# -- _login flow gaps -- @pytest.mark.serial @@ -47,7 +43,7 @@ class TestLoginOpenBrowser: @pytest.mark.asyncio async def test_open_browser_calls_typer_launch(self, isolate_credentials_file): """When open_browser=True, typer.launch is called with the auth URL.""" - mock_client = _make_mock_client(status="APPROVED", apiKey="key-123") + mock_client = _make_mock_client() with ( patch( @@ -55,143 +51,17 @@ async def test_open_browser_calls_typer_launch(self, isolate_credentials_file): return_value=mock_client, ), patch("runpod_flash.cli.commands.login.typer.launch") as mock_launch, - patch( - "runpod_flash.cli.commands.login.asyncio.sleep", new_callable=AsyncMock - ), - patch("runpod_flash.cli.commands.login.console"), + patch("runpod_flash.cli.commands.login.console") as mock_console, ): - await _fresh_login_module()._login(open_browser=True, timeout_seconds=5) + mock_console.input.return_value = "rpa_key-123" + await _fresh_login_module()._login(open_browser=True) mock_launch.assert_called_once() url = mock_launch.call_args[0][0] assert "req-123" in url -@pytest.mark.serial -class TestLoginConsumedStatus: - """Test CONSUMED status handling.""" - - @pytest.mark.asyncio - async def test_consumed_with_api_key_saves_key(self, isolate_credentials_file): - """CONSUMED with a valid apiKey saves credentials and succeeds.""" - mock_client = _make_mock_client(status="CONSUMED", apiKey="consumed-key") - - with ( - patch( - "runpod_flash.cli.commands.login.RunpodGraphQLClient", - return_value=mock_client, - ), - patch( - "runpod_flash.cli.commands.login.asyncio.sleep", new_callable=AsyncMock - ), - patch("runpod_flash.cli.commands.login.console"), - ): - await _fresh_login_module()._login(open_browser=False, timeout_seconds=5) - - assert isolate_credentials_file.exists() - assert "consumed-key" in isolate_credentials_file.read_text() - - @pytest.mark.asyncio - async def test_consumed_without_api_key_raises(self): - """CONSUMED without an apiKey raises RuntimeError.""" - mock_client = _make_mock_client(status="CONSUMED", apiKey=None) - - with ( - patch( - "runpod_flash.cli.commands.login.RunpodGraphQLClient", - return_value=mock_client, - ), - patch( - "runpod_flash.cli.commands.login.asyncio.sleep", new_callable=AsyncMock - ), - patch("runpod_flash.cli.commands.login.console"), - ): - with pytest.raises(RuntimeError, match="login failed: consumed"): - await _fresh_login_module()._login( - open_browser=False, timeout_seconds=5 - ) - - -@pytest.mark.serial -class TestLoginExpiresAtDeadline: - """Test expiresAt deadline capping.""" - - @pytest.mark.asyncio - async def test_expires_at_caps_deadline(self, tmp_path): - """When expiresAt is earlier than timeout, deadline uses expiresAt.""" - # Set expiresAt to 1 second in the past so we immediately timeout - past = (dt.datetime.now(dt.timezone.utc) - dt.timedelta(seconds=1)).isoformat() - mock_client = AsyncMock() - mock_client.create_flash_auth_request.return_value = { - "id": "req-456", - "expiresAt": past, - } - # Status returns PENDING so we enter the polling loop - mock_client.get_flash_auth_request_status.return_value = { - "status": "PENDING", - "apiKey": None, - } - mock_client.__aenter__.return_value = mock_client - - with ( - patch( - "runpod_flash.cli.commands.login.RunpodGraphQLClient", - return_value=mock_client, - ), - patch( - "runpod_flash.cli.commands.login.asyncio.sleep", - new_callable=AsyncMock, - ), - patch( - "runpod_flash.cli.commands.login.console", - ), - ): - with pytest.raises(RuntimeError, match="login timed out"): - await _fresh_login_module()._login( - open_browser=False, timeout_seconds=600 - ) - - -@pytest.mark.serial -class TestLoginTimeout: - """Test timeout branch.""" - - @pytest.mark.asyncio - async def test_timeout_raises_runtime_error(self): - """When deadline is reached, RuntimeError is raised.""" - # Use expiresAt set in the past to force immediate timeout - past = (dt.datetime.now(dt.timezone.utc) - dt.timedelta(seconds=1)).isoformat() - mock_client = AsyncMock() - mock_client.create_flash_auth_request.return_value = { - "id": "req-789", - "expiresAt": past, - } - mock_client.get_flash_auth_request_status.return_value = { - "status": "PENDING", - "apiKey": None, - } - mock_client.__aenter__.return_value = mock_client - - with ( - patch( - "runpod_flash.cli.commands.login.RunpodGraphQLClient", - return_value=mock_client, - ), - patch( - "runpod_flash.cli.commands.login.asyncio.sleep", - new_callable=AsyncMock, - ), - patch( - "runpod_flash.cli.commands.login.console", - ), - ): - with pytest.raises(RuntimeError, match="login timed out"): - await _fresh_login_module()._login( - open_browser=False, timeout_seconds=600 - ) - - -# ── login_command CLI wrapper ──────────────────────────────────────────── +# -- login_command CLI wrapper -- @pytest.mark.serial @@ -200,28 +70,42 @@ class TestLoginCommand: def test_login_command_raises_exit_on_error(self): """login_command raises typer.Exit(1) when _login raises RuntimeError.""" - with patch( - "runpod_flash.cli.commands.login.asyncio.run", - side_effect=RuntimeError("auth failed"), - ): + login = AsyncMock(side_effect=RuntimeError("auth failed")) + with patch("runpod_flash.cli.commands.login._login", new=login): with pytest.raises(typer.Exit) as exc_info: - _fresh_login_module().login_command(no_open=True, timeout=5.0) + _fresh_login_module().login_command(no_open=True) assert exc_info.value.exit_code == 1 + def test_login_command_handles_keyboard_interrupt(self): + login = AsyncMock(side_effect=KeyboardInterrupt) + with ( + patch("runpod_flash.cli.commands.login._login", new=login), + patch("runpod_flash.cli.commands.login.print_error") as mock_print_error, + patch("runpod_flash.cli.commands.login.console") as mock_console, + ): + with pytest.raises(typer.Exit) as exc_info: + _fresh_login_module().login_command(no_open=True) + + assert exc_info.value.exit_code == 130 + mock_console.print.assert_called_once_with() + mock_print_error.assert_called_once_with(mock_console, "login cancelled") + def test_login_command_succeeds(self): """login_command succeeds when _login completes normally.""" - with patch("runpod_flash.cli.commands.login.asyncio.run"): - # Should not raise - _fresh_login_module().login_command(no_open=True, timeout=5.0) + login = AsyncMock() + with patch("runpod_flash.cli.commands.login._login", new=login): + _fresh_login_module().login_command(no_open=True) + + login.assert_awaited_once_with(open_browser=False) -# ── GraphQL auth methods ──────────────────────────────────────────────── +# -- GraphQL auth methods -- @pytest.mark.serial class TestGraphQLAuthMethods: - """Direct tests for create_flash_auth_request and get_flash_auth_request_status.""" + """Direct tests for GraphQL auth helpers.""" @pytest.mark.asyncio async def test_create_flash_auth_request(self): @@ -246,7 +130,6 @@ async def test_create_flash_auth_request(self): assert result["status"] == "PENDING" client._execute_graphql.assert_called_once() - # Verify the mutation string mutation = client._execute_graphql.call_args[0][0] assert "createFlashAuthRequest" in mutation @@ -263,47 +146,31 @@ async def test_create_flash_auth_request_empty_response(self): assert result == {} @pytest.mark.asyncio - async def test_get_flash_auth_request_status(self): - """get_flash_auth_request_status sends query with request_id.""" + async def test_verify_api_key_returns_true_for_authenticated_user(self): + """verify_api_key returns true when the API returns a current user.""" with patch.dict(os.environ, {"RUNPOD_API_KEY": "test-key"}): from runpod_flash.core.api.runpod import RunpodGraphQLClient client = RunpodGraphQLClient() client._execute_graphql = AsyncMock( - return_value={ - "flashAuthRequestStatus": { - "id": "req-xyz", - "status": "APPROVED", - "apiKey": "new-key", - "expiresAt": None, - } - } + return_value={"myself": {"__typename": "User"}} ) - result = await client.get_flash_auth_request_status("req-xyz") - - assert result["status"] == "APPROVED" - assert result["apiKey"] == "new-key" - - # Verify variables passed - call_args = client._execute_graphql.call_args - variables = call_args[0][1] - assert variables["flashAuthRequestId"] == "req-xyz" + assert await client.verify_api_key() is True @pytest.mark.asyncio - async def test_get_flash_auth_request_status_empty_response(self): - """get_flash_auth_request_status returns empty dict when key missing.""" + async def test_verify_api_key_returns_false_without_user(self): + """verify_api_key returns false when the current user is unavailable.""" with patch.dict(os.environ, {"RUNPOD_API_KEY": "test-key"}): from runpod_flash.core.api.runpod import RunpodGraphQLClient client = RunpodGraphQLClient() - client._execute_graphql = AsyncMock(return_value={}) + client._execute_graphql = AsyncMock(return_value={"myself": None}) - result = await client.get_flash_auth_request_status("req-missing") - assert result == {} + assert await client.verify_api_key() is False -# ── GraphQL session without API key ────────────────────────────────────── +# -- GraphQL session without API key -- @pytest.mark.serial @@ -312,11 +179,6 @@ class TestGraphQLSessionWithoutApiKey: @pytest.mark.asyncio async def test_fresh_user_no_key_anywhere(self): - """Fresh user: no env var, no credentials file -- no auth header. - - The autouse isolate_credentials_file fixture ensures no credentials - file exists and RUNPOD_API_KEY is deleted from the environment. - """ from runpod_flash.core.api.runpod import RunpodGraphQLClient client = RunpodGraphQLClient(require_api_key=False) @@ -330,7 +192,6 @@ async def test_fresh_user_no_key_anywhere(self): @pytest.mark.asyncio async def test_session_includes_auth_header_when_key_set(self): - """Session created with Authorization header when api_key is provided.""" with patch.dict(os.environ, {"RUNPOD_API_KEY": "my-key"}): from runpod_flash.core.api.runpod import RunpodGraphQLClient @@ -346,12 +207,6 @@ async def test_session_includes_auth_header_when_key_set(self): @pytest.mark.asyncio async def test_no_auth_header_when_require_api_key_false_despite_env_var(self): - """Re-login with RUNPOD_API_KEY env var set must not send auth. - - This is the exact bug scenario: flash login sets require_api_key=False - but RUNPOD_API_KEY in the environment was leaking into the session, - causing the server to see an authenticated user instead of a guest. - """ with patch.dict(os.environ, {"RUNPOD_API_KEY": "existing-key"}): from runpod_flash.core.api.runpod import RunpodGraphQLClient @@ -368,12 +223,6 @@ async def test_no_auth_header_when_require_api_key_false_despite_env_var(self): async def test_no_auth_header_when_credentials_file_has_key( self, isolate_credentials_file ): - """Re-login after prior flash login (key in credentials file) must not send auth. - - A previous successful flash login writes the API key to the credentials file. - Running flash login again must still act as a guest -- the stored key must - not leak into the session. - """ isolate_credentials_file.parent.mkdir(parents=True, exist_ok=True) isolate_credentials_file.write_text( '[default]\napi_key = "previously-saved-key"\n' @@ -394,11 +243,6 @@ async def test_no_auth_header_when_credentials_file_has_key( async def test_no_auth_header_when_both_env_var_and_credentials_file( self, isolate_credentials_file ): - """Re-login with both env var and credentials file must not send auth. - - Covers the force re-login scenario where both RUNPOD_API_KEY and a - credentials file with a stored key are present simultaneously. - """ isolate_credentials_file.parent.mkdir(parents=True, exist_ok=True) isolate_credentials_file.write_text('[default]\napi_key = "file-key"\n')