diff --git a/src/runpod_flash/cli/commands/login.py b/src/runpod_flash/cli/commands/login.py index 2483566b..5a1f50a1 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,23 @@ 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 - - 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()}") + api_key = console.input("Paste the API key shown after authorization: ").strip() - if dt.datetime.now(dt.timezone.utc) >= deadline: - raise RuntimeError("login timed out") + if not api_key: + raise RuntimeError("no api key provided") - 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) diff --git a/src/runpod_flash/core/api/runpod.py b/src/runpod_flash/core/api/runpod.py index 42c23d87..962ec68f 100644 --- a/src/runpod_flash/core/api/runpod.py +++ b/src/runpod_flash/core/api/runpod.py @@ -860,20 +860,6 @@ 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]: - query = """ - query flashAuthRequestStatus($flashAuthRequestId: String!) { - flashAuthRequestStatus(flashAuthRequestId: $flashAuthRequestId) { - id - status - expiresAt - apiKey - } - } - """ - result = await self._execute_graphql(query, {"flashAuthRequestId": request_id}) - return result.get("flashAuthRequestStatus", {}) - async def close(self): """Close the HTTP session.""" if self.session and not self.session.closed: diff --git a/src/runpod_flash/core/resources/request_logs.py b/src/runpod_flash/core/resources/request_logs.py index fa551703..fca191e1 100644 --- a/src/runpod_flash/core/resources/request_logs.py +++ b/src/runpod_flash/core/resources/request_logs.py @@ -1,3 +1,4 @@ +import json import logging import re from dataclasses import dataclass @@ -34,6 +35,67 @@ class QBRequestLogBatch: ready_worker_ids: List[str] = field(default_factory=list) +@dataclass +class SSEEvent: + id: str + data: dict[str, Any] + + +@dataclass +class LogEvent: + source: str + line: str + ts: str + + +def parse_sse_event(data: str) -> Optional[SSEEvent]: + """ + Parses an SSE line into a dictionary + """ + if not data: + return None + + try: + event_id_line, data_line = filter(bool, data.split("\n")) + event_id = event_id_line.split(":", 1)[1].strip() + data_json = data_line.split(":", 1)[1].strip() + data = json.loads(data_json) + return SSEEvent(id=event_id, data=data) + except Exception as e: + log.error("Failed to parse SSE event: %s", e) + return None + + +def parse_log_event(data: dict[str, Any]) -> Optional[LogEvent]: + """ + Parses a log event from a dictionary + """ + try: + return LogEvent(source=data["source"], line=data["line"], ts=data["ts"]) + except Exception as e: + log.error("Failed to parse log event: %s", e) + return None + + +async def stream_pod_logs(pod_id: str, tail: int = 0): + """ + Streams logs from pod using SSE + """ + if tail < 0: + raise ValueError("tail must be greater than 0") + + url = f"{RUNPOD_HAPI_URL}/v1/pod/{pod_id}/logs?stream=true&tail={tail}" + + async with get_authenticated_httpx_client() as client: + async with client.get(url) as response: + async for line in response.aiter_lines(): + event = parse_sse_event(line) + if event: + log_event = parse_log_event(event.data) + if log_event: + yield log_event + + class QBRequestLogFetcher: def __init__( self, diff --git a/tests/unit/test_login.py b/tests/unit/test_login.py index 9d061a27..df04580a 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,13 @@ 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.__aenter__.return_value = client return client @@ -85,42 +64,39 @@ 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) - - async def test_login_approved_saves_key(self, isolate_credentials_file): - mock_client = _make_mock_client(status="APPROVED", apiKey="fresh-api-key") - _login = _get_login_fn() - - with patch( - "runpod_flash.cli.commands.login.RunpodGraphQLClient", - return_value=mock_client, - ): - await _login(open_browser=False, timeout_seconds=5) + mock_console.input.return_value = "pasted-api-key" + await _login(open_browser=False) assert isolate_credentials_file.exists() - assert "fresh-api-key" in isolate_credentials_file.read_text() + assert "pasted-api-key" in isolate_credentials_file.read_text() - async def test_login_expired(self): - mock_client = _make_mock_client(status="EXPIRED", apiKey=None) + 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, ): - with pytest.raises(RuntimeError, match="login failed: expired"): - await _login(open_browser=False, timeout_seconds=5) + mock_console.input.return_value = " " + with pytest.raises(RuntimeError, match="no api key provided"): + 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 +105,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..78b5a0e3 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,18 @@ 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.__aenter__.return_value = client return client -# ── _login flow gaps ───────────────────────────────────────────────────── +# -- _login flow gaps -- @pytest.mark.serial @@ -47,7 +42,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 +50,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 = "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 @@ -205,23 +74,22 @@ def test_login_command_raises_exit_on_error(self): side_effect=RuntimeError("auth failed"), ): 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_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) + _fresh_login_module().login_command(no_open=True) -# ── 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 create_flash_auth_request.""" @pytest.mark.asyncio async def test_create_flash_auth_request(self): @@ -246,7 +114,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 @@ -262,48 +129,8 @@ async def test_create_flash_auth_request_empty_response(self): result = await client.create_flash_auth_request() assert result == {} - @pytest.mark.asyncio - async def test_get_flash_auth_request_status(self): - """get_flash_auth_request_status sends query with request_id.""" - 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, - } - } - ) - - 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" - - @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.""" - 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={}) - - result = await client.get_flash_auth_request_status("req-missing") - assert result == {} - - -# ── GraphQL session without API key ────────────────────────────────────── +# -- GraphQL session without API key -- @pytest.mark.serial @@ -312,11 +139,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 +152,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 +167,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 +183,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 +203,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')