From 5867dfa235e4b3db12837242581347181fdef95f Mon Sep 17 00:00:00 2001 From: fnwoke-lightning Date: Wed, 4 Nov 2026 13:45:52 -0500 Subject: [PATCH 01/27] chore: vendor the M4-9/M4-10 lightning_cloud changes Copy the per-request credential source and PKCE login from grid's lightning-cloud package into the SDK's vendored copy, and switch the CLI to the vendored login. --- python/lightning_sdk/api/llm_api.py | 34 +- python/lightning_sdk/cli/entrypoint.py | 70 +- python/lightning_sdk/cli/legacy/entrypoint.py | 13 +- python/lightning_sdk/cli/utils/auth.py | 381 +---------- .../lightning_cloud/cli_login.py | 634 ++++++++++++++++++ .../lightning_cloud/credentials.py | 117 ++++ python/lightning_sdk/lightning_cloud/env.py | 18 + python/lightning_sdk/lightning_cloud/login.py | 212 ++---- .../lightning_cloud/rest_client.py | 62 +- python/pyproject.toml | 1 + python/tests/api/test_agent_api.py | 10 +- python/tests/api/test_billing_api.py | 8 +- .../api/test_data_connection_credentials.py | 4 +- python/tests/api/test_filesystem_api.py | 26 +- python/tests/api/test_k8s_api.py | 18 +- python/tests/api/test_mmt_api.py | 32 +- python/tests/api/test_org_api.py | 4 +- python/tests/api/test_studio_api.py | 84 +-- python/tests/api/test_studio_start_failure.py | 2 +- python/tests/api/test_teamspace_api.py | 100 +-- python/tests/api/test_user_api.py | 28 +- python/tests/cli/test_entrypoint.py | 157 ++--- python/tests/cli/utils/test_auth.py | 434 ++---------- python/tests/conftest.py | 3 +- .../core/studio/test_studio_cloud_account.py | 4 +- .../tests/core/studio/test_studio_commands.py | 8 +- .../core/studio/test_studio_duplicate.py | 4 +- .../core/studio/test_studio_environment.py | 12 +- python/tests/core/studio/test_studio_init.py | 28 +- python/tests/core/studio/test_studio_ports.py | 4 +- .../studio/test_studio_prevent_refetch.py | 4 +- .../tests/core/studio/test_studio_rename.py | 2 +- python/tests/core/studio/test_studio_start.py | 14 +- .../tests/core/studio/test_studio_status.py | 2 +- python/tests/core/studio/test_studio_vm.py | 2 +- python/tests/core/test_agents.py | 6 +- python/tests/core/test_deployment.py | 62 +- python/tests/core/test_job.py | 118 ++-- python/tests/core/test_k8s_cluster.py | 10 +- python/tests/core/test_llm.py | 30 +- python/tests/core/test_mmt.py | 50 +- python/tests/core/test_models.py | 32 +- python/tests/core/test_org.py | 18 +- .../core/test_permissions_integration.py | 6 +- python/tests/core/test_sandbox.py | 92 +-- python/tests/core/test_sandbox_config.py | 18 +- python/tests/core/test_teamspace.py | 116 ++-- python/tests/core/test_user.py | 28 +- .../tests/lightning_cloud/test_cli_login.py | 114 ++++ .../tests/lightning_cloud/test_credentials.py | 77 +++ python/tests/lightning_cloud/test_login.py | 18 +- python/tests/utils/test_resolve.py | 12 +- 52 files changed, 1763 insertions(+), 1580 deletions(-) create mode 100644 python/lightning_sdk/lightning_cloud/cli_login.py create mode 100644 python/lightning_sdk/lightning_cloud/credentials.py create mode 100644 python/tests/lightning_cloud/test_cli_login.py create mode 100644 python/tests/lightning_cloud/test_credentials.py diff --git a/python/lightning_sdk/api/llm_api.py b/python/lightning_sdk/api/llm_api.py index 712bf1933..501ae6281 100644 --- a/python/lightning_sdk/api/llm_api.py +++ b/python/lightning_sdk/api/llm_api.py @@ -4,16 +4,15 @@ import json import os import threading -import time import warnings from typing import Any, AsyncGenerator, Dict, Generator, List, Optional, Union -from urllib.parse import urlencode from rich.prompt import Confirm from lightning_sdk.api.utils import cached_lightning_client -from lightning_sdk.lightning_cloud import env -from lightning_sdk.lightning_cloud.login import Auth, AuthServer +from lightning_sdk.lightning_cloud import cli_login +from lightning_sdk.lightning_cloud.credentials import browser_login +from lightning_sdk.lightning_cloud.login import Auth from lightning_sdk.lightning_cloud.openapi.models import ( StreamResultOfV1ConversationResponseChunk, V1ConversationResponseChunk, @@ -22,19 +21,6 @@ V1ResponseChoiceDelta, ) -LITAI_CODE = os.environ.get("LITAI_CODE", "x334uv8t7v") - - -class _AuthServer(AuthServer): - def __init__(self, model: str, *args: Any, **kwargs: Any) -> None: - self._model = model - super().__init__(*args, **kwargs) - - def get_auth_url(self, port: int) -> str: - redirect_uri = f"http://localhost:{port}/login-complete" - params = urlencode({"redirectTo": redirect_uri, "okbhrt": LITAI_CODE, "litaimodel": self._model}) - return f"{env.LIGHTNING_CLOUD_URL}/sign-in?{params}" - class _AuthLitAI(Auth): def __init__(self, model: str, shall_confirm: bool = False) -> None: @@ -55,15 +41,19 @@ def _run_server(self) -> None: "Login cancelled. Please login at https://platform.lightning.ai/sign-up. " "Or run `lightning login` to login." ) - print("Opening browser for authentication...") - print("Please come back to the terminal after logging in.") - time.sleep(1) - _AuthServer(self._model).login_with_browser(self) + cli_login.login() def authenticate(model: str, shall_confirm: bool = True) -> None: auth = _AuthLitAI(model, shall_confirm) - auth.authenticate() + with browser_login(False): + try: + auth.authenticate() + except ValueError: + pass + else: + return + auth._run_server() class LLMApi: diff --git a/python/lightning_sdk/cli/entrypoint.py b/python/lightning_sdk/cli/entrypoint.py index 3cb4d8ee7..313636915 100644 --- a/python/lightning_sdk/cli/entrypoint.py +++ b/python/lightning_sdk/cli/entrypoint.py @@ -2,7 +2,10 @@ import os import sys +from contextlib import suppress +from typing import Optional +import requests import rich_click as click from lightning_sdk import __version__ @@ -45,13 +48,9 @@ from lightning_sdk.cli.lockdown import apply_lockdown from lightning_sdk.cli.ls import ls from lightning_sdk.cli.rm import rm -from lightning_sdk.cli.utils.auth import ( - browser_authentication, - login_with_pasted_code, - logout_session, - oidc_login_enabled, -) +from lightning_sdk.cli.utils.auth import browser_authentication from lightning_sdk.cli.utils.logging import CommandLoggingGroup, logging_excepthook +from lightning_sdk.lightning_cloud import cli_login from lightning_sdk.lightning_cloud.login import Auth from lightning_sdk.utils.resolve import _get_authed_user, in_studio @@ -104,13 +103,22 @@ def main_cli(ctx: click.Context) -> None: @click.option( "--no-browser", is_flag=True, - hidden=True, help="Sign in without a browser on this machine: approve on another device, then paste the code it shows.", ) -def login(no_browser: bool) -> None: +@click.option( + "--org", + "orgs", + metavar="A,B", + help="Comma-separated organizations to ask for when the login is approved.", +) +@click.option("--project", help="Teamspace to ask for when the login is approved.") +@click.option( + "--insecure-storage", + is_flag=True, + help="Store the login in a file only you can read when no OS keychain is available.", +) +def login(no_browser: bool, orgs: Optional[str], project: Optional[str], insecure_storage: bool) -> None: """Sign in to Lightning AI.""" - if no_browser and not oidc_login_enabled(): - raise click.UsageError("`lightning login --no-browser` is not available yet.") auth = Auth() if auth.api_key or auth.auth_token: click.echo('"lightning login" is not required within a Studio or when credentials are set in the environment') @@ -120,21 +128,18 @@ def login(no_browser: bool) -> None: # this is unexpected, as we automatically auth within a Studio raise RuntimeError("Unable to login within a Studio. Did you change your shell setup?") from None - # always start a fresh session: a saved one may have expired - auth.clear() + org_list = [org.strip() for org in orgs.split(",") if org.strip()] if orgs else None try: - if no_browser: - login_with_pasted_code(auth) - else: - with browser_authentication(True): - auth.authenticate() + cli_login.login( + no_browser=no_browser, + insecure_storage=insecure_storage, + orgs=org_list, + project=project, + ) + except cli_login.LoginError: + raise except ConnectionError: raise RuntimeError(f"Unable to connect to {_cloud_url()}. Please check your internet connection.") from None - except ValueError: - # the browser flow ended without the callback saving a token (see the warning it logs) - raise RuntimeError( - "Login failed: the browser sign-in finished without returning a token. Try running `lightning login` again" - ) from None try: click.echo(f'Successfully logged in as "{_get_authed_user().name}".') @@ -143,13 +148,32 @@ def login(no_browser: bool) -> None: click.echo("Successfully logged in.") +def _saved_session_token() -> Optional[str]: + """A `sess_` token saved by an older SDK, so logout can revoke it too.""" + auth = Auth() + try: + if auth.load() and isinstance(auth.auth_token, str) and auth.auth_token.startswith("sess_"): + return auth.auth_token + except Exception: + pass + return None + + @main_cli.command def logout() -> None: """Sign out of the current session.""" + legacy = _saved_session_token() try: - logout_session(Auth()) + cli_login.logout() except Exception as e: raise RuntimeError(f"Logout failed: {e}") from e + if legacy: + with suppress(requests.RequestException): + requests.post( + f"{_cloud_url().rstrip('/')}/v1/auth/logout", + headers={"Authorization": f"Bearer {legacy}"}, + timeout=10, + ) click.echo("Successfully logged out.") diff --git a/python/lightning_sdk/cli/legacy/entrypoint.py b/python/lightning_sdk/cli/legacy/entrypoint.py index f36108b72..6e9cdc6c4 100644 --- a/python/lightning_sdk/cli/legacy/entrypoint.py +++ b/python/lightning_sdk/cli/legacy/entrypoint.py @@ -30,7 +30,7 @@ from lightning_sdk.cli.utils.auth import browser_authentication from lightning_sdk.cli.utils.coloring import CustomHelpFormatter from lightning_sdk.constants import _LIGHTNING_DEBUG -from lightning_sdk.lightning_cloud.login import Auth +from lightning_sdk.lightning_cloud import cli_login def _notify_exception(exception_type: Type[BaseException], value: BaseException, tb: Optional[TracebackType]) -> None: @@ -73,12 +73,10 @@ def main_cli(ctx: click.Context) -> None: @main_cli.command def login() -> None: """Login to Lightning AI Studios.""" - auth = Auth() - auth.clear() - try: - with browser_authentication(True): - auth.authenticate() + cli_login.login() + except cli_login.LoginError: + raise except ConnectionError: raise RuntimeError(f"Unable to connect to {_cloud_url()}. Please check your internet connection.") from None @@ -86,8 +84,7 @@ def login() -> None: @main_cli.command def logout() -> None: """Logout from Lightning AI Studios.""" - auth = Auth() - auth.clear() + cli_login.logout() # additional commands diff --git a/python/lightning_sdk/cli/utils/auth.py b/python/lightning_sdk/cli/utils/auth.py index 5de4c68fd..8555a328d 100644 --- a/python/lightning_sdk/cli/utils/auth.py +++ b/python/lightning_sdk/cli/utils/auth.py @@ -1,384 +1,27 @@ -"""Non-interactive authentication helpers for CLI commands.""" +"""Authentication helpers shared by the CLI commands.""" -import base64 -import hashlib -import hmac -import json -import os -import secrets -from contextlib import contextmanager, suppress -from contextvars import ContextVar -from pathlib import Path -from typing import Generator, Optional, Tuple -from urllib.parse import urlencode +from contextlib import contextmanager +from typing import Generator -import requests import rich_click as click -from lightning_sdk.lightning_cloud import env -from lightning_sdk.lightning_cloud.login import Auth, AuthServer, logger - -_BROWSER_AUTH_ALLOWED = ContextVar("lightning_browser_auth_allowed", default=True) -_ORIGINAL_AUTH_SAVE = Auth.save - -_CLIENT_ID = "lightning-cli" -_PASTE_ATTEMPTS = 3 - -# The authorization-code login needs GRID_AUTH_CLI_LOGIN_ENABLED on the control plane, which production -# doesn't set yet. Until it does, `lightning login` uses the session-token login, and this opts in to test it. -OIDC_LOGIN_ENV = "LIGHTNING_CLI_OIDC_LOGIN" - - -def oidc_login_enabled() -> bool: - return env.truthy(os.getenv(OIDC_LOGIN_ENV)) - - -def _generate_state() -> str: - return base64.urlsafe_b64encode(os.urandom(16)).rstrip(b"=").decode("ascii") - - -def _origin() -> str: - """Platform origin without a trailing slash, so paths are not `//cli-login`.""" - return env.LIGHTNING_CLOUD_URL.rstrip("/") - - -def _pkce_pair() -> Tuple[str, str]: - """Return a PKCE `(verifier, S256 challenge)` pair.""" - verifier = secrets.token_urlsafe(64) - challenge = base64.urlsafe_b64encode(hashlib.sha256(verifier.encode("ascii")).digest()).rstrip(b"=") - return verifier, challenge.decode("ascii") - - -def _authorize_url(redirect_uri: str, state: str, challenge: str) -> str: - query = urlencode( - { - "response_type": "code", - "client_id": _CLIENT_ID, - "redirect_uri": redirect_uri, - "scope": "openid", - "state": state, - "code_challenge": challenge, - "code_challenge_method": "S256", - } - ) - return f"{_origin()}/oidc/authorize?{query}" - - -class _CodeRejectedError(Exception): - """The token endpoint refused the code; `str(e)` is its OAuth error.""" - - -def _exchange_code(code: str, verifier: str, redirect_uri: str) -> str: - """Trade an approved code and its PKCE verifier for a CLI session token.""" - try: - response = requests.post( - f"{_origin()}/oidc/token", - data={ - "grant_type": "authorization_code", - "client_id": _CLIENT_ID, - "code": code, - "code_verifier": verifier, - "redirect_uri": redirect_uri, - }, - timeout=10, - ) - except requests.ConnectionError as e: - raise ConnectionError(str(e)) from e - body = _json_body(response) - if response.ok and body.get("access_token"): - return body["access_token"] - raise _CodeRejectedError(body.get("error") or f"HTTP {response.status_code}") - - -def _save_with_session_token( - self: Auth, - token: str = "", - user_id: str = "", - api_key: str = "", - auth_token: str = "", -) -> None: - """Persist session tokens (`token`) as `auth_token` so the SDK/CLI sends `Bearer sess_...`.""" - if not auth_token and token and (not api_key or token.startswith("sess_")): - auth_token = token - if token.startswith("sess_"): - api_key = "" - _ORIGINAL_AUTH_SAVE( - self, - token=token, - user_id=user_id, - api_key=api_key, - auth_token=auth_token, - ) - with suppress(OSError): - os.chmod(self.secrets_file, 0o600) - - -def _login_with_authorization_code(self: AuthServer, auth: Auth) -> None: - """Sign in with an authorization code and PKCE, delivered to a callback on 127.0.0.1.""" - import webbrowser - - import uvicorn - from fastapi import FastAPI - from starlette.background import BackgroundTask - from starlette.responses import RedirectResponse - - from lightning_sdk.lightning_cloud.utils.network import find_free_network_port - - expected_state = _generate_state() - verifier, challenge = _pkce_pair() - app = FastAPI() - port = find_free_network_port() - redirect_uri = f"http://127.0.0.1:{port}/callback" - url = _authorize_url(redirect_uri, expected_state, challenge) - - try: - requests.head(url, timeout=10) - except requests.ConnectionError as e: - raise requests.ConnectionError( - f"No internet connection available. Please connect to a stable internet connection \n{e}" - ) from e - except requests.RequestException as e: - raise requests.RequestException( - f"An error occurred with the request. Please report this issue to Lightning Team \n{e}" - ) from e - - logger.info(f"login started for lightning.ai, opening {url}") - if not webbrowser.open(url): - if os.environ.get("LIGHTNING_DEPLOYMENT_ID"): - raise RuntimeError( - "Failed to authenticate to Lightning. Ensure that you have selected 'Include SDK credentials' in the " - "'Environment' section of the deployment settings." - ) - click.echo("No browser could be opened on this machine.") - login_with_pasted_code(auth) - return - - @app.get("/callback") - def save_token(code: str = "", state: str = "", error: str = "") -> RedirectResponse: - failure = "" - if not state or not hmac.compare_digest(state, expected_state): - failure = "state mismatch" - elif error or not code: - failure = error or "no code in the callback" - else: - try: - auth.save(auth_token=_exchange_code(code, verifier, redirect_uri)) - except (_CodeRejectedError, ConnectionError) as e: - failure = str(e) - - if failure: - logger.warning(f"Authentication Failed ({failure}). Run `lightning login` again, or update the CLI.") - target = f"{_origin()}/cli-login-failed" - else: - logger.info("Authentication Successful") - target = f"{_origin()}/cli-login-successful" - - return RedirectResponse( - url=target, - background=BackgroundTask(stop_server), - ) - - def stop_server() -> None: - server.should_exit = True - - server = uvicorn.Server(config=uvicorn.Config(app, host="127.0.0.1", port=port, log_level="error")) - server.run() - - -def _login_with_session_token(self: AuthServer, auth: Auth) -> None: - """Sign in through the sign-in page, which hands a session token to a callback on localhost.""" - import webbrowser - - import uvicorn - from fastapi import FastAPI, Query - from starlette.background import BackgroundTask - from starlette.responses import RedirectResponse - - from lightning_sdk.lightning_cloud.utils.network import find_free_network_port - - expected_state = _generate_state() - app = FastAPI() - port = find_free_network_port() - base_url = self.get_auth_url(port) - redirect_uri = f"http://localhost:{port}/login-complete" - if urlencode({"redirectTo": redirect_uri}) in base_url: - stateful_redirect = f"{redirect_uri}?{urlencode({'state': expected_state})}" - url = base_url.replace(urlencode({"redirectTo": redirect_uri}), urlencode({"redirectTo": stateful_redirect}), 1) - else: - url = base_url - - try: - requests.head(url, timeout=10) - except requests.ConnectionError as e: - raise requests.ConnectionError( - f"No internet connection available. Please connect to a stable internet connection \n{e}" - ) from e - except requests.RequestException as e: - raise requests.RequestException( - f"An error occurred with the request. Please report this issue to Lightning Team \n{e}" - ) from e - - logger.info(f"login started for lightning.ai, opening {url}") - if not webbrowser.open(url): - deployment_id = os.environ.get("LIGHTNING_DEPLOYMENT_ID") - if deployment_id: - raise RuntimeError( - "Failed to authenticate to Lightning. Ensure that you have selected 'Include SDK credentials' in the " - "'Environment' section of the deployment settings." - ) - raise RuntimeError( - "Failed to authenticate to Lightning. When running without access to a browser, " - "'LIGHTNING_API_KEY' should be exported." - ) - - @app.get("/login-complete") - async def save_token( - token: str = "", - key: str = "", - state: str = "", - user_id: str = Query("", alias="userID"), - ) -> RedirectResponse: - state_valid = bool(state) and hmac.compare_digest(state, expected_state) - if token and state_valid: - auth.save(token=token, user_id=user_id, api_key=key) - logger.info("Authentication Successful") - target = f"{env.LIGHTNING_CLOUD_URL}/cli-login-successful" - else: - reason = "no token in the callback" if not token else "state mismatch" - logger.warning( - f"Authentication Failed ({reason}). If the browser was already signed in, the sign-in page " - "could not mint a CLI token. Run `lightning login` again, or update the CLI." - ) - target = f"{env.LIGHTNING_CLOUD_URL}/cli-login-failed" - - return RedirectResponse( - url=target, - background=BackgroundTask(stop_server), - ) - - def stop_server() -> None: - server.should_exit = True - - server = uvicorn.Server(config=uvicorn.Config(app, port=port, log_level="error")) - server.run() - - -def _login_with_browser(self: AuthServer, auth: Auth) -> None: - if oidc_login_enabled(): - _login_with_authorization_code(self, auth) - else: - _login_with_session_token(self, auth) - - -Auth.save = _save_with_session_token -AuthServer.login_with_browser = _login_with_browser - - -def login_with_pasted_code(auth: Auth) -> None: - """Sign in on a machine without a browser: approve on any device, then paste the code shown there.""" - redirect_uri = f"{_origin()}/cli-login" - verifier, challenge = _pkce_pair() - click.echo( - "To sign in, open this address in a browser on any device and approve the request:\n\n" - f" {_authorize_url(redirect_uri, _generate_state(), challenge)}\n\n" - "The page then shows a code. Only paste a code from a page you opened this way; " - "never give it to anyone, as it gives full access to your Lightning account." - ) - for attempt in range(_PASTE_ATTEMPTS): - code = click.prompt("Code").strip() - try: - auth.save(auth_token=_exchange_code(code, verifier, redirect_uri)) - return - except _CodeRejectedError as e: - if str(e) != "invalid_grant": - raise RuntimeError(f"Login failed: {e}.") from None - if attempt < _PASTE_ATTEMPTS - 1: - click.echo("That code was not accepted. It works once and expires after a minute. Try again.") - raise RuntimeError("Login failed: the code was not accepted. Run `lightning login --no-browser` again.") - - -def _json_body(response: requests.Response) -> dict: - try: - body = response.json() - except ValueError: - return {} - return body if isinstance(body, dict) else {} - - -def logout_session(auth: Auth) -> None: - """Revoke the active CLI session token on the server and clear local credentials.""" - token = getattr(auth, "auth_token", None) - if not isinstance(token, str) or not token: - secrets_file = getattr(auth, "secrets_file", None) - if isinstance(secrets_file, Path) and secrets_file.exists(): - try: - with secrets_file.open() as creds: - data = json.load(creds) - loaded = data.get("auth_token") - if isinstance(loaded, str) and loaded: - token = loaded - except Exception: - token = None - if isinstance(token, str) and token: - with suppress(requests.RequestException): - requests.post( - f"{_origin()}/v1/auth/logout", - headers={"Authorization": f"Bearer {token}"}, - timeout=10, - ) - auth.clear() +from lightning_sdk.lightning_cloud.credentials import browser_login +from lightning_sdk.lightning_cloud.login import Auth @contextmanager def browser_authentication(allowed: bool) -> Generator[None, None, None]: """Temporarily control whether missing credentials may start browser authentication.""" - token = _BROWSER_AUTH_ALLOWED.set(allowed) - authenticate = Auth.authenticate - - def guarded_authenticate(auth: Auth) -> Optional[str]: - if _BROWSER_AUTH_ALLOWED.get(): - return authenticate(auth) - - has_own_run_server = "_run_server" in auth.__dict__ - run_server = auth.__dict__.get("_run_server") - - def blocked_run_server() -> None: - raise ValueError("No Lightning credentials are available. Run `lightning login` first.") - - auth.__dict__["_run_server"] = blocked_run_server - try: - return authenticate(auth) - finally: - if has_own_run_server: - auth.__dict__["_run_server"] = run_server - else: - auth.__dict__.pop("_run_server", None) - - Auth.authenticate = guarded_authenticate - try: + with browser_login(allowed): yield - finally: - Auth.authenticate = authenticate - _BROWSER_AUTH_ALLOWED.reset(token) def require_auth_header() -> str: """Return available credentials without falling back to browser authentication.""" - auth = Auth() try: - header = auth.auth_header - except ValueError: - header = None - - if header: - return header - - if auth.load(): - try: - header = auth.auth_header - except ValueError: - header = None - if header: - return header - - raise click.UsageError("No Lightning credentials are available. Run `lightning login` first.") + header = Auth().auth_header + except (ValueError, ConnectionError) as err: + raise click.UsageError(str(err)) from err + if not header: + raise click.UsageError("No Lightning credentials are available. Run `lightning login` first.") + return header diff --git a/python/lightning_sdk/lightning_cloud/cli_login.py b/python/lightning_sdk/lightning_cloud/cli_login.py new file mode 100644 index 000000000..fa03ec7d4 --- /dev/null +++ b/python/lightning_sdk/lightning_cloud/cli_login.py @@ -0,0 +1,634 @@ +# This file is vendored. +# Do not edit it directly; changes will be overwritten. +# Make changes in the internal upstream repository instead. + +import base64 +import contextlib +import dataclasses +import getpass +import hashlib +import hmac +import json +import os +import pathlib +import re +import secrets +import sys +import tempfile +import time +import webbrowser +from datetime import datetime +from typing import Any, Dict, Iterator, Optional, Sequence +from urllib.parse import parse_qs, urlencode, urlsplit + +import requests + +from lightning_sdk.lightning_cloud import env +from lightning_sdk.lightning_cloud.login import Auth + +CLIENT_ID = "lightning-cli" +_KEYRING_SERVICE = CLIENT_ID +_REFRESH_MARGIN = 300.0 +_PASTE_TRIES = 3 + +_NOT_LOGGED_IN = "You are not logged in. Run `lightning login`." +_LOGIN_ENDED = "Your login has expired or was revoked. Run `lightning login`." +_LOGIN_REUSED = ( + "This login was revoked because a copy of it was used elsewhere. " + "Run `lightning login` and consider rotating any keys this machine held." +) +_NO_KEYCHAIN = ( + "No OS keychain is available to store the login. Run " + "`lightning login --insecure-storage` to keep it in a file only you can read." +) +_PASTE_WARNING = ( + "The code gives access to your Lightning account. Paste it only into the " + "CLI you just started, never into a website, form, chat or email." +) +_DEPLOYMENT = ( + "Failed to authenticate to Lightning. Ensure that you have selected " + "'Include SDK credentials' in the 'Environment' section of the deployment settings." +) + + +class LoginError(ConnectionError): + """A login could not be made, used or ended. The message says what to do next.""" + + +class _Unreachable(LoginError): + pass + + +class _Refused(LoginError): + def __init__(self, error: str, description: str) -> None: + super().__init__(f"Lightning refused the login: {description or error}.") + self.error = error + self.description = description + + +@dataclasses.dataclass(frozen=True) +class LoginStatus: + url: str + storage: str + access_expires_at: float + refresh_expires_at: float + grant: Dict[str, Any] + + +@dataclasses.dataclass(frozen=True) +class Login(LoginStatus): + access_token: str = dataclasses.field(repr=False) + refresh_token: str = dataclasses.field(repr=False) + + +def login( + no_browser: bool = False, + insecure_storage: bool = False, + orgs: Optional[Sequence[str]] = None, + project: Optional[str] = None, +) -> Login: + """Log in through the browser, or by pasting the code the login page shows.""" + if insecure_storage: + storage = _FileStorage.kind + elif _keychain_available(): + storage = _KeychainStorage.kind + else: + raise LoginError(_NO_KEYCHAIN) + attempt = _Attempt(_origin(), storage, orgs, project) + if not no_browser: + done = _browser_login(attempt) + if done is not None: + return done + if os.environ.get("LIGHTNING_DEPLOYMENT_ID"): + raise LoginError(_DEPLOYMENT) + return _paste_back_login(attempt) + + +def token(min_remaining: float = 1800) -> str: + """An access token with at least ``min_remaining`` seconds left, refreshed first if needed.""" + login = _fresh_login(_origin(), min_remaining) + if login is None: + raise LoginError(_NOT_LOGGED_IN) + if login.access_expires_at - time.time() < min_remaining: + raise LoginError( + f"Couldn't refresh the login at {login.url}. Check your connection and try again." + ) + return login.access_token + + +def current_login() -> Optional[LoginStatus]: + """The saved login for this Lightning URL, without its tokens.""" + login = _load(_origin()) + if login is None: + return None + return LoginStatus( + url=login.url, + storage=login.storage, + access_expires_at=login.access_expires_at, + refresh_expires_at=login.refresh_expires_at, + grant=login.grant, + ) + + +def logout() -> None: + """Revoke the saved login on the server, then forget it on this machine.""" + with _locked(): + login = _load(_origin(), cached=False) + revoked = login is None or _revoke(login.refresh_token) + if login is not None: + _forget() + Auth.secrets_file.unlink(missing_ok=True) + if not revoked: + raise LoginError( + "Logged out on this machine, but Lightning couldn't be reached to " + "revoke the login. It may stay active until it expires." + ) + + +def _request_token(rejected: Optional[str] = None) -> Optional[str]: + """The saved login's access token for a request, or None without a login. + + ``rejected`` is a token the server refused, which forces a refresh unless + another process has already replaced it. + """ + login = _fresh_login(_origin(), _REFRESH_MARGIN, rejected) + return None if login is None else login.access_token + + +class _Attempt: + """One login: its state, its PKCE verifier and where it will be stored.""" + + def __init__( + self, + origin: str, + storage: str, + orgs: Optional[Sequence[str]], + project: Optional[str], + ) -> None: + self.origin = origin + self.storage = storage + self.state = secrets.token_urlsafe(32) + self._verifier = secrets.token_urlsafe(48) + self._hints = {} + if orgs: + self._hints["orgs_hint"] = orgs if isinstance(orgs, str) else ",".join(orgs) + if project: + self._hints["project_hint"] = project + + def authorize_url(self, redirect_uri: str) -> str: + digest = hashlib.sha256(self._verifier.encode()).digest() + query = { + "response_type": "code", + "client_id": CLIENT_ID, + "redirect_uri": redirect_uri, + "scope": "openid", + "state": self.state, + "code_challenge": base64.urlsafe_b64encode(digest).rstrip(b"=").decode(), + "code_challenge_method": "S256", + **self._hints, + } + return f"{self.origin}/oidc/authorize?{urlencode(query)}" + + def complete(self, code: str, redirect_uri: str) -> Login: + body = _token_endpoint({ + "grant_type": "authorization_code", + "code": code, + "redirect_uri": redirect_uri, + "code_verifier": self._verifier, + "storage": self.storage, + }) + new = _login_from(body, self.storage, self.origin) + try: + with _locked(): + _save(new) + except LoginError: + _revoke(new.refresh_token) + raise + Auth.secrets_file.unlink(missing_ok=True) + return new + + +def _browser_login(attempt: _Attempt) -> Optional[Login]: + with _callback_server(attempt) as (redirect_uri, wait): + url = attempt.authorize_url(redirect_uri) + if not webbrowser.open(url): + return None + _say( + f"Opening {url} in your browser to log in. " + "If it doesn't open, run `lightning login --no-browser`." + ) + return wait() + + +@contextlib.contextmanager +def _callback_server(attempt: _Attempt) -> Iterator[Any]: + import http.server + import socketserver + + outcome = [] + + class Callback(http.server.BaseHTTPRequestHandler): + def do_GET(self) -> None: + parts = urlsplit(self.path) + if parts.path != "/callback": + self.send_error(404) + return + query = {key: values[0] for key, values in parse_qs(parts.query).items()} + try: + outcome.append(_callback_login(attempt, query, redirect_uri)) + page = "cli-login-successful" + except Exception as err: + outcome.append(err) + page = "cli-login-failed" + self.send_response(302) + self.send_header("Location", f"{attempt.origin}/{page}") + self.end_headers() + + def log_message(self, format: str, *args: Any) -> None: + # The default logs the request line, which carries the code. + pass + + with socketserver.TCPServer(("127.0.0.1", 0), Callback) as server: + # A short timeout lets Ctrl-C interrupt the wait. + server.timeout = 0.5 + redirect_uri = f"http://127.0.0.1:{server.server_address[1]}/callback" + + def wait() -> Login: + while not outcome: + server.handle_request() + if isinstance(outcome[0], Exception): + raise outcome[0] + return outcome[0] + + yield redirect_uri, wait + + +def _callback_login(attempt: _Attempt, query: Dict[str, str], redirect_uri: str) -> Login: + if not hmac.compare_digest(query.get("state", "").encode(), attempt.state.encode()): + raise LoginError( + "The login callback didn't match this login, so nothing was saved. " + "Run `lightning login` again." + ) + if "error" in query: + raise LoginError("The login was denied, so nothing was saved.") + return attempt.complete(query.get("code", ""), redirect_uri) + + +def _paste_back_login(attempt: _Attempt) -> Login: + redirect_uri = f"{attempt.origin}/cli-login" + _say( + "Open this URL on any device to log in:\n\n" + f" {attempt.authorize_url(redirect_uri)}\n\n{_PASTE_WARNING}\n" + ) + for _ in range(_PASTE_TRIES): + try: + code = getpass.getpass("Paste the code from the login page: ").strip() + except EOFError: + raise LoginError("No code was entered. Run `lightning login` to try again.") from None + if not code: + _say("No code was entered.") + continue + try: + return attempt.complete(code, redirect_uri) + except _Refused as err: + if err.error != "invalid_grant": + raise + _say("That code didn't work. It may have expired or been mistyped.") + raise LoginError( + f"No code worked after {_PASTE_TRIES} tries. Run `lightning login` to start again." + ) + + +def _fresh_login(url: str, min_remaining: float, rejected: Optional[str] = None) -> Optional[Login]: + login = _load(url) + if login is None or _usable(login, min_remaining, rejected): + return login + # Refresh tokens rotate, and presenting a retired one revokes the whole + # login. So processes sharing it refresh one at a time, and each re-reads + # the login under the lock in case another just refreshed it. + with _locked(): + login = _load(url, cached=False) + if login is None or _usable(login, min_remaining, rejected): + return login + try: + body = _token_endpoint({ + "grant_type": "refresh_token", + "refresh_token": login.refresh_token, + }) + except _Unreachable: + if rejected is None and login.access_expires_at > time.time(): + return login + raise + except _Refused as err: + _forget() + message = _LOGIN_REUSED if "reused" in err.description else _LOGIN_ENDED + raise LoginError(message) from None + fresh = _login_from(body, login.storage, url) + _save(fresh) + return fresh + + +def _usable(login: Login, min_remaining: float, rejected: Optional[str]) -> bool: + if rejected is not None: + return login.access_token != rejected + return login.access_expires_at - time.time() >= min_remaining + + +def _keychain_available() -> bool: + try: + import keyring + + return keyring.get_keyring().priority > 0 + except Exception: + return False + + +class _KeychainStorage: + """Tokens in the OS keychain, read only when the login file's version changes.""" + + kind = "keychain" + + def __init__(self) -> None: + self._cache = None + + def write(self, url: str, tokens: Dict[str, str]) -> Dict[str, str]: + import keyring + + try: + keyring.set_password(_KEYRING_SERVICE, url, json.dumps(tokens)) + except Exception as err: + raise LoginError(_NO_KEYCHAIN) from err + return {} + + def read(self, url: str, metadata: Dict[str, Any], cached: bool) -> Optional[Dict[str, str]]: + key = (url, metadata.get("version")) + if cached and self._cache is not None and self._cache[0] == key: + return self._cache[1] + import keyring + + try: + tokens = _tokens(json.loads(keyring.get_password(_KEYRING_SERVICE, url) or "null")) + except Exception: + return None + if tokens is not None and tokens["version"] == key[1]: + self._cache = (key, tokens) + return tokens + + def forget(self, url: str) -> None: + self._cache = None + import keyring + + with contextlib.suppress(Exception): + keyring.delete_password(_KEYRING_SERVICE, url) + + +class _FileStorage: + """Tokens in the login file itself, for --insecure-storage.""" + + kind = "file" + + def write(self, url: str, tokens: Dict[str, str]) -> Dict[str, str]: + return tokens + + def read(self, url: str, metadata: Dict[str, Any], cached: bool) -> Optional[Dict[str, str]]: + return _tokens(metadata) + + def forget(self, url: str) -> None: + pass + + +_STORAGES = {storage.kind: storage for storage in (_KeychainStorage(), _FileStorage())} + + +def _tokens(data: Any) -> Optional[Dict[str, str]]: + if not isinstance(data, dict): + return None + tokens = {key: data.get(key) for key in ("version", "access_token", "refresh_token")} + if not all(isinstance(value, str) and value for value in tokens.values()): + return None + return tokens + + +def _load(url: str, cached: bool = True) -> Optional[Login]: + metadata = _read_metadata() + if metadata is None or metadata.get("url") != url: + return None + storage = _STORAGES.get(metadata.get("storage")) + if storage is None: + return None + tokens = storage.read(url, metadata, cached) + if tokens is None: + return None + try: + access_expires_at = float(metadata["access_expires_at"]) + refresh_expires_at = float(metadata["refresh_expires_at"]) + except (KeyError, TypeError, ValueError): + return None + if tokens["version"] != metadata.get("version"): + # A save stopped between the keychain and the login file. The keychain + # holds the newer pair, so refresh with it to rewrite both. + access_expires_at = 0.0 + return Login( + url=url, + storage=storage.kind, + access_expires_at=access_expires_at, + refresh_expires_at=refresh_expires_at, + grant=metadata.get("grant") or {}, + access_token=tokens["access_token"], + refresh_token=tokens["refresh_token"], + ) + + +def _save(login: Login) -> None: + previous = _read_metadata() + if previous and (previous.get("url"), previous.get("storage")) != (login.url, login.storage): + _forget_secrets(previous) + version = secrets.token_hex(16) + tokens = { + "version": version, + "access_token": login.access_token, + "refresh_token": login.refresh_token, + } + stored_here = _STORAGES[login.storage].write(login.url, tokens) + _write_metadata({ + "version": version, + "url": login.url, + "storage": login.storage, + "access_expires_at": login.access_expires_at, + "refresh_expires_at": login.refresh_expires_at, + "grant": login.grant, + **stored_here, + }) + + +def _forget() -> None: + metadata = _read_metadata() + _retrying(_unlink, _login_path()) + if metadata is not None: + _forget_secrets(metadata) + + +def _forget_secrets(metadata: Dict[str, Any]) -> None: + storage = _STORAGES.get(metadata.get("storage")) + if storage is not None and isinstance(metadata.get("url"), str): + storage.forget(metadata["url"]) + + +def _login_path() -> pathlib.Path: + return pathlib.Path(env.LIGHTNING_LOGIN_PATH) + + +def _read_metadata() -> Optional[Dict[str, Any]]: + try: + data = json.loads(_login_path().read_text()) + except (OSError, ValueError): + return None + return data if isinstance(data, dict) else None + + +def _write_metadata(data: Dict[str, Any]) -> None: + path = _login_path() + path.parent.mkdir(parents=True, exist_ok=True) + # mkstemp creates the file readable only by the user, and the rename keeps that. + fd, tmp = tempfile.mkstemp(dir=path.parent, prefix=f".{path.name}.") + try: + with os.fdopen(fd, "w") as file: + json.dump(data, file) + _retrying(os.replace, tmp, path) + except BaseException: + with contextlib.suppress(OSError): + os.unlink(tmp) + raise + + +def _unlink(path: pathlib.Path) -> None: + path.unlink(missing_ok=True) + + +def _retrying(operation: Any, *args: Any) -> Any: + for _ in range(50): + try: + return operation(*args) + except PermissionError: + # Windows refuses to replace or delete a file another process has open. + time.sleep(0.02) + return operation(*args) + + +@contextlib.contextmanager +def _locked() -> Iterator[None]: + path = _login_path() + path.parent.mkdir(parents=True, exist_ok=True) + fd = os.open(f"{path}.lock", os.O_RDWR | os.O_CREAT, 0o600) + try: + _lock(fd) + try: + yield + finally: + _unlock(fd) + finally: + os.close(fd) + + +if os.name == "nt": + import msvcrt + + def _lock(fd: int) -> None: + while True: + try: + msvcrt.locking(fd, msvcrt.LK_NBLCK, 1) + return + except OSError: + time.sleep(0.05) + + def _unlock(fd: int) -> None: + msvcrt.locking(fd, msvcrt.LK_UNLCK, 1) +else: + import fcntl + + def _lock(fd: int) -> None: + fcntl.flock(fd, fcntl.LOCK_EX) + + def _unlock(fd: int) -> None: + fcntl.flock(fd, fcntl.LOCK_UN) + + +def _origin() -> str: + return env.LIGHTNING_CLOUD_URL.rstrip("/") + + +def _post(path: str, form: Dict[str, str]) -> requests.Response: + origin = _origin() + try: + return requests.post( + f"{origin}{path}", + data={"client_id": CLIENT_ID, **form}, + headers={"User-Agent": env.USER_AGENT}, + timeout=(env.LIGHTNING_CLOUD_CONNECT_TIMEOUT, env.LIGHTNING_CLOUD_READ_TIMEOUT), + ) + except requests.RequestException as err: + raise _Unreachable(f"Couldn't reach {origin}. Check your connection and try again.") from err + + +def _token_endpoint(form: Dict[str, str]) -> Dict[str, Any]: + response = _post("/oidc/token", form) + try: + body = response.json() + except ValueError: + body = None + if not isinstance(body, dict): + body = {} + if response.status_code == 200: + return body + error = str(body.get("error", "")) + if response.status_code == 400 and error in ("invalid_grant", "access_denied"): + raise _Refused(error, str(body.get("error_description", ""))) + if response.status_code >= 500 or response.status_code == 429: + # The server is briefly unavailable or asking us to slow down. Unlike a + # refusal, that says nothing about the token we hold, so a refresh that + # hits it must not fail a request the token still covers. + raise _Unreachable( + f"Lightning answered the login request with {response.status_code}. Try again shortly." + ) + raise LoginError(f"Lightning answered the login request with {response.status_code} {error}.") + + +def _revoke(token: str) -> bool: + try: + return _post("/oidc/revoke", {"token": token}).status_code == 200 + except _Unreachable: + return False + + +def _login_from(body: Dict[str, Any], storage: str, url: str) -> Login: + try: + return Login( + url=url, + storage=storage, + access_expires_at=time.time() + float(body["expires_in"]), + refresh_expires_at=_epoch(body["refresh_expires_at"]), + grant=body.get("lightning_grant") or {}, + access_token=str(body["access_token"]), + refresh_token=str(body["refresh_token"]), + ) + except (KeyError, TypeError, ValueError): + raise LoginError( + "Lightning sent a login this SDK can't read. Upgrade the SDK and run `lightning login`." + ) from None + + +# Python 3.9's datetime.fromisoformat reads neither "Z" nor nanoseconds, and Go writes both. +_RFC3339 = re.compile(r"(\d{4}-\d\d-\d\dT\d\d:\d\d:\d\d)(\.\d+)?(Z|[+-]\d\d:\d\d)") + + +def _epoch(text: str) -> float: + match = _RFC3339.fullmatch(text) + if match is None: + raise ValueError("not an RFC 3339 time") + whole, fraction, zone = match.groups() + zone = "+00:00" if zone == "Z" else zone + return datetime.fromisoformat(whole + zone).timestamp() + float(fraction or 0) + + +def _say(message: str) -> None: + print(message, file=sys.stderr) diff --git a/python/lightning_sdk/lightning_cloud/credentials.py b/python/lightning_sdk/lightning_cloud/credentials.py new file mode 100644 index 000000000..7ac1090e8 --- /dev/null +++ b/python/lightning_sdk/lightning_cloud/credentials.py @@ -0,0 +1,117 @@ +# This file is vendored. +# Do not edit it directly; changes will be overwritten. +# Make changes in the internal upstream repository instead. + +import contextlib +import contextvars +import os +import pathlib +from typing import Iterator, Optional + +from lightning_sdk.lightning_cloud import cli_login, env +from lightning_sdk.lightning_cloud.login import Auth, Keys + +_MISSING_CREDENTIALS = ( + "We couldn't find any credentials linked to your account. " + "Please try logging in using the CLI command `lightning login`" +) +_LOGIN_FIRST = "No Lightning credentials are available. Run `lightning login` first." + +_browser_login_allowed = contextvars.ContextVar( + "lightning_browser_login_allowed", default=True +) + + +@contextlib.contextmanager +def browser_login(allowed: bool) -> Iterator[None]: + """Allow or forbid opening a browser to log in when no credential is found.""" + token = _browser_login_allowed.set(allowed) + try: + yield + finally: + _browser_login_allowed.reset(token) + + +def _present(value: Optional[str]) -> Optional[str]: + if value is None: + return None + text = value.strip() + if text == "": + return None + return text + + +class CredentialSource: + def __init__(self) -> None: + self._may_open_browser = True + self._login_token: Optional[str] = None + + def header(self) -> str: + token = self._token() + if token is None and self._may_open_browser: + if not _browser_login_allowed.get(): + raise ValueError(_LOGIN_FIRST) + self._may_open_browser = False + Auth()._run_server() + token = self._token() + if not token: + raise ValueError(_MISSING_CREDENTIALS) + return f"Bearer {token}" + + def refresh(self) -> bool: + """Refresh the login after the server refused its access token. + + Only the saved login can refresh, and only when its token was the one + last handed out. + """ + rejected = self._login_token + if rejected is None: + return False + fresh = self._saved_login_token(rejected) + return fresh is not None and fresh != rejected + + def _token(self) -> Optional[str]: + self._login_token = None + for read in ( + self._from_token_file, + self._from_env_token, + self._from_saved_login, + ): + token = read() + if token: + return token + return None + + def _from_token_file(self) -> Optional[str]: + path = os.environ.get(env.LIGHTNING_AUTH_TOKEN_FILE) + if not path: + return None + try: + raw = pathlib.Path(path).read_text() + except FileNotFoundError: + return None + return _present(raw) + + def _from_env_token(self) -> Optional[str]: + return _present(os.environ.get(Keys.AUTH_TOKEN.value)) + + def _from_saved_login(self) -> Optional[str]: + self._login_token = self._saved_login_token() + if self._login_token: + return self._login_token + auth = Auth() + if not auth.load(): + return None + token = _present(auth.auth_token) + if token: + return token + return _present(auth.api_key) + + def _saved_login_token(self, rejected: Optional[str] = None) -> Optional[str]: + try: + return cli_login._request_token(rejected) + except cli_login.LoginError: + # Every process sharing a refused login would otherwise open a + # browser. Saying "run `lightning login`" once is enough. + self._may_open_browser = False + raise diff --git a/python/lightning_sdk/lightning_cloud/env.py b/python/lightning_sdk/lightning_cloud/env.py index aa546c03c..b0b6bf487 100644 --- a/python/lightning_sdk/lightning_cloud/env.py +++ b/python/lightning_sdk/lightning_cloud/env.py @@ -3,9 +3,12 @@ # Make changes in the internal upstream repository instead. import os +from importlib.metadata import PackageNotFoundError, version from pathlib import Path from typing import Optional +from lightning_sdk.lightning_cloud.__version__ import __version__ + def truthy(val) -> bool: if val in [ @@ -21,6 +24,15 @@ def truthy(val) -> bool: return bool(val) +def _user_agent() -> str: + # A lightning-sdk agent is checked against the SDK's minimum version, so + # only the installed SDK may claim that name. + try: + return f"lightning-sdk/{version('lightning_sdk')}" + except PackageNotFoundError: + return f"lightning-cloud/{__version__}" + + def _float_env(name: str, default: Optional[float]) -> Optional[float]: raw = os.getenv(name) if not raw: @@ -40,6 +52,9 @@ def _float_env(name: str, default: Optional[float]) -> Optional[float]: SSL_CA_CERT = os.getenv("REQUESTS_CA_BUNDLE", default=os.getenv("SSL_CERT_FILE", default=None)) VERSION = os.getenv("VERSION", "0.0.1") +USER_AGENT = _user_agent() +# Name only. Read the path from the environment on each request. +LIGHTNING_AUTH_TOKEN_FILE = "LIGHTNING_AUTH_TOKEN_FILE" DEBUG = truthy(os.getenv("DEBUG")) CONTEXT = os.getenv("CONTEXT", "staging-3") LIGHTNING_SETTINGS_PATH = os.getenv( @@ -48,6 +63,9 @@ def _float_env(name: str, default: Optional[float]) -> Optional[float]: LIGHTNING_CREDENTIAL_PATH = os.getenv( 'LIGHTNING_CREDENTIAL_PATH', str(Path.home() / '.lightning' / 'credentials.json')) +LIGHTNING_LOGIN_PATH = os.getenv( + 'LIGHTNING_LOGIN_PATH', + str(Path.home() / '.lightning' / 'login.json')) DOT_IGNORE_FILENAME = ".lightningignore" diff --git a/python/lightning_sdk/lightning_cloud/login.py b/python/lightning_sdk/lightning_cloud/login.py index 3815a3270..2fea997bf 100644 --- a/python/lightning_sdk/lightning_cloud/login.py +++ b/python/lightning_sdk/lightning_cloud/login.py @@ -2,7 +2,6 @@ # Do not edit it directly; changes will be overwritten. # Make changes in the internal upstream repository instead. -import base64 from enum import Enum import json import logging @@ -10,9 +9,6 @@ import pathlib from dataclasses import dataclass from typing import Optional, Literal -from urllib.parse import urlencode - -import webbrowser import requests @@ -56,9 +52,6 @@ def __post_init__(self): if getattr(self, key.suffix) is None: setattr(self, key.suffix, os.environ.get(key.value, None)) - # used by authenticate method - self._with_env_var = bool(self.api_key or self.auth_token) - def load(self) -> bool: """Load credentials from disk and update properties with credentials. @@ -75,30 +68,33 @@ def load(self) -> bool: setattr(self, key.suffix, credentials.get(key.suffix, None)) return True - def save(self, - token: str = "", - user_id: str = "", - api_key: str = "", - auth_token: str = "", - # username: str = "", - ) -> None: - """save credentials to disk.""" + def save( + self, + user_id: Optional[str] = None, + api_key: Optional[str] = None, + auth_token: Optional[str] = None, + # username: str = "", + ) -> None: + """Replace the saved login with the given fields.""" + given = { + Keys.USER_ID.suffix: user_id, + Keys.API_KEY.suffix: api_key, + Keys.AUTH_TOKEN.suffix: auth_token, + } + stored = {key: value for key, value in given.items() if value} + self.secrets_file.parent.mkdir(exist_ok=True, parents=True) - with self.secrets_file.open("w") as f: - json.dump( - { - # f"{Keys.USERNAME.suffix}": username, - f"{Keys.USER_ID.suffix}": user_id, - f"{Keys.API_KEY.suffix}": api_key, - f"{Keys.AUTH_TOKEN.suffix}": auth_token, - }, - f, - ) + payload = json.dumps(stored).encode() + fd = os.open(self.secrets_file, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) + try: + os.write(fd, payload) + finally: + os.close(fd) + # os.open's mode applies only to a new file, so an existing one needs chmod. + os.chmod(self.secrets_file, 0o600) - # self.username = username - self.user_id = user_id - self.api_key = api_key - self.auth_token = auth_token + for key in Keys: + setattr(self, key.suffix, stored.get(key.suffix)) logger.debug("credentials saved successfully") @classmethod @@ -111,24 +107,25 @@ def clear(cls) -> None: logger.debug("credentials removed successfully") def get_auth_header(self, override: Optional[AuthOverride] = None) -> Optional[str]: - """Get authentication header with optional override. - - By default, uses the automatic priority selection (auth_token > api_key). - You can override this for specific cases where you need a different auth method. - + """Get a Bearer authentication header, with an optional override. + + By default, returns the credential every request would send: the file + named by ``LIGHTNING_AUTH_TOKEN_FILE``, then ``LIGHTNING_AUTH_TOKEN``, + then the saved login. It never opens a browser. + Parameters ---------- override : AuthOverride, optional - Override the default authentication method: - - "auth_token": Force use of JWT auth token (Bearer) - - "api_key": Force use of API key (Basic/Bearer) - - None: Use automatic selection (default) - + Use a token held by this instance instead: + - "auth_token": this instance's auth token + - "api_key": this instance's API key + - None: the credential every request would send (default) + Returns ------- Optional[str] The authorization header for the specified method. - + Raises ------ ValueError @@ -141,7 +138,7 @@ def get_auth_header(self, override: Optional[AuthOverride] = None) -> Optional[s "Please use token_login() method first." ) return f"Bearer {self.auth_token}" - + elif override == "api_key": if not self.api_key: raise ValueError( @@ -149,63 +146,40 @@ def get_auth_header(self, override: Optional[AuthOverride] = None) -> Optional[s "Please set LIGHTNING_API_KEY environment variable " "or use authenticate() method." ) - if self.user_id: - token = f"{self.user_id}:{self.api_key}" - return f"Basic {base64.b64encode(token.encode('ascii')).decode('ascii')}" # noqa E501 return f"Bearer {self.api_key}" - + elif override is None: - # Use the original automatic selection logic (default behavior) - if self.auth_token: - return f"Bearer {self.auth_token}" - elif self.api_key: - if self.user_id: - token = f"{self.user_id}:{self.api_key}" - return f"Basic {base64.b64encode(token.encode('ascii')).decode('ascii')}" # noqa E501 - return f"Bearer {self.api_key}" - else: - raise ValueError( - "No authentication credentials available. Please authenticate first using " - "token_login() or authenticate() methods." - ) - + from lightning_sdk.lightning_cloud.credentials import CredentialSource, browser_login + + with browser_login(False): + return CredentialSource().header() + else: raise ValueError(f"Invalid authentication override: {override}") @property def auth_header(self) -> Optional[str]: - """authentication header used by lightning-cloud client (automatic selection).""" + """The credential every request would send, without opening a browser.""" return self.get_auth_header() def _run_server(self) -> None: - """start a server to complete authentication.""" - AuthServer().login_with_browser(self) + """Log in through the browser, or by pasting a code when none opens.""" + from lightning_sdk.lightning_cloud.cli_login import login + + login() def authenticate(self) -> Optional[str]: """Performs end to end authentication flow. Returns ---------- - authorization header to use when authentication completes. + authorization header to use when authentication completes. With no + credential, this opens a browser to log in, unless browser login is + turned off. """ - if self._with_env_var: - logger.debug("successfully loaded credentials from env") - return self.auth_header + from lightning_sdk.lightning_cloud.credentials import CredentialSource - if not self.load(): - logger.debug( - "failed to load credentials, opening browser to get new.") - self._run_server() - return self.auth_header - - elif self.auth_token: - return self.auth_header - elif self.api_key: - return self.auth_header - - raise ValueError( - "We couldn't find any credentials linked to your account. Please try logging in using the CLI command `lightning login`" - ) + return CredentialSource().header() def token_login(self, token_key: str, save_token: bool = True) -> Optional[str]: """Performs token-based authentication. @@ -273,7 +247,7 @@ def token_login(self, token_key: str, save_token: bool = True) -> Optional[str]: else: logger.info("Successfully authenticated using auth token (not saved to disk).") - return self.auth_header + return f"Bearer {jwt_token}" def refresh_token(self, duration: int = 43200) -> Optional[str]: """Refreshes the current JWT token. @@ -343,7 +317,7 @@ def refresh_token(self, duration: int = 43200) -> Optional[str]: self.save(auth_token=new_jwt_token) logger.info("Successfully refreshed JWT token.") - return self.auth_header + return f"Bearer {new_jwt_token}" def create_api_client(self, override: Optional[AuthOverride] = None) -> 'ApiClient': """Create an API client with optional authentication override. @@ -378,77 +352,3 @@ def create_api_client(self, override: Optional[AuthOverride] = None) -> 'ApiClie client.set_default_header('Authorization', auth_header) return client - - -class AuthServer: - - def get_auth_url(self, port: int) -> str: - redirect_uri = f"http://localhost:{port}/login-complete" - params = urlencode(dict(redirectTo=redirect_uri)) - return f"{env.LIGHTNING_CLOUD_URL}/sign-in?{params}" - - def login_with_browser(self, auth: Auth) -> None: - # Deferred: the login callback server is only needed here, and these - # imports are too slow to pay on every CLI startup. - import uvicorn - from fastapi import FastAPI, Query, Request - from starlette.background import BackgroundTask - from starlette.responses import RedirectResponse - - from lightning_sdk.lightning_cloud.utils.network import find_free_network_port - - app = FastAPI() - port = find_free_network_port() - url = self.get_auth_url(port) - try: - # check if server is reachable or catch any network errors - requests.head(url) - except requests.ConnectionError as e: - raise requests.ConnectionError( - f"No internet connection available. Please connect to a stable internet connection \n{e}" # noqa E501 - ) - except requests.RequestException as e: - raise requests.RequestException( - f"An error occurred with the request. Please report this issue to Lightning Team \n{e}" # noqa E501 - ) - - logger.info(f"login started for lightning.ai, opening {url}") - ok = webbrowser.open(url) - if not ok: - # can't open a browser, authentication failed - deployment_id = os.environ.get("LIGHTNING_DEPLOYMENT_ID", None) - if deployment_id is not None and deployment_id != "": - raise RuntimeError("Failed to authenticate to Lightning. Ensure that you have selected 'Include SDK credentials' in the 'Environment' section of the deployment settings.") - raise RuntimeError("Failed to authenticate to Lightning. When running without access to a browser, 'LIGHTNING_API_KEY' should be exported.") - - @app.get("/login-complete") - async def save_token(request: Request, - token="", - key="", - user_id: str = Query("", alias="userID")): - if token: - auth.save(token=token, - # username=user_id, - user_id=user_id, - api_key=key) - logger.info("Authentication Successful") - else: - logger.warning( - "Authentication Failed. This is most likely because you're using an older version of the CLI. \n" # noqa E501 - "Please try to update the CLI or open an issue with this information \n" # noqa E501 - f"expected token in {request.query_params.items()}") - - # Include the credentials in the redirect so that UI will also be logged in - params = urlencode(dict(token=token, key=key, userID=user_id)) - - return RedirectResponse( - url=f"{env.LIGHTNING_CLOUD_URL}/me/apps?{params}", - # The response background task is being executed right after the server finished writing the response - background=BackgroundTask(stop_server)) - - def stop_server(): - server.should_exit = True - - server = uvicorn.Server( - config=uvicorn.Config(app, port=port, log_level="error")) - server.run() diff --git a/python/lightning_sdk/lightning_cloud/rest_client.py b/python/lightning_sdk/lightning_cloud/rest_client.py index 3193821c1..4c12c9fae 100644 --- a/python/lightning_sdk/lightning_cloud/rest_client.py +++ b/python/lightning_sdk/lightning_cloud/rest_client.py @@ -12,7 +12,7 @@ import urllib3 from lightning_sdk.lightning_cloud import env -from lightning_sdk.lightning_cloud.login import Auth +from lightning_sdk.lightning_cloud.credentials import CredentialSource from lightning_sdk.lightning_cloud.openapi import ( ApiClient, AssistantsServiceApi, @@ -61,7 +61,7 @@ def create_swagger_client(check_context: bool = True, with_auth: bool = True): If true, check if the context is set. It's only false for APIs that doesn't need the context information i.e login with_auth: bool - If true, it uses user's Authorization for the api client + If true, GridRestClient attaches Authorization on each request. """ if check_context and not env.CONTEXT: raise RuntimeError( @@ -80,9 +80,7 @@ def create_swagger_client(check_context: bool = True, with_auth: bool = True): env.LIGHTNING_CLOUD_READ_TIMEOUT, ) api_client = ApiClient(configuration) - if with_auth: - api_client.default_headers["Authorization"] = Auth().authenticate() - api_client.user_agent = f"Grid-CLI-{env.VERSION}" + api_client.user_agent = env.USER_AGENT return api_client @@ -119,8 +117,15 @@ class GridRestClient( ): def __init__(self, api_client: Optional[ApiClient] = None, with_auth: bool = True): - api_client = api_client if api_client else create_swagger_client(with_auth=with_auth) - api_client.request = request_auth_warning_wrapper(api_client.request) + api_client = ( + api_client if api_client else create_swagger_client(with_auth=with_auth) + ) + if with_auth: + api_client.request = credential_request_wrapper( + api_client.request, CredentialSource() + ) + else: + api_client.request = request_auth_warning_wrapper(api_client.request) super().__init__(api_client) @@ -202,7 +207,10 @@ def __init__(self, retry: bool = True, max_tries: Optional[int] = 10, with_auth: bool = True) -> None: - super().__init__(api_client=create_swagger_client(with_auth=with_auth)) + super().__init__( + api_client=create_swagger_client(with_auth=with_auth), + with_auth=with_auth, + ) if retry: for base_class in GridRestClient.__mro__: for name, attribute in base_class.__dict__.items(): @@ -261,6 +269,42 @@ def _without_none(values: Dict[str, Any]) -> Dict[str, Any]: return {key: value for key, value in values.items() if value is not None} +_LOGIN_MESSAGE = "Authentication failed. Please run `lightning login`." + + +def _has_authorization(headers: Optional[dict]) -> bool: + if not headers: + return False + return any(str(key).lower() == "authorization" for key in headers) + + +def credential_request_wrapper(func, source): + @functools.wraps(func) + def wrap(*args, **kwargs): + incoming = kwargs.get("headers") + supplied = _has_authorization(incoming) + headers = dict(incoming or {}) + if not supplied: + headers["Authorization"] = source.header() + kwargs["headers"] = headers + try: + return func(*args, **kwargs) + except ApiException as err: + if err.status != 401: + raise + if supplied or not source.refresh(): + raise ConnectionError(_LOGIN_MESSAGE) + headers["Authorization"] = source.header() + try: + return func(*args, **kwargs) + except ApiException as retry_err: + if retry_err.status == 401: + raise ConnectionError(_LOGIN_MESSAGE) + raise + + return wrap + + def request_auth_warning_wrapper(func): @functools.wraps(func) @@ -270,7 +314,7 @@ def wrap(*args, **kwargs): return response except ApiException as err: if err.status == 401: - raise ConnectionError("Authentication failed. Please run `lightning login`.") + raise ConnectionError(_LOGIN_MESSAGE) raise err return wrap diff --git a/python/pyproject.toml b/python/pyproject.toml index 787a955a6..ae1d707d0 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -21,6 +21,7 @@ dependencies = [ # alphabetically sorted "rich-click", # for rich CLI help formatting "docker", # for lit-container-registry "fastapi", # for lightning-cloud + "keyring", # for the CLI login "packaging", # for version checks "pyyaml", # for config handling "requests", # for downloading files via urls diff --git a/python/tests/api/test_agent_api.py b/python/tests/api/test_agent_api.py index 5d2046765..65df45a8d 100644 --- a/python/tests/api/test_agent_api.py +++ b/python/tests/api/test_agent_api.py @@ -3,7 +3,7 @@ from lightning_sdk.api.agents_api import AgentApi -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") def test_get_agent(_mock_auth, internal_agents_api_get_agent_mocker): agent_api = AgentApi() agent = agent_api.get_agent("ag-abc") @@ -11,20 +11,20 @@ def test_get_agent(_mock_auth, internal_agents_api_get_agent_mocker): assert agent.name == "ag-abc" -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") def test_delete_agent(_mock_auth, internal_agents_api_delete_agent_mocker): agent_api = AgentApi() agent_api.delete_agent("ag-abc", "ts-abc") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") def test_update_agent(_mock_auth, internal_agents_api_update_agent_mocker, internal_agents_api_get_agent_mocker): agent_api = AgentApi() new_agent = agent_api.update_agent("ag-abc", "ts-abc", name="test-sdk-1") assert new_agent.name == "test-sdk-1" -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") def test_update_endpoint( _mock_auth, internal_agents_api_update_agent_endpoint_mocker, internal_agents_api_get_agent_endpoint_mocker ): @@ -34,7 +34,7 @@ def test_update_endpoint( assert new_endpoint.openai.api_key == "test-sdk-1" -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") def test_get_agent_endpoint(_mock_auth, internal_agents_api_get_agent_endpoint_mocker): agent_api = AgentApi() endpoint = agent_api._get_agent_endpoint("ep-abc", "ts-abc") diff --git a/python/tests/api/test_billing_api.py b/python/tests/api/test_billing_api.py index 224db25cd..47e150c6b 100644 --- a/python/tests/api/test_billing_api.py +++ b/python/tests/api/test_billing_api.py @@ -292,7 +292,7 @@ def test_get_activity_filter_values_with_project_id(mock_client): "lightning_sdk.api.billing_api._authenticate_and_get_auth_headers", return_value={"Authorization": "Bearer test-token"}, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_session_activity_csv(mock_authenticate, mock_requests_get, tmp_path): mock_response = mock.Mock() mock_response.status_code = 200 @@ -320,7 +320,7 @@ def test_get_session_activity_csv(mock_authenticate, mock_requests_get, tmp_path "lightning_sdk.api.billing_api._authenticate_and_get_auth_headers", return_value={"Authorization": "Bearer test-token"}, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_session_activity_json_default_format(mock_authenticate, mock_requests_get, tmp_path): mock_response = mock.Mock() mock_response.status_code = 200 @@ -344,7 +344,7 @@ def test_get_session_activity_json_default_format(mock_authenticate, mock_reques "lightning_sdk.api.billing_api._authenticate_and_get_auth_headers", return_value={"Authorization": "Bearer test-token"}, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_resource_activity_csv(mock_authenticate, mock_requests_get, tmp_path): mock_response = mock.Mock() mock_response.status_code = 200 @@ -376,7 +376,7 @@ def test_get_resource_activity_csv(mock_authenticate, mock_requests_get, tmp_pat "lightning_sdk.api.billing_api._authenticate_and_get_auth_headers", return_value={"Authorization": "Bearer test-token"}, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_resource_activity_json_default_format(mock_authenticate, mock_requests_get, tmp_path): mock_response = mock.Mock() mock_response.status_code = 200 diff --git a/python/tests/api/test_data_connection_credentials.py b/python/tests/api/test_data_connection_credentials.py index 6875c2430..e1f62dd40 100644 --- a/python/tests/api/test_data_connection_credentials.py +++ b/python/tests/api/test_data_connection_credentials.py @@ -14,7 +14,7 @@ def _api_with_connections(*connections, credentials=None): """Build a TeamspaceApi whose generated client is mocked out.""" - with mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()): + with mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()): api = TeamspaceApi() client = mock.MagicMock() @@ -50,7 +50,7 @@ def test_unknown_connection_name_lists_the_ones_that_exist(): def test_no_connections_at_all_is_not_a_crash(): """An empty teamspace returns no list at all, rather than an empty one.""" - with mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()): + with mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()): api = TeamspaceApi() client = mock.MagicMock() client.data_connection_service_list_data_connections.return_value = V1ListDataConnectionsResponse( diff --git a/python/tests/api/test_filesystem_api.py b/python/tests/api/test_filesystem_api.py index 3ba886a15..1678a2281 100644 --- a/python/tests/api/test_filesystem_api.py +++ b/python/tests/api/test_filesystem_api.py @@ -9,7 +9,7 @@ @mock.patch("requests.get", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") @mock.patch( "lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers", return_value={"Authorization": "Bearer token"}, @@ -31,7 +31,7 @@ def test_download_file(_authenticate_mock, _mock_auth, mock_requests_get, tmpdir @mock.patch("lightning_sdk.api.filesystem_api.tqdm") @mock.patch("lightning_sdk.api.filesystem_api.ThreadPoolExecutor") @mock.patch("lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") def test_download_folder(_mock_auth, authenticate_mock, mock_executor, mock_tqdm, mock_collect, tmpdir): authenticate_mock.return_value = "test-token-123" filesystem_api = FilesystemApi() @@ -60,7 +60,7 @@ def test_download_folder(_mock_auth, authenticate_mock, mock_executor, mock_tqdm @mock.patch("requests.get", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") @mock.patch( "lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers", return_value={"Authorization": "Bearer token"}, @@ -86,7 +86,7 @@ def test_download_file_rejects_an_error_document(_authenticate_mock, _mock_auth, @mock.patch("requests.get", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") @mock.patch( "lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers", return_value={"Authorization": "Bearer token"}, @@ -110,7 +110,7 @@ def test_download_file_retries_a_dropped_connection(_authenticate_mock, _mock_au @mock.patch("requests.get", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") @mock.patch( "lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers", return_value={"Authorization": "Bearer token"}, @@ -139,7 +139,7 @@ def test_download_file_retries_a_truncated_body(_authenticate_mock, _mock_auth, @mock.patch("requests.get", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") @mock.patch( "lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers", return_value={"Authorization": "Bearer token"}, @@ -163,7 +163,7 @@ def test_download_file_does_not_retry_a_denial(_authenticate_mock, _mock_auth, m "lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers", return_value={"Authorization": "Bearer token"}, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") def test_download_folder_does_not_report_success_on_a_partial_download( _mock_auth, _authenticate_mock, inline_executor_cls, tmp_path ): @@ -189,7 +189,7 @@ def download(remote_path, *args, **kwargs): @mock.patch("requests.get", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") @mock.patch( "lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers", return_value={"Authorization": "Bearer token"}, @@ -208,7 +208,7 @@ def test_list_files_returns_tree(_authenticate_mock, _mock_auth, mock_requests_g @mock.patch("requests.get", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") @mock.patch( "lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers", return_value={"Authorization": "Bearer token"}, @@ -226,7 +226,7 @@ def test_list_files_returns_empty_when_tree_missing(_authenticate_mock, _mock_au @mock.patch("requests.get", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") @mock.patch( "lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers", return_value={"Authorization": "Bearer token"}, @@ -251,7 +251,7 @@ def test_list_files_passes_correct_url_and_params(_authenticate_mock, _mock_auth @mock.patch("requests.get", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") @mock.patch( "lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers", return_value={"Authorization": "Bearer token"}, @@ -270,7 +270,7 @@ def test_list_files_non_recursive_by_default(_authenticate_mock, _mock_auth, moc @mock.patch("requests.get", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") @mock.patch( "lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers", return_value={"Authorization": "Bearer token"}, @@ -286,7 +286,7 @@ def test_list_files_raises_on_non_200(_authenticate_mock, _mock_auth, mock_reque @mock.patch("requests.get", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") @mock.patch( "lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers", return_value={"Authorization": "Bearer token"}, diff --git a/python/tests/api/test_k8s_api.py b/python/tests/api/test_k8s_api.py index bdf0be085..4e54e31bd 100644 --- a/python/tests/api/test_k8s_api.py +++ b/python/tests/api/test_k8s_api.py @@ -18,7 +18,7 @@ ] -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") def test_get_billing_usage_with_empty_metrics(_mock_auth): k8s_api = K8sClusterApi("test") get_k8s_mock = mock.Mock(return_value=V1ListClusterMetricsResponse(cluster_metrics=[])) @@ -33,7 +33,7 @@ def test_get_billing_usage_with_empty_metrics(_mock_auth): get_k8s_mock.assert_called_once_with("test") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") def test_get_billing_usage_metrics_with_no_range(_mock_auth): k8s_api = K8sClusterApi("test") get_k8s_mock = mock.Mock( @@ -51,7 +51,7 @@ def test_get_billing_usage_metrics_with_no_range(_mock_auth): get_k8s_mock.assert_called_once_with("test") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") def test_gets_calculates_metrics_with_start_and_end_dates(_mock_auth): k8s_api = K8sClusterApi("user-abc") start_date = datetime(2025, 11, 19, 11, 30) @@ -68,7 +68,7 @@ def test_gets_calculates_metrics_with_start_and_end_dates(_mock_auth): assert result[1]["num_allocated_gpus"] == 4.0 -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") def test_gets_calculates_metrics_with_only_start_date(_mock_auth): k8s_api = K8sClusterApi("user-abc") date = datetime(2025, 11, 19, 12, 30) @@ -83,7 +83,7 @@ def test_gets_calculates_metrics_with_only_start_date(_mock_auth): get_k8s_mock.assert_called_once_with("user-abc", start_date=date) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") def test_gets_calculates_metrics_with_only_end_date(_mock_auth): k8s_api = K8sClusterApi("user-abc") date = datetime(2025, 11, 19, 12, 30) @@ -98,7 +98,7 @@ def test_gets_calculates_metrics_with_only_end_date(_mock_auth): get_k8s_mock.assert_called_once_with("user-abc", end_date=date) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") def test_gets_raises_error_with_nonexistant_cluster(_mock_auth): k8s_api = K8sClusterApi("fakey-abc") date = datetime(2025, 11, 19, 12, 30) @@ -111,7 +111,7 @@ def test_gets_raises_error_with_nonexistant_cluster(_mock_auth): get_k8s_mock.assert_called_once_with("fakey-abc", end=date) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") def test_get_billing_usage_metrics_print_to_stdout(_mock_auth): k8s_api = K8sClusterApi("test") get_k8s_mock = mock.Mock( @@ -127,7 +127,7 @@ def test_get_billing_usage_metrics_print_to_stdout(_mock_auth): get_k8s_mock.assert_called_once_with("test") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") def test_get_billing_usage_with_dates_get_folded_into_same_hour(_mock_auth): k8s_api = K8sClusterApi("test") test_cluster = [ @@ -163,7 +163,7 @@ def test_get_billing_usage_with_dates_get_folded_into_same_hour(_mock_auth): get_k8s_mock.assert_called_once_with("test") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") def test_get_billing_usage_metrics_with_max_gpus_safety_check(_mock_auth): k8s_api = K8sClusterApi("test") get_k8s_mock = mock.Mock( diff --git a/python/tests/api/test_mmt_api.py b/python/tests/api/test_mmt_api.py index f6c3605cf..35ec1f40a 100644 --- a/python/tests/api/test_mmt_api.py +++ b/python/tests/api/test_mmt_api.py @@ -17,7 +17,7 @@ from lightning_sdk.status import Status -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") def test_mmt_v2_submit_job(_mock_auth): job_api = MMTApiV2() @@ -113,7 +113,7 @@ def test_mmt_v2_submit_job(_mock_auth): create_job_mock.assert_called_once_with(project_id="ts-abc", body=body) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") def test_mmt_v2_submit_job_threads_placement_group_id(_mock_auth): job_api = MMTApiV2() create_job_mock = mock.MagicMock() @@ -143,7 +143,7 @@ def test_mmt_v2_submit_job_threads_placement_group_id(_mock_auth): assert body.spec.placement_group_id == "pg-1" -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") def test_mmt_v2_submit_job_threads_keep_minutes(_mock_auth): job_api = MMTApiV2() create_job_mock = mock.MagicMock() @@ -175,7 +175,7 @@ def test_mmt_v2_submit_job_threads_keep_minutes(_mock_auth): @pytest.mark.parametrize("timeout", [None, 0]) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") def test_mmt_v2_submit_job_omits_unset_keep_minutes(_mock_auth, timeout): job_api = MMTApiV2() create_job_mock = mock.MagicMock() @@ -206,7 +206,7 @@ def test_mmt_v2_submit_job_omits_unset_keep_minutes(_mock_auth, timeout): assert body.spec.keep_machine_after_stop is None -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") def test_mmt_v2_submit_job_rejects_negative_keep_minutes(_mock_auth): job_api = MMTApiV2() job_api._client.jobs_service_create_multi_machine_job = mock.MagicMock() @@ -233,7 +233,7 @@ def test_mmt_v2_submit_job_rejects_negative_keep_minutes(_mock_auth): ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") def test_mmt_v2_submit_job_threads_tags(_mock_auth): mmt_api = MMTApiV2() create_job_mock = mock.MagicMock() @@ -263,7 +263,7 @@ def test_mmt_v2_submit_job_threads_tags(_mock_auth): assert body.tags == ["prod"] -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") def test_get_mmt_by_name(_mock_auth): job_api = MMTApiV2() @@ -274,7 +274,7 @@ def test_get_mmt_by_name(_mock_auth): get_job_by_name_mock.assert_called_once_with(name="test-job", project_id="ts-abc") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") def test_get_mmt(_mock_auth): job_api = MMTApiV2() @@ -298,7 +298,7 @@ def test_get_mmt(_mock_auth): ("MultiMachineJob_STATE_DELETED", Status.Stopped), ], ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") def test_translate_state(_mock_auth, internal_state, expected_state): job_api = MMTApiV2() assert job_api._job_state_to_external(internal_state) == expected_state @@ -345,7 +345,7 @@ def test_machine_translate( ), ], ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") def test_mmt_stop(_mock_auth, job_states: List[str], total_calls_get_job: int, called_update_job: bool): job_api = MMTApiV2() @@ -384,7 +384,7 @@ def get_job_side_effect(*args, **kwargs): update_job_mock.assert_not_called() -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") def test_mmt_delete(_mock_auth): job_api = MMTApiV2() @@ -409,7 +409,7 @@ def test_get_studio_name_returns_none_when_cloudspace_missing(): ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") def test_mmt_v2_submit_job_auto_derives_fault_tolerance_from_max_run_attempts(_mock_auth): job_api = MMTApiV2() create_job_mock = mock.MagicMock() @@ -444,7 +444,7 @@ def test_mmt_v2_submit_job_auto_derives_fault_tolerance_from_max_run_attempts(_m assert body.fault_tolerance.strategy == "MULTI_MACHINE_JOB_FAULT_TOLERANCE_STRATEGY_RECREATE_ALL_NODES" -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") def test_mmt_v2_submit_job_omits_fault_tolerance_when_max_run_attempts_unset(_mock_auth): job_api = MMTApiV2() create_job_mock = mock.MagicMock() @@ -476,7 +476,7 @@ def test_mmt_v2_submit_job_omits_fault_tolerance_when_max_run_attempts_unset(_mo assert body.fault_tolerance is None -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") def test_mmt_v2_submit_job_omits_fault_tolerance_when_max_run_attempts_leq_one(_mock_auth): job_api = MMTApiV2() create_job_mock = mock.MagicMock() @@ -508,7 +508,7 @@ def test_mmt_v2_submit_job_omits_fault_tolerance_when_max_run_attempts_leq_one(_ assert body.fault_tolerance is None -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") def test_mmt_set_tags_sends_request(_mock_auth): from lightning_sdk.lightning_cloud.openapi import JobsServiceSetWorkloadTagsBody @@ -526,7 +526,7 @@ def test_mmt_set_tags_sends_request(_mock_auth): ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") def test_mmt_set_tags_sends_empty_list(_mock_auth): from lightning_sdk.lightning_cloud.openapi import JobsServiceSetWorkloadTagsBody diff --git a/python/tests/api/test_org_api.py b/python/tests/api/test_org_api.py index 94e0cd452..76767aff2 100644 --- a/python/tests/api/test_org_api.py +++ b/python/tests/api/test_org_api.py @@ -7,7 +7,7 @@ from lightning_sdk.lightning_cloud.openapi import V1Organization -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_org_api(internal_get_org_api_mocker): org_api = OrgApi() @@ -15,7 +15,7 @@ def test_org_api(internal_get_org_api_mocker): assert isinstance(org, V1Organization) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_org_api_valueerror(internal_get_org_api_mocker): org_api = OrgApi() diff --git a/python/tests/api/test_studio_api.py b/python/tests/api/test_studio_api.py index 80cbeb9b1..aea360e7f 100644 --- a/python/tests/api/test_studio_api.py +++ b/python/tests/api/test_studio_api.py @@ -51,7 +51,7 @@ class _DummyResponse: "lightning_sdk.lightning_cloud.openapi.api.cloud_space_service_api.CloudSpaceServiceApi.cloud_space_service_list_cloud_spaces", autospec=True, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_studio(mock_list_cloudspaces): def list_cloudspaces_side_effect(self, project_id, name): if name in ["st-abc", "st-def"]: @@ -69,7 +69,7 @@ def list_cloudspaces_side_effect(self, project_id, name): "lightning_sdk.lightning_cloud.openapi.api.cloud_space_service_api.CloudSpaceServiceApi.cloud_space_service_list_cloud_spaces", autospec=True, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_studio_error(mock_list_cloudspaces): def list_cloudspaces_side_effect(self, project_id, name): if name in ["st-abc", "st-def"]: @@ -95,7 +95,7 @@ def list_cloudspaces_side_effect(self, project_id, name): "lightning_sdk.lightning_cloud.openapi.api.cloud_space_service_api.CloudSpaceServiceApi.cloud_space_service_create_cloud_space", autospec=True, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_create_studio(mock_create_cloud_space, mock_create_lightning_run, _, cloud_account, sandbox, disable_secrets): def _create_cloudspace_side_effect(self, body, project_id, **kwargs): assert isinstance(body, CloudSpaceServiceCreateCloudSpaceBody) @@ -140,7 +140,7 @@ def _create_lightning_run_side_effect(self, body, project_id, cloudspace_id, **k "lightning_sdk.lightning_cloud.openapi.api.cloud_space_service_api.CloudSpaceServiceApi.cloud_space_service_get_cloud_space_instance_status", autospec=True, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_studio_status(mock_get_status): mock_get_status.return_value = V1GetCloudSpaceInstanceStatusResponse( requested=Externalv1CloudSpaceInstanceStatus( @@ -186,7 +186,7 @@ def test_get_studio_status(mock_get_status): "lightning_sdk.lightning_cloud.openapi.api.cloud_space_service_api.CloudSpaceServiceApi.cloud_space_service_get_cloud_space_instance_status", autospec=True, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_switch_studio_machine(mock_get_status, mock_update_config, mock_switch_instance, machine): mock_get_status.return_value = V1GetCloudSpaceInstanceStatusResponse( requested=Externalv1CloudSpaceInstanceStatus( @@ -217,7 +217,7 @@ def test_switch_studio_machine(mock_get_status, mock_update_config, mock_switch_ "lightning_sdk.lightning_cloud.openapi.api.cloud_space_service_api.CloudSpaceServiceApi.cloud_space_service_switch_cloud_space_instance", autospec=True, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_switch_machine_no_requested(switch_mock, update_mock, status_mock): return_vals = [ # First call: there is a requested machine (we just made the request) @@ -278,7 +278,7 @@ def side_effect(*args, **kwargs): "lightning_sdk.lightning_cloud.openapi.api.cloud_space_service_api.CloudSpaceServiceApi.cloud_space_service_start_cloud_space_instance", autospec=True, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_start_studio(mock_start_instance, mock_get_status): mock_start_instance.return_value = V1StartCloudSpaceInstanceResponse() mock_get_status.return_value = V1GetCloudSpaceInstanceStatusResponse( @@ -308,7 +308,7 @@ def test_start_studio(mock_start_instance, mock_get_status): "lightning_sdk.lightning_cloud.openapi.api.cloud_space_service_api.CloudSpaceServiceApi.cloud_space_service_stop_cloud_space_instance", autospec=True, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_stop_studio(mock_stop_instance, mock_get_status): mock_get_status.return_value = V1GetCloudSpaceInstanceStatusResponse( requested=Externalv1CloudSpaceInstanceStatus( @@ -328,7 +328,7 @@ def test_stop_studio(mock_stop_instance, mock_get_status): "lightning_sdk.lightning_cloud.openapi.api.cloud_space_service_api.CloudSpaceServiceApi.cloud_space_service_execute_command_in_cloud_space", autospec=True, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_run_command(mock_execute_command, mock_get_stream): mock_execute_command.return_value = V1ExecuteCloudSpaceCommandResponse( exit_code=0, output="Command Started Successfully", session_name="session-name" @@ -356,7 +356,7 @@ def test_run_command(mock_execute_command, mock_get_stream): "lightning_sdk.lightning_cloud.openapi.api.cloud_space_service_api.CloudSpaceServiceApi.cloud_space_service_delete_cloud_space", autospec=True, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_delete_studio(mock_delete_cloud_space): mock_delete_cloud_space.return_value = V1DeleteCloudSpaceResponse() @@ -406,7 +406,7 @@ def test_delete_studio(mock_delete_cloud_space): "lightning_sdk.lightning_cloud.openapi.api.cluster_service_api.ClusterServiceApi.cluster_service_list_project_clusters", autospec=True, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_machine( mock_list_project_clusters, mock_list_clusters, mock_list_accelerators, mock_get_config, name, expected_machine ): @@ -689,7 +689,7 @@ def _side_effect(self, project_id, id, **kwargs): "lightning_sdk.lightning_cloud.openapi.api.projects_service_api.ProjectsServiceApi.projects_service_get_project", autospec=True, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_duplicate_user( mock_get_project, mock_search_users, @@ -747,7 +747,7 @@ def test_duplicate_user( "lightning_sdk.lightning_cloud.openapi.api.projects_service_api.ProjectsServiceApi.projects_service_get_project", autospec=True, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_duplicate_org( mock_get_project, mock_get_org, @@ -813,7 +813,7 @@ def test_duplicate_org( "lightning_sdk.lightning_cloud.openapi.api.projects_service_api.ProjectsServiceApi.projects_service_get_project", autospec=True, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_duplicate_user_with_new_name( mock_get_project, mock_search_users, @@ -877,7 +877,7 @@ def test_duplicate_user_with_new_name( "lightning_sdk.lightning_cloud.openapi.api.projects_service_api.ProjectsServiceApi.projects_service_get_project", autospec=True, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_duplicate_org_with_new_name( mock_get_project, mock_get_org, @@ -928,7 +928,7 @@ def test_duplicate_org_with_new_name( "lightning_sdk.api.studio_api._authenticate_and_get_auth_headers", return_value={"Authorization": "Bearer test-token"}, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_tree(mock_authenticate, mock_requests_get): """Test get_tree retrieves directory structure from studio.""" mock_response = mock.MagicMock() @@ -1009,7 +1009,7 @@ def test_get_tree(mock_authenticate, mock_requests_get): "lightning_sdk.api.studio_api._authenticate_and_get_auth_headers", return_value={"Authorization": "Bearer test-token"}, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_path_info(mock_authenticate, mock_requests_get, path, tree_response, expected_result): mock_response = mock.MagicMock() mock_response.json.return_value = tree_response @@ -1108,7 +1108,7 @@ def test_get_path_info(mock_authenticate, mock_requests_get, path, tree_response "lightning_sdk.api.studio_api._authenticate_and_get_auth_headers", return_value={"Authorization": "Bearer test-token"}, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_list_files( mock_authenticate, mock_requests_get, @@ -1159,7 +1159,7 @@ def _blob_upload_create_response(path, upload_id, urls): @mock.patch("requests.put") @mock.patch("lightning_sdk.api.utils.tqdm") @mock.patch("lightning_sdk.api.utils._authenticate_and_get_auth_headers") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_upload_file( authenticate_mock, tqdm_mock, @@ -1237,7 +1237,7 @@ def test_upload_file( @mock.patch("requests.put") @mock.patch("lightning_sdk.api.utils.tqdm") @mock.patch("lightning_sdk.api.utils._authenticate_and_get_auth_headers") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_upload_file_leading_slash( authenticate_mock, tqdm_mock, @@ -1274,7 +1274,7 @@ def test_upload_file_leading_slash( @mock.patch( "lightning_sdk.api.studio_api._authenticate_and_get_auth_headers", return_value={"Authorization": "Bearer token"} ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_download_file(mock_authenticate, mock_requests_get, tmpdir): mock_response = mock.Mock() mock_response.status_code = 200 @@ -1293,7 +1293,7 @@ def test_download_file(mock_authenticate, mock_requests_get, tmpdir): @mock.patch( "lightning_sdk.api.studio_api._authenticate_and_get_auth_headers", return_value={"Authorization": "Bearer token"} ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_download_file_non_200_status(mock_authenticate, mock_requests_get, tmpdir): mock_response = mock.Mock() mock_response.status_code = 404 @@ -1315,7 +1315,7 @@ def test_download_file_non_200_status(mock_authenticate, mock_requests_get, tmpd @mock.patch("lightning_sdk.api.studio_api.tqdm") @mock.patch("lightning_sdk.api.studio_api.ThreadPoolExecutor") @mock.patch("lightning_sdk.api.studio_api._authenticate_and_get_auth_headers") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_download_folder(authenticate_mock, mock_executor, mock_tqdm, mock_collect, tmpdir): authenticate_mock.return_value = {"Authorization": "Bearer test-token-123"} @@ -1364,7 +1364,7 @@ def _error_document_response(status_code: int, content_length: Optional[int] = N @mock.patch("requests.get", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_download_single_studio_file_rejects_an_error_document(mock_requests_get, tmp_path): mock_requests_get.return_value = _error_document_response(503) @@ -1378,7 +1378,7 @@ def test_download_single_studio_file_rejects_an_error_document(mock_requests_get @mock.patch("requests.get", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_download_single_studio_file_rejects_an_error_document_under_200(mock_requests_get, tmp_path): # The observed failure: a 2 MiB object came back as a 278-byte body under a 200, so the # only thing that gives it away is the size the listing reported. @@ -1396,7 +1396,7 @@ def test_download_single_studio_file_rejects_an_error_document_under_200(mock_re @mock.patch( "lightning_sdk.api.studio_api._authenticate_and_get_auth_headers", return_value={"Authorization": "Bearer token"} ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_download_folder_does_not_report_success_on_a_partial_download( _mock_authenticate, inline_executor_cls, tmp_path ): @@ -1425,7 +1425,7 @@ def download(file_info, *args, **kwargs): "lightning_sdk.lightning_cloud.openapi.api.endpoint_service_api.EndpointServiceApi.endpoint_service_create_endpoint", autospec=True, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_add_ports(mock_create_endpoint): mock_create_endpoint.return_value = V1Endpoint( id="endpoint-id", @@ -1445,7 +1445,7 @@ def test_add_ports(mock_create_endpoint): "lightning_sdk.lightning_cloud.openapi.api.cloud_space_service_api.CloudSpaceServiceApi.cloud_space_service_update_cloud_space_sleep_config", autospec=True, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_update_autoshutdown_auto_sleep_on(mock_update_sleep_config): def response(self, id, project_id, body): if id == "st-abc": @@ -1482,7 +1482,7 @@ def response(self, id, project_id, body): "lightning_sdk.lightning_cloud.openapi.api.cloud_space_service_api.CloudSpaceServiceApi.cloud_space_service_update_cloud_space_sleep_config", autospec=True, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_update_autoshutdown_auto_sleep_off(mock_update_sleep_config): def response(self, id, project_id, body): if id == "st-abc": @@ -1519,7 +1519,7 @@ def response(self, id, project_id, body): "lightning_sdk.lightning_cloud.openapi.api.cloud_space_service_api.CloudSpaceServiceApi.cloud_space_service_update_cloud_space_sleep_config", autospec=True, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_update_autoshutdown_idle_shutdown(mock_update_sleep_config): def response(self, id, project_id, body): if id == "st-abc": @@ -1552,7 +1552,7 @@ def response(self, id, project_id, body): assert config.idle_shutdown_seconds == 900 -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_env(): from lightning_sdk.lightning_cloud.openapi import V1EnvVar @@ -1576,7 +1576,7 @@ def test_get_env(): assert result == expected -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_env_empty(): studio_api = StudioApi() @@ -1589,7 +1589,7 @@ def test_get_env_empty(): @mock.patch("lightning_sdk.api.studio_api.StudioApi._update_cloudspace") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_set_env_partial_true(mock_update_cloudspace): from lightning_sdk.lightning_cloud.openapi import V1EnvVar @@ -1629,7 +1629,7 @@ def test_set_env_partial_true(mock_update_cloudspace): @mock.patch("lightning_sdk.api.studio_api.StudioApi._update_cloudspace") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_set_env_partial_false(mock_update_cloudspace): from lightning_sdk.lightning_cloud.openapi import V1EnvVar @@ -1665,7 +1665,7 @@ def test_set_env_partial_false(mock_update_cloudspace): @mock.patch("lightning_sdk.api.studio_api.StudioApi._update_cloudspace") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_set_env_empty_new_env(mock_update_cloudspace): from lightning_sdk.lightning_cloud.openapi import V1EnvVar @@ -1770,7 +1770,7 @@ def test_set_env_empty_new_env(mock_update_cloudspace): "lightning_sdk.api.studio_api.StudioApi._get_machines_for_cloud_account", autospec=True, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_machine_has_capacity(mock_get_machines, machine, accelerators, expected_result): """Test machine_has_capacity method with various scenarios.""" mock_get_machines.return_value = accelerators @@ -1891,7 +1891,7 @@ def test_machine_has_capacity(mock_get_machines, machine, accelerators, expected "lightning_sdk.api.studio_api.StudioApi._get_machines_for_cloud_account", autospec=True, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_machine_is_supported(mock_get_machines, machine, accelerators, expected_result): """Test machine_is_supported method with various scenarios.""" mock_get_machines.return_value = accelerators @@ -1909,7 +1909,7 @@ def test_machine_is_supported(mock_get_machines, machine, accelerators, expected "lightning_sdk.lightning_cloud.openapi.api.endpoint_service_api.EndpointServiceApi.endpoint_service_list_endpoints", autospec=True, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_list_ports(mock_list_endpoints): """Test list_ports returns endpoints from the Studio.""" @@ -2004,7 +2004,7 @@ def test_list_ports(mock_list_endpoints): "lightning_sdk.lightning_cloud.openapi.api.endpoint_service_api.EndpointServiceApi.endpoint_service_list_endpoints", autospec=True, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_port_url(mock_list_endpoints, port, name, endpoints, expected_url, should_raise, error_match): """Test get_port_url with various port and name scenarios.""" mock_list_endpoints.return_value = V1ListEndpointsResponse(endpoints=endpoints) @@ -2070,7 +2070,7 @@ def test_get_port_url(mock_list_endpoints, port, name, endpoints, expected_url, return_value={"Authorization": "Bearer test-token"}, ) @mock.patch("lightning_sdk.api.studio_api.StudioApi.get_path_info", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_remove_file( mock_get_path_info, mock_authenticate, @@ -2152,7 +2152,7 @@ def test_remove_file( return_value={"Authorization": "Bearer test-token"}, ) @mock.patch("lightning_sdk.api.studio_api.StudioApi.get_path_info", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_remove_folder( mock_get_path_info, mock_authenticate, @@ -2197,7 +2197,7 @@ def test_remove_folder( "lightning_sdk.api.studio_api._authenticate_and_get_auth_headers", return_value={"Authorization": "Bearer token"}, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_list_files_follows_cursor_pages(_mock_authenticate, mock_requests_get): """A truncated listing is followed page by page until nextCursor comes back empty.""" pages = [ diff --git a/python/tests/api/test_studio_start_failure.py b/python/tests/api/test_studio_start_failure.py index 56fc9ceca..4cb3a0b50 100644 --- a/python/tests/api/test_studio_start_failure.py +++ b/python/tests/api/test_studio_start_failure.py @@ -28,7 +28,7 @@ def instance(instance_id="new", phase=State.PENDING, message=None, ready=False): @pytest.fixture() def start_api(mocker): - mocker.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) + mocker.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) api = StudioApi() api._client = mock.MagicMock() api._client.cloud_space_service_start_cloud_space_instance.return_value = V1StartCloudSpaceInstanceResponse() diff --git a/python/tests/api/test_teamspace_api.py b/python/tests/api/test_teamspace_api.py index 9bceda76c..2d96ee92e 100644 --- a/python/tests/api/test_teamspace_api.py +++ b/python/tests/api/test_teamspace_api.py @@ -29,7 +29,7 @@ def __init__(self, id: str) -> None: self.id = id -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_teamspace(internal_teamspace_api_mocker): teamspace_api = TeamspaceApi() @@ -37,7 +37,7 @@ def test_get_teamspace(internal_teamspace_api_mocker): assert isinstance(project, V1Project) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_teamspace_error(internal_teamspace_api_mocker): teamspace_api = TeamspaceApi() @@ -45,7 +45,7 @@ def test_get_teamspace_error(internal_teamspace_api_mocker): teamspace_api.get_teamspace("xyz", "org-def") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_list_teamspaces(internal_teamspace_api_list_mocker): teamspace_api = TeamspaceApi() @@ -64,7 +64,7 @@ def test_list_teamspaces(internal_teamspace_api_list_mocker): @mock.patch("lightning_sdk.api.teamspace_api.AuthApi") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_list_studios(mock_auth_api, internal_studio_api_list_mocker): mock_auth_api.return_value.whoami.return_value.user_id = "user-abc" teamspace_api = TeamspaceApi() @@ -76,7 +76,7 @@ def test_list_studios(mock_auth_api, internal_studio_api_list_mocker): assert isinstance(st, V1CloudSpace) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_list_clusters(internal_teamspace_api_cluster_list_mocker): teamspace_api = TeamspaceApi() @@ -87,7 +87,7 @@ def test_list_clusters(internal_teamspace_api_cluster_list_mocker): assert isinstance(cl, V1ProjectClusterBinding) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_create_agent(internal_teamspace_api_create_agent_mocker): teamspace_api = TeamspaceApi() agent = teamspace_api.create_agent( @@ -105,7 +105,7 @@ def test_create_agent(internal_teamspace_api_create_agent_mocker): @mock.patch("lightning_sdk.api.teamspace_api._ModelFileUploader") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_upload_model_file(uploader_mock): teamspace_api = TeamspaceApi() teamspace_api.upload_model_file( @@ -128,7 +128,7 @@ def test_upload_model_file(uploader_mock): uploader_mock().assert_called_with() # .__call__() -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_try_get_cluster_id(): # cluster set via env variable teamspace_api = TeamspaceApi() @@ -159,7 +159,7 @@ def test_try_get_cluster_id(): _ = teamspace_api._determine_cloud_account("teamspace-id") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_create_delete_model(): teamspace_api = TeamspaceApi() # create a content reused in following cases @@ -195,7 +195,7 @@ def test_create_delete_model(): teamspace_api._client.models_store_delete_model.assert_called_with(project_id="ts-abc", model_id="model-id") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_create_delete_model_with_experiment(): teamspace_api = TeamspaceApi() experiment = MyDummyExperiment(id="exp-abc") @@ -238,7 +238,7 @@ def test_create_delete_model_with_experiment(): teamspace_api._client.models_store_delete_model.assert_called_with(project_id="ts-abc", model_id="model-id") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_create_delete_model_version(): teamspace_api = TeamspaceApi() # mock the models_store_list_models and models_store_create_model_version for existing model @@ -281,7 +281,7 @@ def test_create_delete_model_version(): ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_create_delete_model_version_with_experiment(): teamspace_api = TeamspaceApi() # mock the models_store_list_models and models_store_create_model_version for existing model @@ -329,7 +329,7 @@ def test_create_delete_model_version_with_experiment(): @mock.patch("lightning_sdk.api.cloud_account_api.CloudAccountApi") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_list_machines(mock_cloud_account_api_class): mock_cloud_account_api = mock.Mock() mock_cloud_account_api_class.return_value = mock_cloud_account_api @@ -358,7 +358,7 @@ def test_list_machines(mock_cloud_account_api_class): @mock.patch("lightning_sdk.api.cloud_account_api.CloudAccountApi") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_list_machines_with_specific_machine(mock_cloud_account_api_class): mock_cloud_account_api = mock.Mock() mock_cloud_account_api_class.return_value = mock_cloud_account_api @@ -417,7 +417,7 @@ def test_list_machines_with_specific_machine(mock_cloud_account_api_class): ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_model_errors(internal_teamspace_api_mocker): teamspace_api = TeamspaceApi() @@ -444,7 +444,7 @@ def test_get_model_errors(internal_teamspace_api_mocker): "lightning_sdk.api.teamspace_api._authenticate_and_get_auth_headers", return_value={"Authorization": "Bearer test-token"}, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_tree(mock_authenticate, mock_requests_get): """Test get_tree retrieves directory structure from teamspace drive.""" mock_response = mock.MagicMock() @@ -525,7 +525,7 @@ def test_get_tree(mock_authenticate, mock_requests_get): "lightning_sdk.api.teamspace_api._authenticate_and_get_auth_headers", return_value={"Authorization": "Bearer test-token"}, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_path_info(mock_authenticate, mock_requests_get, path, tree_response, expected_result): mock_response = mock.MagicMock() mock_response.json.return_value = tree_response @@ -624,7 +624,7 @@ def test_get_path_info(mock_authenticate, mock_requests_get, path, tree_response "lightning_sdk.api.teamspace_api._authenticate_and_get_auth_headers", return_value={"Authorization": "Bearer test-token"}, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_list_files( mock_authenticate, mock_requests_get, @@ -684,7 +684,7 @@ def _make_upload_teamspace_api(): "lightning_sdk.api.utils._authenticate_and_get_auth_headers", new=mock.MagicMock(return_value={"Authorization": "Bearer test-token"}), ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_upload_file( tqdm_mock, requests_put_mock, @@ -761,7 +761,7 @@ def test_upload_file( "lightning_sdk.api.utils._authenticate_and_get_auth_headers", new=mock.MagicMock(return_value={"Authorization": "Bearer test-token"}), ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_upload_file_uploads_paths_pass_through( requests_put_mock, requests_post_mock, @@ -794,7 +794,7 @@ def test_upload_file_uploads_paths_pass_through( "lightning_sdk.api.utils._authenticate_and_get_auth_headers", new=mock.MagicMock(return_value={"Authorization": "Bearer test-token"}), ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_upload_file_with_headers( tqdm_mock, requests_put_mock, @@ -841,7 +841,7 @@ def test_upload_file_with_headers( "lightning_sdk.api.utils._authenticate_and_get_auth_headers", new=mock.MagicMock(return_value={"Authorization": "Bearer test-token"}), ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_upload_file_without_headers( requests_put_mock, requests_post_mock, @@ -885,7 +885,7 @@ def _stub_response(status_code: int, body: bytes) -> mock.Mock: @mock.patch( "lightning_sdk.api.teamspace_api._authenticate_and_get_auth_headers", return_value={"Authorization": "Bearer token"} ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_download_file(mock_authenticate, mock_requests_get, tmpdir): mock_requests_get.return_value = _stub_response(200, b"data" * 256) @@ -901,7 +901,7 @@ def test_download_file(mock_authenticate, mock_requests_get, tmpdir): @mock.patch( "lightning_sdk.api.teamspace_api._authenticate_and_get_auth_headers", return_value={"Authorization": "Bearer token"} ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_download_file_rejects_an_error_document(_mock_authenticate, mock_requests_get, tmpdir): mock_requests_get.return_value = _stub_response(403, _S3_ERROR_BODY) @@ -915,7 +915,7 @@ def test_download_file_rejects_an_error_document(_mock_authenticate, mock_reques @mock.patch("requests.get", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_download_single_file_rejects_an_error_document_under_200(mock_requests_get, tmp_path): # A 2 MiB object that comes back as a 278-byte body under a 200 is only detectable # against the size the listing reported. @@ -933,7 +933,7 @@ def test_download_single_file_rejects_an_error_document_under_200(mock_requests_ @mock.patch( "lightning_sdk.api.teamspace_api._authenticate_and_get_auth_headers", return_value={"Authorization": "Bearer token"} ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_download_folder_does_not_report_success_on_a_partial_download( _mock_authenticate, inline_executor_cls, tmp_path ): @@ -959,7 +959,7 @@ def download(file_info, *args, **kwargs): @mock.patch("lightning_sdk.api.teamspace_api.tqdm") @mock.patch("lightning_sdk.api.teamspace_api.ThreadPoolExecutor") @mock.patch("lightning_sdk.api.teamspace_api._authenticate_and_get_auth_headers") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_download_folder(authenticate_mock, mock_executor, mock_tqdm, mock_collect, tmpdir): authenticate_mock.return_value = {"Authorization": "Bearer test-token-123"} @@ -990,7 +990,7 @@ def test_download_folder(authenticate_mock, mock_executor, mock_tqdm, mock_colle assert mock_tqdm.call_args.kwargs["total"] == 3000 # 1000 + 2000 -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_secrets(): teamspace_api = TeamspaceApi() @@ -1007,7 +1007,7 @@ def test_get_secrets(): assert secrets["DATABASE_URL"] == "***REDACTED***" -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_set_secret_create_new(): teamspace_api = TeamspaceApi() @@ -1026,7 +1026,7 @@ def test_set_secret_create_new(): @pytest.mark.parametrize("secret_type", [SecretType.HF_TOKEN, "hf_token"]) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_set_secret_create_new_with_type(secret_type): teamspace_api = TeamspaceApi() @@ -1038,7 +1038,7 @@ def test_set_secret_create_new_with_type(secret_type): mock_create.assert_called_once_with("ts-abc", "HF_TOKEN", "hf_xxx", secret_type=V1SecretType.HF_TOKEN) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_set_secret_invalid_type_raises(): teamspace_api = TeamspaceApi() @@ -1048,7 +1048,7 @@ def test_set_secret_invalid_type_raises(): teamspace_api.set_secret("ts-abc", "NAME", "value", secret_type="not_a_type") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_set_secret_update_existing(): teamspace_api = TeamspaceApi() @@ -1065,7 +1065,7 @@ def test_set_secret_update_existing(): mock_update.assert_called_once_with("ts-abc", "secret-1", "new_secret_value") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_set_generic_secret_does_not_update_same_name_typed_secret(): teamspace_api = TeamspaceApi() existing_secrets = [ @@ -1081,7 +1081,7 @@ def test_set_generic_secret_does_not_update_same_name_typed_secret(): mock_create.assert_called_once_with("ts-abc", "API_KEY", "new_secret_value", secret_type=V1SecretType.UNSPECIFIED) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_delete_secret_resolves_generic_secret_id(): teamspace_api = TeamspaceApi() existing_secrets = [ @@ -1097,7 +1097,7 @@ def test_delete_secret_resolves_generic_secret_id(): mock_delete.assert_called_once_with(project_id="ts-abc", id="generic-secret") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_delete_secret_rejects_missing_generic_secret(): teamspace_api = TeamspaceApi() existing_secrets = [ @@ -1113,7 +1113,7 @@ def test_delete_secret_rejects_missing_generic_secret(): @mock.patch("lightning_sdk.api.utils.LightningClient") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_secrets_api_call(mock_client): mock_client().secret_service_list_secrets.return_value.secrets = [ V1Secret(id="secret-1", name="API_KEY"), @@ -1128,7 +1128,7 @@ def test_get_secrets_api_call(mock_client): @mock.patch("lightning_sdk.api.utils.LightningClient") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_create_secret_api_call(mock_client): teamspace_api = TeamspaceApi() @@ -1143,7 +1143,7 @@ def test_create_secret_api_call(mock_client): @mock.patch("lightning_sdk.api.utils.LightningClient") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_create_secret_api_call_with_type(mock_client): teamspace_api = TeamspaceApi() @@ -1154,7 +1154,7 @@ def test_create_secret_api_call_with_type(mock_client): @mock.patch("lightning_sdk.api.utils.LightningClient") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_update_secret_api_call(mock_client): teamspace_api = TeamspaceApi() @@ -1189,7 +1189,7 @@ def test_update_secret_api_call(mock_client): ("INVALID$SECRET", False), # contains dollar sign ], ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_verify_secret_name(secret_name, expected): teamspace_api = TeamspaceApi() result = teamspace_api.verify_secret_name(secret_name) @@ -1201,7 +1201,7 @@ def test_verify_secret_name(secret_name, expected): "lightning_sdk.api.teamspace_api._authenticate_and_get_auth_headers", return_value={"Authorization": "Bearer token"}, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_list_files_follows_cursor_pages(_mock_authenticate, mock_requests_get): """A truncated listing is followed page by page until nextCursor comes back empty.""" pages = [ @@ -1224,7 +1224,7 @@ def test_list_files_follows_cursor_pages(_mock_authenticate, mock_requests_get): assert params[1]["cursor"] == "cursor-1" -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_create_model_version_keeps_metadata(): """Metadata must reach every version, not just the first. @@ -1257,7 +1257,7 @@ def _tags_client(*tags: V1WorkloadTag) -> mock.MagicMock: return client -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_list_jobs_filters_on_tag_ids(): teamspace_api = TeamspaceApi() teamspace_api._client = mock.MagicMock() @@ -1269,7 +1269,7 @@ def test_list_jobs_filters_on_tag_ids(): ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_list_jobs_omits_the_tag_filter_when_unset(): teamspace_api = TeamspaceApi() teamspace_api._client = mock.MagicMock() @@ -1279,7 +1279,7 @@ def test_list_jobs_omits_the_tag_filter_when_unset(): teamspace_api._client.jobs_service_list_jobs.assert_called_once_with(project_id="ts-abc", standalone=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_list_mmts_filters_on_tag_ids_client_side(): tagged = mock.MagicMock(tags=[V1WorkloadTag(id="tag-1", name="prod")]) also_tagged = mock.MagicMock(tags=[V1WorkloadTag(id="tag-2", name="staging")]) @@ -1295,7 +1295,7 @@ def test_list_mmts_filters_on_tag_ids_client_side(): assert teamspace_api.list_mmts("ts-abc") == [tagged, also_tagged, untagged] -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) @pytest.mark.parametrize("name", ["team a", "TEAM A", " team a ", "Team\tA"]) def test_resolve_tag_ids_matches_names_the_way_the_platform_stores_them(name): teamspace_api = TeamspaceApi() @@ -1304,7 +1304,7 @@ def test_resolve_tag_ids_matches_names_the_way_the_platform_stores_them(name): assert teamspace_api.resolve_tag_ids("ts-abc", [name]) == ["tag-1"] -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_resolve_tag_ids_keeps_the_order_the_names_were_given(): teamspace_api = TeamspaceApi() teamspace_api._client = _tags_client( @@ -1316,7 +1316,7 @@ def test_resolve_tag_ids_keeps_the_order_the_names_were_given(): teamspace_api._client.jobs_service_list_workload_tags.assert_called_once_with(project_id="ts-abc") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_resolve_tag_ids_rejects_an_unknown_tag(): teamspace_api = TeamspaceApi() teamspace_api._client = _tags_client( @@ -1328,7 +1328,7 @@ def test_resolve_tag_ids_rejects_an_unknown_tag(): teamspace_api.resolve_tag_ids("ts-abc", ["nightly"]) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_resolve_tag_ids_says_when_the_teamspace_has_no_tags(): teamspace_api = TeamspaceApi() teamspace_api._client = _tags_client() @@ -1337,7 +1337,7 @@ def test_resolve_tag_ids_says_when_the_teamspace_has_no_tags(): teamspace_api.resolve_tag_ids("ts-abc", ["prod"]) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_list_workload_tags_tolerates_a_teamspace_with_no_tags(): teamspace_api = TeamspaceApi() teamspace_api._client = mock.MagicMock() diff --git a/python/tests/api/test_user_api.py b/python/tests/api/test_user_api.py index 3fb744112..a9774175f 100644 --- a/python/tests/api/test_user_api.py +++ b/python/tests/api/test_user_api.py @@ -11,7 +11,7 @@ ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_user_api(internal_user_api_mocker, monkeypatch): user_api = UserApi() @@ -19,7 +19,7 @@ def test_user_api(internal_user_api_mocker, monkeypatch): assert isinstance(user, V1GetUserResponse) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_user_api_valueerror(internal_user_api_mocker, monkeypatch): monkeypatch.setenv("LIGHTNING_USERNAME", "other-dummy") user_api = UserApi() @@ -33,7 +33,7 @@ def test_user_api_valueerror(internal_user_api_mocker, monkeypatch): autospec=True, return_value=V1GetUserResponse(features=V1UserFeatures(plugin_sweeps=True)), ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_user_api_get_feature_flags(mocker): user_api = UserApi() feature_flags = user_api._get_feature_flags() @@ -44,7 +44,7 @@ def test_user_api_get_feature_flags(mocker): assert feature_flags.plugin_sweeps -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_secrets(): user_api = UserApi() @@ -61,7 +61,7 @@ def test_get_secrets(): assert secrets["DATABASE_URL"] == "***REDACTED***" -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_set_secret_create_new(): user_api = UserApi() @@ -77,7 +77,7 @@ def test_set_secret_create_new(): mock_create.assert_called_once_with("NEW_SECRET", "secret_value") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_set_secret_update_existing(): user_api = UserApi() @@ -94,7 +94,7 @@ def test_set_secret_update_existing(): mock_update.assert_called_once_with("secret-1", "new_secret_value") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_set_secret_does_not_update_same_name_typed_secret(): user_api = UserApi() existing_secrets = [ @@ -110,7 +110,7 @@ def test_set_secret_does_not_update_same_name_typed_secret(): mock_create.assert_called_once_with("API_KEY", "new_secret_value") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_delete_secret_resolves_generic_secret_id(): user_api = UserApi() existing_secrets = [ @@ -126,7 +126,7 @@ def test_delete_secret_resolves_generic_secret_id(): mock_delete.assert_called_once_with(id="generic-secret") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_delete_secret_rejects_missing_generic_secret(): user_api = UserApi() existing_secrets = [ @@ -142,7 +142,7 @@ def test_delete_secret_rejects_missing_generic_secret(): @mock.patch("lightning_sdk.api.utils.LightningClient") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_secrets_api_call(mock_client): mock_client().secret_service_list_user_secrets.return_value.secrets = [ V1Secret(id="secret-1", name="API_KEY"), @@ -157,7 +157,7 @@ def test_get_secrets_api_call(mock_client): @mock.patch("lightning_sdk.api.utils.LightningClient") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_create_secret_api_call(mock_client): user_api = UserApi() @@ -170,7 +170,7 @@ def test_create_secret_api_call(mock_client): @mock.patch("lightning_sdk.api.utils.LightningClient") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_update_secret_api_call(mock_client): user_api = UserApi() @@ -204,7 +204,7 @@ def test_update_secret_api_call(mock_client): ("INVALID$SECRET", False), # contains dollar sign ], ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_verify_secret_name(secret_name, expected): user_api = UserApi() result = user_api.verify_secret_name(secret_name) @@ -212,7 +212,7 @@ def test_verify_secret_name(secret_name, expected): @mock.patch("lightning_sdk.api.utils.LightningClient") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_create_teamspace(mock_client): user_api = UserApi() diff --git a/python/tests/cli/test_entrypoint.py b/python/tests/cli/test_entrypoint.py index b2195a33c..260a90f69 100644 --- a/python/tests/cli/test_entrypoint.py +++ b/python/tests/cli/test_entrypoint.py @@ -9,8 +9,8 @@ from lightning_sdk.cli.entrypoint import login, main_cli from lightning_sdk.cli.sandbox.commands import run_sandbox_command -from lightning_sdk.cli.utils.auth import OIDC_LOGIN_ENV from lightning_sdk.cli.vm.ssh import ssh_vm +from lightning_sdk.lightning_cloud import cli_login from lightning_sdk.lightning_cloud.login import Auth from tests.cli.help import assert_help_contains, command_text, mock_command_logging, run_cli @@ -20,10 +20,9 @@ def _clear_auth(monkeypatch, tmp_path) -> None: monkeypatch.setattr(Auth, "secrets_file", tmp_path / "credentials.json") - monkeypatch.delenv("LIGHTNING_USER_ID", raising=False) - monkeypatch.delenv("LIGHTNING_API_KEY", raising=False) - monkeypatch.delenv("LIGHTNING_AUTH_TOKEN", raising=False) - monkeypatch.delenv(OIDC_LOGIN_ENV, raising=False) + monkeypatch.setenv("LIGHTNING_AUTH_TOKEN_FILE", str(tmp_path / "missing-token")) + for key in ("LIGHTNING_USER_ID", "LIGHTNING_API_KEY", "LIGHTNING_AUTH_TOKEN"): + monkeypatch.delenv(key, raising=False) @pytest.mark.parametrize( @@ -47,110 +46,104 @@ def test_resource_commands_never_start_browser_authentication(monkeypatch, tmp_p assert "Run `lightning login`" in str(result.exception) -def test_explicit_login_can_start_browser_authentication(monkeypatch, tmp_path) -> None: +def test_explicit_login_starts_the_login_flow(monkeypatch, tmp_path) -> None: _clear_auth(monkeypatch, tmp_path) - browser_auth_started = False - - def _complete_browser_auth(auth: Auth) -> None: - nonlocal browser_auth_started - browser_auth_started = True - auth.api_key = "api-key" - - monkeypatch.setattr(Auth, "_run_server", _complete_browser_auth) + login = MagicMock() + monkeypatch.setattr("lightning_sdk.cli.entrypoint.cli_login.login", login) + monkeypatch.setattr("lightning_sdk.cli.entrypoint._get_authed_user", lambda: MagicMock(name="Test User")) result = CliRunner().invoke(main_cli, ["login"]) assert result.exit_code == 0, result.output - assert browser_auth_started + login.assert_called_once_with(no_browser=False, insecure_storage=False, orgs=None, project=None) -def test_login_reports_success_and_failure(monkeypatch, tmp_path) -> None: +def test_login_passes_every_option_to_the_login_flow(monkeypatch, tmp_path) -> None: _clear_auth(monkeypatch, tmp_path) - user = MagicMock() - user.name = "Test User" - monkeypatch.setattr("lightning_sdk.cli.entrypoint._get_authed_user", lambda: user) + login = MagicMock() + monkeypatch.setattr("lightning_sdk.cli.entrypoint.cli_login.login", login) + monkeypatch.setattr("lightning_sdk.cli.entrypoint._get_authed_user", lambda: MagicMock(name="Test User")) - def _save_token(auth: Auth) -> None: - auth.save(token="sess_user-1.secret", user_id="user-1") - - monkeypatch.setattr(Auth, "_run_server", _save_token) - ok = CliRunner().invoke(main_cli, ["login"]) - assert ok.exit_code == 0, ok.output - assert 'Successfully logged in as "Test User".' in ok.output + result = CliRunner().invoke( + main_cli, + ["login", "--no-browser", "--org", "acme, beta", "--project", "team", "--insecure-storage"], + ) - # the browser callback finished without saving a token - _clear_auth(monkeypatch, tmp_path / "second") - monkeypatch.setattr(Auth, "_run_server", lambda _auth: None) - failed = CliRunner().invoke(main_cli, ["login"]) - assert failed.exit_code == 1 - assert "Login failed" in str(failed.exception) + assert result.exit_code == 0, result.output + login.assert_called_once_with( + no_browser=True, + insecure_storage=True, + orgs=["acme", "beta"], + project="team", + ) -def test_login_no_browser_starts_a_fresh_paste_login_instead_of_the_browser(monkeypatch, tmp_path) -> None: +def test_login_reports_success(monkeypatch, tmp_path) -> None: _clear_auth(monkeypatch, tmp_path) - monkeypatch.setenv(OIDC_LOGIN_ENV, "1") - Auth().save(auth_token="sess_expired") user = MagicMock() user.name = "Test User" monkeypatch.setattr("lightning_sdk.cli.entrypoint._get_authed_user", lambda: user) - monkeypatch.setattr( - Auth, - "_run_server", - lambda _self: (_ for _ in ()).throw(AssertionError("browser auth reached")), - ) - - def _approve(auth: Auth) -> None: - assert not Auth.secrets_file.exists() - auth.save(auth_token="sess_new") + monkeypatch.setattr("lightning_sdk.cli.entrypoint.cli_login.login", MagicMock()) - monkeypatch.setattr("lightning_sdk.cli.entrypoint.login_with_pasted_code", _approve) - - result = CliRunner().invoke(main_cli, ["login", "--no-browser"]) + result = CliRunner().invoke(main_cli, ["login"]) assert result.exit_code == 0, result.output assert 'Successfully logged in as "Test User".' in result.output - auth = Auth() - assert auth.load() - assert auth.auth_token == "sess_new" -def test_login_no_browser_is_refused_without_the_opt_in(monkeypatch, tmp_path) -> None: +def test_login_reports_a_refused_login(monkeypatch, tmp_path) -> None: _clear_auth(monkeypatch, tmp_path) - Auth().save(auth_token="sess_kept") monkeypatch.setattr( - "lightning_sdk.cli.entrypoint.login_with_pasted_code", - lambda _auth: (_ for _ in ()).throw(AssertionError("paste login reached")), + "lightning_sdk.cli.entrypoint.cli_login.login", + MagicMock(side_effect=cli_login.LoginError("No OS keychain is available")), ) - result = CliRunner().invoke(main_cli, ["login", "--no-browser"]) + result = CliRunner().invoke(main_cli, ["login"]) - assert result.exit_code == 2 - assert "not available yet" in result.output - auth = Auth() - assert auth.load() - assert auth.auth_token == "sess_kept" + assert result.exit_code == 1 + assert "No OS keychain is available" in str(result.exception) def test_logout_reports_success_and_failure(monkeypatch, tmp_path) -> None: _clear_auth(monkeypatch, tmp_path) + logout = MagicMock() + monkeypatch.setattr("lightning_sdk.cli.entrypoint.cli_login.logout", logout) + ok = CliRunner().invoke(main_cli, ["logout"]) assert ok.exit_code == 0, ok.output assert "Successfully logged out." in ok.output + logout.assert_called_once_with() - monkeypatch.setattr("lightning_sdk.cli.entrypoint.logout_session", MagicMock(side_effect=OSError("denied"))) + monkeypatch.setattr( + "lightning_sdk.cli.entrypoint.cli_login.logout", + MagicMock(side_effect=cli_login.LoginError("nobody home")), + ) failed = CliRunner().invoke(main_cli, ["logout"]) assert failed.exit_code == 1 - assert "Logout failed: denied" in str(failed.exception) + assert "Logout failed: nobody home" in str(failed.exception) + + +def test_logout_revokes_a_legacy_session(monkeypatch, tmp_path) -> None: + _clear_auth(monkeypatch, tmp_path) + Auth().save(auth_token="sess_old") + monkeypatch.setattr("lightning_sdk.cli.entrypoint.cli_login.logout", MagicMock()) + + with patch("lightning_sdk.cli.entrypoint.requests.post") as post: + result = CliRunner().invoke(main_cli, ["logout"]) + + assert result.exit_code == 0, result.output + post.assert_called_once() + assert post.call_args.kwargs["headers"] == {"Authorization": "Bearer sess_old"} def test_cli_browser_auth_policy_is_reset_after_command(monkeypatch, tmp_path) -> None: _clear_auth(monkeypatch, tmp_path) browser_auth_starts = 0 - def _complete_browser_auth(auth: Auth) -> None: + def _complete_browser_auth(_auth: Auth) -> None: nonlocal browser_auth_starts browser_auth_starts += 1 - auth.api_key = "api-key" + monkeypatch.setenv("LIGHTNING_AUTH_TOKEN", "token") monkeypatch.setattr(Auth, "_run_server", _complete_browser_auth) @@ -159,7 +152,7 @@ def _complete_browser_auth(auth: Auth) -> None: assert result.exit_code == 1 assert "Run `lightning login`" in str(result.exception) assert browser_auth_starts == 0 - assert Auth().authenticate().startswith("Bearer ") + assert Auth().authenticate() == "Bearer token" assert browser_auth_starts == 1 @@ -317,7 +310,10 @@ def test_login_help(): result_text = assert_help_contains("lightning login --help", "Usage: lightning login [OPTIONS]") assert "Sign in to Lightning AI." in result_text - assert "--no-browser" not in result_text + assert "--no-browser" in result_text + assert "--org" in result_text + assert "--project" in result_text + assert "--insecure-storage" in result_text @mock_command_logging @@ -327,40 +323,19 @@ def test_logout_help(): assert "Sign out of the current session." in result_text -def test_login_replaces_existing_session(monkeypatch, tmp_path): - _clear_auth(monkeypatch, tmp_path) - Auth().save(auth_token="sess_expired") - - def _complete_browser_auth(auth: Auth) -> None: - assert not Auth.secrets_file.exists() - auth.save(auth_token="sess_new") - - monkeypatch.setattr(Auth, "_run_server", _complete_browser_auth) - - result = CliRunner().invoke(main_cli, ["login"]) - - assert result.exit_code == 0, result.output - auth = Auth() - assert auth.load() - assert auth.auth_token == "sess_new" - - @pytest.mark.parametrize("env_var", ["LIGHTNING_AUTH_TOKEN", "LIGHTNING_API_KEY"]) def test_login_not_required_with_credentials_in_env(monkeypatch, tmp_path, env_var): _clear_auth(monkeypatch, tmp_path) monkeypatch.setenv(env_var, "from-env") - Auth().save(auth_token="sess_saved") - - def _fail_browser_auth(auth: Auth) -> None: - raise AssertionError("browser authentication should not start") - - monkeypatch.setattr(Auth, "_run_server", _fail_browser_auth) + monkeypatch.setattr( + "lightning_sdk.cli.entrypoint.cli_login.login", + MagicMock(side_effect=AssertionError("login should not start")), + ) result = CliRunner().invoke(main_cli, ["login"]) assert result.exit_code == 0, result.output assert "credentials are set in the environment" in result.output - assert Auth.secrets_file.exists() @patch.dict(os.environ, {"LIGHTNING_INTERACTIVE": "true", "LIGHTNING_CLOUD_SPACE_ID": "cloud-space-id"}, clear=True) diff --git a/python/tests/cli/utils/test_auth.py b/python/tests/cli/utils/test_auth.py index ec22912c8..360b847e2 100644 --- a/python/tests/cli/utils/test_auth.py +++ b/python/tests/cli/utils/test_auth.py @@ -1,99 +1,43 @@ -import base64 -import hashlib -import json from unittest.mock import MagicMock, PropertyMock, patch -from urllib.parse import parse_qs, urlparse import pytest import rich_click as click from lightning_sdk.cli.legacy.deploy._auth import _AuthLitServe, _AuthMode -from lightning_sdk.cli.utils.auth import ( - OIDC_LOGIN_ENV, - browser_authentication, - login_with_pasted_code, - logout_session, - require_auth_header, -) +from lightning_sdk.cli.utils.auth import browser_authentication, require_auth_header +from lightning_sdk.lightning_cloud.credentials import CredentialSource from lightning_sdk.lightning_cloud.login import Auth -def _response(status_code: int, body=None) -> MagicMock: - response = MagicMock(status_code=status_code, ok=status_code < 400) - response.json.return_value = body - return response - - -def _login_env(monkeypatch, tmp_path): - """Point credentials at tmp_path and return the secrets file.""" - secrets_file = tmp_path / "credentials.json" - monkeypatch.setattr(Auth, "secrets_file", secrets_file) - for key in ( - "LIGHTNING_USER_ID", - "LIGHTNING_API_KEY", - "LIGHTNING_AUTH_TOKEN", - "LIGHTNING_DEPLOYMENT_ID", - OIDC_LOGIN_ENV, - ): +def _no_credentials(monkeypatch, tmp_path) -> None: + for key in ("LIGHTNING_USER_ID", "LIGHTNING_API_KEY", "LIGHTNING_AUTH_TOKEN"): monkeypatch.delenv(key, raising=False) - return secrets_file - - -def _challenge(verifier: str) -> str: - return base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest()).rstrip(b"=").decode() - - -def _capture_browser_login(monkeypatch, tmp_path, token_response) -> dict: - """Run the loopback login with a fake server and browser; return what was captured.""" - captured: dict = {"secrets_file": _login_env(monkeypatch, tmp_path)} - monkeypatch.setenv(OIDC_LOGIN_ENV, "1") - - class FakeServer: - def __init__(self, config) -> None: - captured["config"] = config - captured["route"] = next(r for r in config.app.routes if getattr(r, "path", None) == "/callback") - - def run(self) -> None: - captured["run"](captured["route"].endpoint) - - post = MagicMock(side_effect=token_response) - monkeypatch.setattr("lightning_sdk.cli.utils.auth._generate_state", lambda: "fixed-state") - monkeypatch.setattr("lightning_sdk.cli.utils.auth.requests.head", MagicMock()) - monkeypatch.setattr("lightning_sdk.cli.utils.auth.requests.post", post) - monkeypatch.setattr("webbrowser.open", lambda url: captured.setdefault("auth_url", url) and True) - monkeypatch.setattr("uvicorn.Server", FakeServer) - captured["post"] = post - return captured - - -def test_require_auth_header_uses_environment_credentials() -> None: - """An environment header must be used without loading or authenticating.""" - auth = MagicMock(auth_header="Basic encoded") - - with patch("lightning_sdk.cli.utils.auth.Auth", return_value=auth): - assert require_auth_header() == "Basic encoded" - - auth.load.assert_not_called() - auth.authenticate.assert_not_called() + monkeypatch.setenv("LIGHTNING_AUTH_TOKEN_FILE", str(tmp_path / "missing-token")) + monkeypatch.setattr(Auth, "secrets_file", tmp_path / "credentials.json") + monkeypatch.setattr( + Auth, + "_run_server", + lambda self: (_ for _ in ()).throw(AssertionError("browser authentication reached")), + ) -def test_require_auth_header_loads_saved_credentials() -> None: - """Saved credentials must work without starting browser authentication.""" - auth = MagicMock() - type(auth).auth_header = PropertyMock(side_effect=[ValueError("missing"), "Bearer token"]) - auth.load.return_value = True +def test_require_auth_header_uses_the_credential_source() -> None: + """An available header must be returned without loading or authenticating.""" + auth = MagicMock(auth_header="Bearer token") with patch("lightning_sdk.cli.utils.auth.Auth", return_value=auth): assert require_auth_header() == "Bearer token" + auth.load.assert_not_called() auth.authenticate.assert_not_called() def test_require_auth_header_instructs_login_without_opening_browser() -> None: """Missing credentials must stop ordinary commands and direct users to login.""" auth = MagicMock() - type(auth).auth_header = PropertyMock(side_effect=ValueError("missing")) - auth.load.return_value = False + type(auth).auth_header = PropertyMock( + side_effect=ValueError("No Lightning credentials are available. Run `lightning login` first.") + ) with patch("lightning_sdk.cli.utils.auth.Auth", return_value=auth), pytest.raises( click.UsageError, match="lightning login" @@ -103,338 +47,36 @@ def test_require_auth_header_instructs_login_without_opening_browser() -> None: auth.authenticate.assert_not_called() -def test_browser_authentication_blocks_subclass_callback_server() -> None: - """The guard must also block SDK auth subclasses with custom callback servers.""" +def test_browser_authentication_blocks_the_source_from_opening_a_browser(monkeypatch, tmp_path) -> None: + _no_credentials(monkeypatch, tmp_path) - class CustomAuth(Auth): - browser_started = False + with browser_authentication(False), pytest.raises(ValueError, match="Run `lightning login` first"): + CredentialSource().header() - def load(self) -> bool: - return False - def _run_server(self) -> None: - self.browser_started = True +def test_browser_authentication_allows_the_source_to_open_a_browser(monkeypatch, tmp_path) -> None: + _no_credentials(monkeypatch, tmp_path) + started = [] + monkeypatch.setattr(Auth, "_run_server", lambda self: started.append(True)) - auth = CustomAuth(api_key="", auth_token="") + with browser_authentication(True), pytest.raises(ValueError, match="couldn't find any credentials"): + CredentialSource().header() - with browser_authentication(False), pytest.raises(ValueError, match="lightning login"): - auth.authenticate() + assert started == [True] - assert not auth.browser_started - -def test_litserve_auth_uses_saved_credentials_without_starting_callback_server() -> None: - """Deployment authentication must not start callback auth when credentials exist.""" +def test_litserve_auth_uses_the_credential_source(monkeypatch) -> None: + """Deployment authentication must not start callback auth when a credential exists.""" + monkeypatch.setenv("LIGHTNING_AUTH_TOKEN", "token") auth = _AuthLitServe(_AuthMode.DEPLOY) - auth.api_key = "key" - with patch.object( - auth, - "_run_server", - side_effect=AssertionError("callback server must not start"), - ): - assert auth.authenticate().startswith("Bearer ") + with patch.object(auth, "_run_server", side_effect=AssertionError("callback server must not start")): + assert auth.authenticate() == "Bearer token" -def test_litserve_auth_instructs_login_when_callback_auth_would_be_required() -> None: +def test_litserve_auth_instructs_login_when_callback_auth_would_be_required(monkeypatch, tmp_path) -> None: """Deployment callback auth must tell the user to run the explicit login command.""" - with pytest.raises(click.UsageError, match="lightning login"): - _AuthLitServe(_AuthMode.DEPLOY)._run_server() - - -def test_auth_save_stores_session_token_as_auth_token(monkeypatch, tmp_path) -> None: - """Browser login callback passes the session token as `token`; it must be saved as `auth_token`.""" - secrets_file = tmp_path / "credentials.json" - monkeypatch.setattr(Auth, "secrets_file", secrets_file) - monkeypatch.delenv("LIGHTNING_USER_ID", raising=False) - monkeypatch.delenv("LIGHTNING_API_KEY", raising=False) - monkeypatch.delenv("LIGHTNING_AUTH_TOKEN", raising=False) - - auth = Auth() - auth.save(token="sess_user-1.secret", user_id="user-1", api_key="") - - assert auth.auth_token == "sess_user-1.secret" - assert auth.api_key == "" - assert auth.auth_header == "Bearer sess_user-1.secret" - assert json.loads(secrets_file.read_text()) == { - "user_id": "user-1", - "api_key": "", - "auth_token": "sess_user-1.secret", - } - assert (secrets_file.stat().st_mode & 0o777) == 0o600 - - -def test_logout_session_revokes_server_session_and_clears_credentials(monkeypatch, tmp_path) -> None: - """`lightning logout` revokes the saved session token on the server and removes local credentials.""" - secrets_file = tmp_path / "credentials.json" - monkeypatch.setattr(Auth, "secrets_file", secrets_file) - monkeypatch.delenv("LIGHTNING_USER_ID", raising=False) - monkeypatch.delenv("LIGHTNING_API_KEY", raising=False) - monkeypatch.delenv("LIGHTNING_AUTH_TOKEN", raising=False) - - auth = Auth() - auth.save(token="sess_user-1.secret", user_id="user-1") - - with patch("lightning_sdk.cli.utils.auth.requests.post") as post_mock: - logout_session(Auth()) - - post_mock.assert_called_once() - assert post_mock.call_args.kwargs["headers"] == {"Authorization": "Bearer sess_user-1.secret"} - assert not secrets_file.exists() - - -def test_login_with_browser_exchanges_the_code_with_pkce_on_loopback(monkeypatch, tmp_path) -> None: - """The callback trades the code and the verifier matching the sent challenge for a session.""" - from lightning_sdk.lightning_cloud.login import AuthServer - - captured = _capture_browser_login( - monkeypatch, tmp_path, [_response(200, {"access_token": "sess_user-1.secret", "token_type": "Bearer"})] - ) - captured["run"] = lambda endpoint: captured.setdefault("location", endpoint(code="the-code", state="fixed-state")) - - auth = Auth() - AuthServer().login_with_browser(auth) - - url = urlparse(captured["auth_url"]) - query = {k: v[0] for k, v in parse_qs(url.query).items()} - assert url.path == "/oidc/authorize" - assert query["client_id"] == "lightning-cli" - assert query["response_type"] == "code" - assert query["scope"] == "openid" - assert query["state"] == "fixed-state" - assert query["code_challenge_method"] == "S256" - assert query["redirect_uri"] == f"http://127.0.0.1:{captured['config'].port}/callback" - assert captured["config"].host == "127.0.0.1" - - data = captured["post"].call_args.kwargs["data"] - assert captured["post"].call_args.args[0].endswith("/oidc/token") - assert data["grant_type"] == "authorization_code" - assert data["code"] == "the-code" - assert data["redirect_uri"] == query["redirect_uri"] - assert _challenge(data["code_verifier"]) == query["code_challenge"] - - assert captured["location"].headers["location"].endswith("/cli-login-successful") - assert auth.auth_token == "sess_user-1.secret" - assert json.loads(captured["secrets_file"].read_text())["auth_token"] == "sess_user-1.secret" - assert (captured["secrets_file"].stat().st_mode & 0o777) == 0o600 - - -@pytest.mark.parametrize( - "params", - [ - {"code": "the-code", "state": ""}, - {"code": "the-code", "state": "forged-state"}, - {"code": "", "state": "fixed-state", "error": "access_denied"}, - {"code": "", "state": "fixed-state"}, - ], - ids=["no state", "forged state", "denied", "no code"], -) -def test_login_with_browser_refuses_a_bad_callback_without_exchanging(monkeypatch, tmp_path, params) -> None: - from lightning_sdk.lightning_cloud.login import AuthServer - - captured = _capture_browser_login(monkeypatch, tmp_path, []) - captured["run"] = lambda endpoint: captured.setdefault("location", endpoint(**params)) - - auth = Auth() - AuthServer().login_with_browser(auth) - - assert captured["location"].headers["location"].endswith("/cli-login-failed") - captured["post"].assert_not_called() - assert not captured["secrets_file"].exists() - assert not auth.auth_token - - -def test_login_with_browser_reports_a_refused_exchange(monkeypatch, tmp_path) -> None: - from lightning_sdk.lightning_cloud.login import AuthServer - - captured = _capture_browser_login(monkeypatch, tmp_path, [_response(400, {"error": "invalid_grant"})]) - captured["run"] = lambda endpoint: captured.setdefault("location", endpoint(code="used", state="fixed-state")) - - AuthServer().login_with_browser(Auth()) - - assert captured["location"].headers["location"].endswith("/cli-login-failed") - assert not captured["secrets_file"].exists() - - -def test_login_with_browser_without_a_browser_falls_back_to_paste(monkeypatch, tmp_path) -> None: - from lightning_sdk.lightning_cloud.login import AuthServer + _no_credentials(monkeypatch, tmp_path) - secrets_file = _login_env(monkeypatch, tmp_path) - monkeypatch.setenv(OIDC_LOGIN_ENV, "1") - monkeypatch.setattr("lightning_sdk.cli.utils.auth.requests.head", MagicMock()) - monkeypatch.setattr("webbrowser.open", lambda url: False) - monkeypatch.setattr("uvicorn.Server", lambda config: (_ for _ in ()).throw(AssertionError("server started"))) - monkeypatch.setattr( - "lightning_sdk.cli.utils.auth.login_with_pasted_code", lambda auth: auth.save(auth_token="sess_x") - ) - - AuthServer().login_with_browser(Auth()) - - assert json.loads(secrets_file.read_text())["auth_token"] == "sess_x" - - -def test_login_with_browser_in_a_deployment_without_a_browser_raises(monkeypatch, tmp_path) -> None: - from lightning_sdk.lightning_cloud.login import AuthServer - - _login_env(monkeypatch, tmp_path) - monkeypatch.setenv("LIGHTNING_DEPLOYMENT_ID", "dep") - monkeypatch.setattr("lightning_sdk.cli.utils.auth.requests.head", MagicMock()) - monkeypatch.setattr("webbrowser.open", lambda url: False) - - with pytest.raises(RuntimeError, match="Include SDK credentials"): - AuthServer().login_with_browser(Auth()) - - -def test_login_with_browser_defaults_to_the_session_token_login(monkeypatch, tmp_path) -> None: - """Without the opt-in, the sign-in page hands a session token to `/login-complete`, and the callback - redirects to `/cli-login-successful` without forwarding it in the URL.""" - import asyncio - - from lightning_sdk.lightning_cloud.login import AuthServer - - _login_env(monkeypatch, tmp_path) - - captured: dict = {} - - class FakeServer: - def __init__(self, config) -> None: - captured["app"] = config.app - - def run(self) -> None: - route = next(r for r in captured["app"].routes if getattr(r, "path", None) == "/login-complete") - resp = asyncio.run( - route.endpoint(token="sess_user-1.secret", key="", state="fixed-state", user_id="user-1") - ) - captured["location"] = resp.headers["location"] - - monkeypatch.setattr("lightning_sdk.cli.utils.auth._generate_state", lambda: "fixed-state") - monkeypatch.setattr("lightning_sdk.cli.utils.auth.requests.head", MagicMock()) - monkeypatch.setattr("webbrowser.open", lambda url: captured.setdefault("auth_url", url) or True) - monkeypatch.setattr("uvicorn.Server", FakeServer) - - auth = Auth() - AuthServer().login_with_browser(auth) - - assert "/sign-in?" in captured["auth_url"] - assert "state=fixed-state" in captured["auth_url"] or "state%3Dfixed-state" in captured["auth_url"] - assert captured["location"].endswith("/cli-login-successful") - assert "token=" not in captured["location"] - assert auth.auth_token == "sess_user-1.secret" - - -def test_session_token_login_rejects_callback_without_matching_state(monkeypatch, tmp_path) -> None: - """Omitting or forging `state` must not store the callback token or key.""" - import asyncio - - from lightning_sdk.lightning_cloud.login import AuthServer - - secrets_file = _login_env(monkeypatch, tmp_path) - - captured: dict = {} - - class FakeServer: - def __init__(self, config) -> None: - captured["app"] = config.app - - def run(self) -> None: - captured["route"] = next(r for r in captured["app"].routes if getattr(r, "path", None) == "/login-complete") - - monkeypatch.setattr("lightning_sdk.cli.utils.auth._generate_state", lambda: "fixed-state") - monkeypatch.setattr("lightning_sdk.cli.utils.auth.requests.head", MagicMock()) - monkeypatch.setattr("webbrowser.open", lambda url: True) - monkeypatch.setattr("uvicorn.Server", FakeServer) - - auth = Auth() - AuthServer().login_with_browser(auth) - - for state in ("", "forged-state"): - resp = asyncio.run( - captured["route"].endpoint( - token="sess_attacker.secret", key="attacker-key", state=state, user_id="attacker" - ) - ) - assert resp.headers["location"].endswith("/cli-login-failed") - - assert not secrets_file.exists() - assert not auth.auth_token - assert not auth.api_key - - -def _paste_login(monkeypatch, tmp_path, codes: list, responses: list) -> tuple: - secrets_file = _login_env(monkeypatch, tmp_path) - post = MagicMock(side_effect=responses) - monkeypatch.setattr("lightning_sdk.cli.utils.auth.requests.post", post) - monkeypatch.setattr("lightning_sdk.cli.utils.auth.click.prompt", MagicMock(side_effect=codes)) - return secrets_file, post - - -def test_login_with_pasted_code_strips_a_trailing_slash_from_the_origin(monkeypatch, tmp_path, capsys) -> None: - monkeypatch.setattr("lightning_sdk.lightning_cloud.env.LIGHTNING_CLOUD_URL", "http://localhost:9800/") - _paste_login(monkeypatch, tmp_path, ["the-code"], [_response(200, {"access_token": "sess_user-1.secret"})]) - - login_with_pasted_code(Auth()) - - url = next(word for word in capsys.readouterr().out.split() if "/oidc/authorize" in word) - query = {k: v[0] for k, v in parse_qs(urlparse(url).query).items()} - assert url.startswith("http://localhost:9800/oidc/authorize?") - assert query["redirect_uri"] == "http://localhost:9800/cli-login" - assert query["scope"] == "openid" - - -def test_login_with_pasted_code_prints_the_authorize_url_and_saves_the_session(monkeypatch, tmp_path, capsys) -> None: - secrets_file, post = _paste_login( - monkeypatch, tmp_path, [" the-code \n"], [_response(200, {"access_token": "sess_user-1.secret"})] - ) - - login_with_pasted_code(Auth()) - - output = capsys.readouterr().out - url = next(word for word in output.split() if "/oidc/authorize" in word) - query = {k: v[0] for k, v in parse_qs(urlparse(url).query).items()} - assert query["client_id"] == "lightning-cli" - assert query["scope"] == "openid" - assert query["redirect_uri"].endswith("/cli-login") - assert "//cli-login" not in query["redirect_uri"] - assert query["code_challenge_method"] == "S256" - assert "never give it to anyone" in output - - data = post.call_args.kwargs["data"] - assert data["code"] == "the-code" - assert data["redirect_uri"] == query["redirect_uri"] - assert _challenge(data["code_verifier"]) == query["code_challenge"] - assert json.loads(secrets_file.read_text())["auth_token"] == "sess_user-1.secret" - - -def test_login_with_pasted_code_prompts_again_on_invalid_grant(monkeypatch, tmp_path) -> None: - secrets_file, post = _paste_login( - monkeypatch, - tmp_path, - ["typo", "the-code"], - [_response(400, {"error": "invalid_grant"}), _response(200, {"access_token": "sess_user-1.secret"})], - ) - - login_with_pasted_code(Auth()) - - assert [c.kwargs["data"]["code"] for c in post.call_args_list] == ["typo", "the-code"] - assert json.loads(secrets_file.read_text())["auth_token"] == "sess_user-1.secret" - - -def test_login_with_pasted_code_gives_up_after_three_bad_codes(monkeypatch, tmp_path) -> None: - refused = [_response(400, {"error": "invalid_grant"})] * 3 - secrets_file, post = _paste_login(monkeypatch, tmp_path, ["a", "b", "c"], refused) - - with pytest.raises(RuntimeError, match="not accepted"): - login_with_pasted_code(Auth()) - - assert post.call_count == 3 - assert not secrets_file.exists() - - -def test_login_with_pasted_code_stops_on_any_other_error(monkeypatch, tmp_path) -> None: - secrets_file, post = _paste_login(monkeypatch, tmp_path, ["a"], [_response(404, {})]) - - with pytest.raises(RuntimeError, match="HTTP 404"): - login_with_pasted_code(Auth()) - - assert post.call_count == 1 - assert not secrets_file.exists() + with browser_authentication(False), pytest.raises(click.UsageError, match="lightning login"): + _AuthLitServe(_AuthMode.DEPLOY)._run_server() diff --git a/python/tests/conftest.py b/python/tests/conftest.py index b603c137b..2e3d183b6 100644 --- a/python/tests/conftest.py +++ b/python/tests/conftest.py @@ -86,6 +86,7 @@ def _clear_lightning_env_vars(monkeypatch): if name.startswith("LIGHTNING_"): monkeypatch.delenv(name, raising=False) monkeypatch.setenv("LIGHTNING_API_KEY", "dummy-api-key") + monkeypatch.setenv("LIGHTNING_AUTH_TOKEN", "dummy-auth-token") @pytest.fixture(autouse=True) @@ -1779,7 +1780,7 @@ def slurm_jobs_user_service_create_user_slurm_job_side_effect(self, project_id, @pytest.fixture() def mocker_auth(mocker): mocker.patch( - "lightning_sdk.lightning_cloud.rest_client.Auth", + "lightning_sdk.lightning_cloud.credentials.Auth", autospec=True, ) diff --git a/python/tests/core/studio/test_studio_cloud_account.py b/python/tests/core/studio/test_studio_cloud_account.py index bb4cd379b..13c13333a 100644 --- a/python/tests/core/studio/test_studio_cloud_account.py +++ b/python/tests/core/studio/test_studio_cloud_account.py @@ -63,7 +63,7 @@ def _list_cloudspaces_side_effect(*args, **kwargs): "lightning_sdk.lightning_cloud.openapi.api.cloud_space_service_api.CloudSpaceServiceApi.cloud_space_service_switch_cloud_space_instance", autospec=True, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_studio_switch_cloud_account( mock_switch_cloudspace_instance, mock_update_cloudspace_instance_config, @@ -168,7 +168,7 @@ def test_studio_switch_cloud_account( "lightning_sdk.lightning_cloud.openapi.api.cloud_space_service_api.CloudSpaceServiceApi.cloud_space_service_switch_cloud_space_instance", autospec=True, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_studio_switch_cloud_account_not_global( mock_switch_cloudspace_instance, mock_update_cloudspace_instance_config, diff --git a/python/tests/core/studio/test_studio_commands.py b/python/tests/core/studio/test_studio_commands.py index 52f653fb4..e1fa07483 100644 --- a/python/tests/core/studio/test_studio_commands.py +++ b/python/tests/core/studio/test_studio_commands.py @@ -61,7 +61,7 @@ class _DummyResponse: ) @mock.patch("lightning_sdk.api.org_api.OrgApi.get_org", autospec=True) @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.get_teamspace", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_run_command( mock_get_teamspace, mock_get_org, @@ -174,7 +174,7 @@ def _create_lightning_run_side_effect(*args, **kwargs): ) @mock.patch("lightning_sdk.api.org_api.OrgApi.get_org", autospec=True) @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.get_teamspace", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_run_command_error( mock_get_teamspace, mock_get_org, @@ -283,7 +283,7 @@ def _create_lightning_run_side_effect(*args, **kwargs): ) @mock.patch("lightning_sdk.api.org_api.OrgApi.get_org", autospec=True) @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.get_teamspace", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_run_command_exit_code( mock_get_teamspace, mock_get_org, @@ -397,7 +397,7 @@ def _create_lightning_run_side_effect(*args, **kwargs): ) @mock.patch("lightning_sdk.api.org_api.OrgApi.get_org", autospec=True) @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.get_teamspace", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_run_command_and_detach( mock_get_teamspace, mock_get_org, diff --git a/python/tests/core/studio/test_studio_duplicate.py b/python/tests/core/studio/test_studio_duplicate.py index 12057ab26..bcea92633 100644 --- a/python/tests/core/studio/test_studio_duplicate.py +++ b/python/tests/core/studio/test_studio_duplicate.py @@ -38,7 +38,7 @@ def _list_cloudspaces_side_effect(*args, **kwargs): ) @mock.patch("lightning_sdk.api.org_api.OrgApi.get_org", autospec=True) @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.get_teamspace", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_studio_duplicate_machine( mock_get_teamspace, mock_get_org, @@ -130,7 +130,7 @@ def _create_lightning_run_side_effect(*args, **kwargs): ) @mock.patch("lightning_sdk.api.org_api.OrgApi.get_org", autospec=True) @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.get_teamspace", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_studio_duplicate_with_new_name( mock_get_teamspace, mock_get_org, diff --git a/python/tests/core/studio/test_studio_environment.py b/python/tests/core/studio/test_studio_environment.py index 59d2dfc97..41c8bff14 100644 --- a/python/tests/core/studio/test_studio_environment.py +++ b/python/tests/core/studio/test_studio_environment.py @@ -16,7 +16,7 @@ def _list_cloudspaces_side_effect(*args, **kwargs): return _list_cloudspaces_side_effect -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_studio_env_property(): """Test that Studio.env property calls StudioApi.get_env correctly.""" from lightning_sdk.lightning_cloud.openapi import V1EnvVar @@ -44,7 +44,7 @@ def test_studio_env_property(): studio._studio_api.get_env.assert_called_once_with(mock_studio) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_studio_set_env_partial_true(): """Test that Studio.set_env calls StudioApi.set_env correctly with partial=True.""" mock_studio = V1CloudSpace(id="st-abc", name="st-abc", cluster_id="c-abc") @@ -62,7 +62,7 @@ def test_studio_set_env_partial_true(): studio._studio_api.set_env.assert_called_once_with(mock_studio, "ts-abc", new_env, partial=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_studio_set_env_partial_false(): """Test that Studio.set_env calls StudioApi.set_env correctly with partial=False.""" mock_studio = V1CloudSpace(id="st-abc", name="st-abc", cluster_id="c-abc") @@ -80,7 +80,7 @@ def test_studio_set_env_partial_false(): studio._studio_api.set_env.assert_called_once_with(mock_studio, "ts-abc", new_env, partial=False) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_studio_set_env_validates_names_and_accepts_empty_values(): studio = Studio.__new__(Studio) studio._studio = V1CloudSpace(id="st-abc", name="st-abc", cluster_id="c-abc") @@ -99,7 +99,7 @@ def test_studio_set_env_validates_names_and_accepts_empty_values(): assert studio._studio_api.set_env.call_count == 1 -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_studio_delete_env_preserves_unrelated_variables(): studio = Studio.__new__(Studio) studio._studio = V1CloudSpace(id="st-abc", name="st-abc", cluster_id="c-abc") @@ -114,7 +114,7 @@ def test_studio_delete_env_preserves_unrelated_variables(): studio._studio_api.set_env.assert_called_once_with(studio._studio, "ts-abc", {"KEEP": "yes"}, partial=False) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_studio_delete_env_handles_final_missing_and_invalid_names(): studio = Studio.__new__(Studio) studio._studio = V1CloudSpace(id="st-abc", name="st-abc", cluster_id="c-abc") diff --git a/python/tests/core/studio/test_studio_init.py b/python/tests/core/studio/test_studio_init.py index 98566738d..6d8529065 100644 --- a/python/tests/core/studio/test_studio_init.py +++ b/python/tests/core/studio/test_studio_init.py @@ -29,7 +29,7 @@ from lightning_sdk.studio import Studio -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_studio_exposes_id(internal_studio_init_mocker): studio = Studio(name="st-abc", teamspace="org-abc/ts-abc") @@ -46,7 +46,7 @@ def _list_cloudspaces_side_effect(*args, **kwargs): return _list_cloudspaces_side_effect -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_studio_exposes_placement_group_id(internal_studio_init_mocker): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") studio._studio = V1CloudSpace( @@ -58,7 +58,7 @@ def test_studio_exposes_placement_group_id(internal_studio_init_mocker): assert studio.placement_group_id == "pg-1" -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_studio_placement_group_id_is_none_without_compute_config(internal_studio_init_mocker): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") studio._studio = V1CloudSpace(id="st-abc", name="st-abc") @@ -89,7 +89,7 @@ def test_studio_placement_group_id_is_none_without_compute_config(internal_studi ) @mock.patch("lightning_sdk.api.org_api.OrgApi.get_org", autospec=True) @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.get_teamspace", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_studio_init( mock_get_teamspace, mock_get_org, @@ -201,7 +201,7 @@ def _get_status_side_effect(*args, **kwargs): @mock.patch.dict(os.environ, {"LIGHTNING_TEAMSPACE": ""}, clear=False) @mock.patch("lightning_sdk.utils.config.Config.get_value") @mock.patch("lightning_sdk.studio._resolve_teamspace") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_studio_init_no_teamspace(mock_resolve_teamspace, mock_config_get_value): # Mock config to return None for teamspace_name mock_config_get_value.return_value = None @@ -240,7 +240,7 @@ def test_studio_init_no_teamspace(mock_resolve_teamspace, mock_config_get_value) ) @mock.patch("lightning_sdk.api.org_api.OrgApi.get_org", autospec=True) @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.get_teamspace", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_studio_init_detects_current_studio_from_env( mock_get_teamspace, mock_get_org, @@ -331,7 +331,7 @@ def _get_status_side_effect(*args, **kwargs): ) @mock.patch("lightning_sdk.api.org_api.OrgApi.get_org", autospec=True) @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.get_teamspace", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_studio_init_uses_current_studio_cloud_account( mock_get_teamspace, mock_get_org, @@ -419,7 +419,7 @@ def _get_status_side_effect(*args, **kwargs): ) @mock.patch("lightning_sdk.api.org_api.OrgApi.get_org", autospec=True) @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.get_teamspace", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_studio_init_inherits_current_studio_template( mock_get_teamspace, mock_get_org, @@ -504,7 +504,7 @@ def _get_status_side_effect(*args, **kwargs): ) @mock.patch("lightning_sdk.api.org_api.OrgApi.get_org", autospec=True) @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.get_teamspace", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_studio_init_defaults_to_first_enabled_template_outside_studio( mock_get_teamspace, mock_get_org, @@ -586,7 +586,7 @@ def _get_status_side_effect(*args, **kwargs): "lightning_sdk.lightning_cloud.openapi.api.cluster_service_api.ClusterServiceApi.cluster_service_list_clusters", autospec=True, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_studio_init_with_cloud_provider( mock_list_clusters, mock_list_project_clusters, @@ -697,7 +697,7 @@ def _create_lightning_run_side_effect(*args, **kwargs): "lightning_sdk.lightning_cloud.openapi.api.cluster_service_api.ClusterServiceApi.cluster_service_list_clusters", autospec=True, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_studio_init_with_cloud_provider_string( mock_list_clusters, mock_list_project_clusters, @@ -800,7 +800,7 @@ def _create_lightning_run_side_effect(*args, **kwargs): "lightning_sdk.lightning_cloud.openapi.api.cluster_service_api.ClusterServiceApi.cluster_service_list_clusters", autospec=True, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_studio_init_with_cloud_for_existing_studio( mock_list_clusters, mock_list_project_clusters, @@ -888,7 +888,7 @@ def test_studio_init_with_cloud_for_existing_studio( ) @mock.patch("lightning_sdk.api.org_api.OrgApi.get_org", autospec=True) @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.get_teamspace", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_studio_init_opens_existing_studio_when_current_studio_in_different_teamspace( mock_get_teamspace, mock_get_org, @@ -962,7 +962,7 @@ def test_studio_init_opens_existing_studio_when_current_studio_in_different_team ) @mock.patch("lightning_sdk.api.org_api.OrgApi.get_org", autospec=True) @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.get_teamspace", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_studio_init_creates_studio_when_current_studio_in_different_teamspace( mock_get_teamspace, mock_get_org, diff --git a/python/tests/core/studio/test_studio_ports.py b/python/tests/core/studio/test_studio_ports.py index 877ccde22..8a00333b8 100644 --- a/python/tests/core/studio/test_studio_ports.py +++ b/python/tests/core/studio/test_studio_ports.py @@ -75,7 +75,7 @@ def _list_cloudspaces_side_effect(*args, **kwargs): @mock.patch("requests.put", autospec=True) @mock.patch("lightning_sdk.api.org_api.OrgApi.get_org", autospec=True) @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.get_teamspace", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_add_port( mock_get_teamspace, mock_get_org, @@ -181,7 +181,7 @@ def _create_endpoint_side_effect(*args, **kwargs): ) @mock.patch("lightning_sdk.api.org_api.OrgApi.get_org", autospec=True) @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.get_teamspace", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_list_ports( mock_get_teamspace, mock_get_org, diff --git a/python/tests/core/studio/test_studio_prevent_refetch.py b/python/tests/core/studio/test_studio_prevent_refetch.py index e0747f0a4..01f6e34a4 100644 --- a/python/tests/core/studio/test_studio_prevent_refetch.py +++ b/python/tests/core/studio/test_studio_prevent_refetch.py @@ -61,7 +61,7 @@ ) @mock.patch("lightning_sdk.api.org_api.OrgApi.get_org", autospec=True) @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.get_teamspace", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_status_and_machine_use_cached_code_status_when_refetch_prevented( mock_get_teamspace, mock_get_org, @@ -148,7 +148,7 @@ def test_status_and_machine_use_cached_code_status_when_refetch_prevented( ) @mock.patch("lightning_sdk.api.org_api.OrgApi.get_org", autospec=True) @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.get_teamspace", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_status_still_refetches_when_not_prevented( mock_get_teamspace, mock_get_org, diff --git a/python/tests/core/studio/test_studio_rename.py b/python/tests/core/studio/test_studio_rename.py index 0c812e323..92fa1330f 100644 --- a/python/tests/core/studio/test_studio_rename.py +++ b/python/tests/core/studio/test_studio_rename.py @@ -14,7 +14,7 @@ def _list_cloudspaces_side_effect(*args, **kwargs): return _list_cloudspaces_side_effect -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_studio_rename(): """Test that Studio.rename calls StudioApi._update_cloudspace and _update_studio_reference correctly.""" # Setup diff --git a/python/tests/core/studio/test_studio_start.py b/python/tests/core/studio/test_studio_start.py index 332efd7a6..4a83b92a1 100644 --- a/python/tests/core/studio/test_studio_start.py +++ b/python/tests/core/studio/test_studio_start.py @@ -87,7 +87,7 @@ def _list_cloudspaces_side_effect(*args, **kwargs): "lightning_sdk.api.studio_api.StudioApi.machine_is_supported", autospec=True, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_studio_start( mock_machine_is_supported, mock_get_teamspace, @@ -283,7 +283,7 @@ def _create_lightning_run_side_effect(body, project_id, cloudspace_id, **kwargs) ) @mock.patch("lightning_sdk.api.org_api.OrgApi.get_org", autospec=True) @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.get_teamspace", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_studio_start_on_demand_machine( mock_get_teamspace, mock_get_org, @@ -476,7 +476,7 @@ def _create_lightning_run_side_effect(*args, **kwargs): ) @mock.patch("lightning_sdk.api.org_api.OrgApi.get_org", autospec=True) @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.get_teamspace", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_studio_start_interruptible_override( mock_get_teamspace, mock_get_org, @@ -668,7 +668,7 @@ def _create_lightning_run_side_effect(*args, **kwargs): ) @mock.patch("lightning_sdk.api.org_api.OrgApi.get_org", autospec=True) @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.get_teamspace", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_studio_start_different_machine( mock_get_teamspace, mock_get_org, @@ -878,7 +878,7 @@ def _create_lightning_run_side_effect(*args, **kwargs): "lightning_sdk.api.studio_api.StudioApi.machine_is_supported", autospec=True, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_studio_start_uses_current_studio_machine_when_inside_running_studio( mock_machine_is_supported, mock_get_teamspace, @@ -1077,7 +1077,7 @@ def _get_cloud_space_side_effect(*args, **kwargs): ) @mock.patch("lightning_sdk.api.org_api.OrgApi.get_org", autospec=True) @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.get_teamspace", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_studio_start_ignores_stopped_current_studio_machine( mock_get_teamspace, mock_get_org, @@ -1253,7 +1253,7 @@ def _get_cloud_space_side_effect(*args, **kwargs): ) @mock.patch("lightning_sdk.api.org_api.OrgApi.get_org", autospec=True) @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.get_teamspace", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_studio_start_explicit_machine_overrides_current_studio_machine( mock_get_teamspace, mock_get_org, diff --git a/python/tests/core/studio/test_studio_status.py b/python/tests/core/studio/test_studio_status.py index fa738edfe..e94fd1077 100644 --- a/python/tests/core/studio/test_studio_status.py +++ b/python/tests/core/studio/test_studio_status.py @@ -58,7 +58,7 @@ def _list_cloudspaces_side_effect(*args, **kwargs): ) @mock.patch("lightning_sdk.api.org_api.OrgApi.get_org", autospec=True) @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.get_teamspace", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_studio_status( mock_get_teamspace, mock_get_org, diff --git a/python/tests/core/studio/test_studio_vm.py b/python/tests/core/studio/test_studio_vm.py index bc45e7a6b..087701e4e 100644 --- a/python/tests/core/studio/test_studio_vm.py +++ b/python/tests/core/studio/test_studio_vm.py @@ -43,7 +43,7 @@ def _list_cloudspaces_side_effect(*args, **kwargs): ) @mock.patch("lightning_sdk.api.org_api.OrgApi.get_org", autospec=True) @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.get_teamspace", autospec=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_studio_vm_cls_name( mock_get_teamspace, mock_get_org, diff --git a/python/tests/core/test_agents.py b/python/tests/core/test_agents.py index 30e0bc7ae..defbd06db 100644 --- a/python/tests/core/test_agents.py +++ b/python/tests/core/test_agents.py @@ -4,20 +4,20 @@ from lightning_sdk.lightning_cloud.openapi import V1Assistant -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_agent_init(internal_agents_api_get_agent_mocker): agent = Agent("ag-abc") assert isinstance(agent._agent, V1Assistant) assert agent._agent.name == "ag-abc" -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_agent_delete(internal_agents_api_delete_agent_mocker, internal_agents_api_get_agent_mocker): agent = Agent("ag-abc") agent.delete() -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_agent_update( internal_agents_api_update_agent_mocker, internal_agents_api_get_agent_mocker, diff --git a/python/tests/core/test_deployment.py b/python/tests/core/test_deployment.py index 1dda66677..1e9b648b5 100644 --- a/python/tests/core/test_deployment.py +++ b/python/tests/core/test_deployment.py @@ -27,7 +27,7 @@ from lightning_sdk.machine import Machine -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) @patch("lightning_sdk.deployment.login.Auth", new=MagicMock()) def test_deployment_resolve_teamspace(monkeypatch): mock_org = MagicMock() @@ -229,7 +229,7 @@ def test_deployment_secret_mutations_reject_missing_invalid_and_uncreated(): deployment.update.assert_not_called() -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) @patch("lightning_sdk.deployment.login.Auth", new=MagicMock()) def test_unstarted_deployment_id_is_none(monkeypatch): _mock_deployment_dependencies(monkeypatch, MagicMock()) @@ -240,7 +240,7 @@ def test_unstarted_deployment_id_is_none(monkeypatch): assert deployment.id is None -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) @patch("lightning_sdk.deployment.login.Auth", new=MagicMock()) def test_existing_deployment_exposes_id(monkeypatch): _mock_deployment_dependencies(monkeypatch, MagicMock()) @@ -251,7 +251,7 @@ def test_existing_deployment_exposes_id(monkeypatch): assert deployment.id == "dep-123" -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) @patch("lightning_sdk.deployment.login.Auth", new=MagicMock()) def test_deployment_does_not_authenticate_or_look_up_user_eagerly(monkeypatch): """The constructor must not eagerly authenticate or resolve the authenticated user. @@ -273,7 +273,7 @@ def test_deployment_does_not_authenticate_or_look_up_user_eagerly(monkeypatch): assert deployment._teamspace is teamspace_mock -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) @patch("lightning_sdk.deployment.login.Auth", new=MagicMock()) def test_deployment_uses_explicit_owner(monkeypatch): """The deployment owner comes from the explicitly provided user/org, not the auth token.""" @@ -292,7 +292,7 @@ def test_deployment_uses_explicit_owner(monkeypatch): assert deployment._user is None -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) @patch("lightning_sdk.deployment.login.Auth", new=MagicMock()) def test_deployment_raises_when_teamspace_cannot_be_resolved(monkeypatch): """If neither a teamspace nor an owner can be resolved, fail with an actionable error.""" @@ -307,7 +307,7 @@ def test_deployment_raises_when_teamspace_cannot_be_resolved(monkeypatch): deployment_module.Deployment(name="my-deploy") -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) @patch("lightning_sdk.deployment.login.Auth", new=MagicMock()) def test_deployment_export_request_captures_wrapper(tmp_path): result = object() @@ -347,7 +347,7 @@ def test_deployment_export_request_captures_wrapper(tmp_path): ) -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) @patch("lightning_sdk.deployment.login.Auth", new=MagicMock()) def test_to_autoscaling(): with pytest.raises(ValueError, match="The autoscaling metric is required. Currently supported metrics are"): @@ -540,7 +540,7 @@ def test_restore_autoscale_returns_one_config(): assert restored.max_replicas == 3 -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) @patch("lightning_sdk.deployment.login.Auth", new=MagicMock()) def test_to_env(): env = deployment_api_module.to_env({"key": "value"}) @@ -554,7 +554,7 @@ def test_to_env(): assert env == [V1EnvVar(name="ENV_ALIAS", from_secret="source_secret")] -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) @patch("lightning_sdk.deployment.login.Auth", new=MagicMock()) def test_restore_env_preserves_secret_alias(): restored = deployment_api_module.restore_env( @@ -579,7 +579,7 @@ def test_to_env_distinguishes_no_change_from_clear(): assert deployment_api_module.to_env([]) == [] -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) @patch("lightning_sdk.deployment.login.Auth", new=MagicMock()) def test_to_endpoint(): with pytest.raises(ValueError, match="At least one port is required to reach your deployment."): @@ -618,7 +618,7 @@ def test_to_endpoint(): assert endpoint.custom_domain == "custom_domain" -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) @patch("lightning_sdk.deployment.login.Auth", new=MagicMock()) def test_deployment_start_first_time(monkeypatch): monkeypatch.setattr(deployment_module, "User", MagicMock()) @@ -704,7 +704,7 @@ def fn(*_, **__): assert deployment.running_replicas == 4 -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) @patch("lightning_sdk.deployment.login.Auth", new=MagicMock()) def test_deployment_start_with_cloud(monkeypatch): monkeypatch.setattr(deployment_module, "User", MagicMock()) @@ -799,7 +799,7 @@ def fn(*_, **__): ) -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) @patch("lightning_sdk.deployment.login.Auth", new=MagicMock()) def test_deployment_start_already_exist(monkeypatch): monkeypatch.setattr(deployment_module, "User", MagicMock()) @@ -827,7 +827,7 @@ def test_deployment_start_already_exist(monkeypatch): assert call_args.kwargs["body"].desired_state == deployment_module.V1DeploymentState.RUNNING -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) @patch("lightning_sdk.deployment.login.Auth", new=MagicMock()) def test_deployment_update(monkeypatch): monkeypatch.setattr(deployment_module, "User", MagicMock()) @@ -896,7 +896,7 @@ def test_deployment_update(monkeypatch): assert deployment.command == "python server2.py" -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) @patch("lightning_sdk.deployment.login.Auth", new=MagicMock()) def test_deployment_update_injects_secret_alias(monkeypatch): monkeypatch.setattr(deployment_module, "User", MagicMock()) @@ -926,7 +926,7 @@ def test_deployment_update_injects_secret_alias(monkeypatch): assert spec.env == [V1EnvVar(name="ENV_ALIAS", from_secret="source_secret")] -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) @patch("lightning_sdk.deployment.login.Auth", new=MagicMock()) def test_deployment_update_name(monkeypatch): monkeypatch.setattr(deployment_module, "User", MagicMock()) @@ -951,7 +951,7 @@ def test_deployment_update_name(monkeypatch): assert deployment.name == "new_entrypoint" -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) @patch("lightning_sdk.deployment.login.Auth", new=MagicMock()) def test_deployment_update_replicas(monkeypatch): monkeypatch.setattr(deployment_module, "User", MagicMock()) @@ -978,7 +978,7 @@ def test_deployment_update_replicas(monkeypatch): assert deployment.replicas == 0 -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) @patch("lightning_sdk.deployment.login.Auth", new=MagicMock()) def test_deployment_update_strategy(monkeypatch): monkeypatch.setattr(deployment_module, "User", MagicMock()) @@ -1005,7 +1005,7 @@ def test_deployment_update_strategy(monkeypatch): assert deployment._deployment.strategy is not None -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) @patch("lightning_sdk.deployment.login.Auth", new=MagicMock()) def test_deployment_stop(monkeypatch): monkeypatch.setattr(deployment_module, "User", MagicMock()) @@ -1058,7 +1058,7 @@ def update(*_, body, **__): assert deployment.image == "my-image" -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) @patch("lightning_sdk.deployment.login.Auth", new=MagicMock()) def test_deployment_get(monkeypatch): monkeypatch.setattr(deployment_module, "User", MagicMock()) @@ -1094,7 +1094,7 @@ def test_deployment_get(monkeypatch): requests_mock.Session().get.assert_called() -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) @patch("lightning_sdk.deployment.login.Auth", new=MagicMock()) def test_deployment_post(monkeypatch): monkeypatch.setattr(deployment_module, "User", MagicMock()) @@ -1130,7 +1130,7 @@ def test_deployment_post(monkeypatch): requests_mock.Session().post.assert_called_with(f"{url}/api/chat", json=json, verify=False) -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) @patch("lightning_sdk.deployment.login.Auth", new=MagicMock()) def test_deployment_put(monkeypatch): monkeypatch.setattr(deployment_module, "User", MagicMock()) @@ -1165,7 +1165,7 @@ def test_deployment_put(monkeypatch): requests_mock.Session().put.assert_called_with(f"{url}/", verify=False) -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) @patch("lightning_sdk.deployment.login.Auth", new=MagicMock()) def test_deployment_delete(monkeypatch): monkeypatch.setattr(deployment_module, "User", MagicMock()) @@ -1200,7 +1200,7 @@ def test_deployment_delete(monkeypatch): requests_mock.Session().delete.assert_called_with(f"{url}/", verify=False) -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) @patch("lightning_sdk.deployment.login.Auth", new=MagicMock()) def test_to_spec(): with pytest.raises(ValueError, match="The cloud account should be defined."): @@ -1220,7 +1220,7 @@ def test_to_spec(): deployment_api_module.to_spec("cluster-id", Machine.CPU, "image", None, None) -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) @patch("lightning_sdk.deployment.login.Auth", new=MagicMock()) def test_compose_commands(): commands = ["python server.py &", "python server.py &"] @@ -1232,7 +1232,7 @@ def test_compose_commands(): assert command == "( python server.py & ) && ls && ( python server.py & )" -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) @patch("lightning_sdk.deployment.login.Auth", new=MagicMock()) def test_to_health_check_empty(): health_check = deployment_api_module.to_health_check() @@ -1245,7 +1245,7 @@ def test_to_health_check_empty(): assert health_check is None -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) @patch("lightning_sdk.deployment.login.Auth", new=MagicMock()) def test_apply_change(): to_health_check() @@ -1260,7 +1260,7 @@ def test_apply_change(): assert apply_change(spec, "readiness_probe", to_health_check(None)) -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) @patch("lightning_sdk.deployment.login.Auth", new=MagicMock()) def test_deployment_start_with_path_mappings(monkeypatch): monkeypatch.setattr(deployment_module, "User", MagicMock()) @@ -1301,7 +1301,7 @@ def fn(*_, **__): assert len(spec.path_mappings) == 2 -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) @patch("lightning_sdk.deployment.login.Auth", new=MagicMock()) def test_deployment_update_with_path_mappings(monkeypatch): monkeypatch.setattr(deployment_module, "User", MagicMock()) @@ -1338,7 +1338,7 @@ def test_deployment_update_with_path_mappings(monkeypatch): assert len(spec.path_mappings) == 2 -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) @patch("lightning_sdk.deployment.login.Auth", new=MagicMock()) def test_deployment_start_byom_builds_spec(monkeypatch): monkeypatch.setattr(deployment_module, "User", MagicMock()) diff --git a/python/tests/core/test_job.py b/python/tests/core/test_job.py index a8256d732..3e66a025c 100644 --- a/python/tests/core/test_job.py +++ b/python/tests/core/test_job.py @@ -23,7 +23,7 @@ @mock.patch.dict(os.environ, clear=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_job_init( internal_studio_init_mocker, internal_studio_status_mocker, @@ -37,7 +37,7 @@ def test_job_init( @mock.patch.dict(os.environ, clear=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_job_init_error( internal_studio_init_mocker, internal_studio_status_mocker, @@ -66,7 +66,7 @@ def init_job_error(): ("j-stu", Status.Completed), # "completed" ], ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_job_status( internal_job_api_mocker_get_job_status_v2, internal_studio_init_mocker, @@ -84,7 +84,7 @@ def test_job_status( @mock.patch.dict(os.environ, clear=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_stop_job( internal_job_api_mocker_stop_job, internal_studio_init_mocker, @@ -102,7 +102,7 @@ def test_stop_job( @mock.patch.dict(os.environ, clear=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_delete_job( internal_job_api_mocker_delete_job_v2, internal_studio_init_mocker, @@ -126,7 +126,7 @@ def test_delete_job( @pytest.mark.parametrize("command", [None, "echo hello"]) @pytest.mark.parametrize("env", [None, {"key": "value"}]) @pytest.mark.parametrize("interruptible", [True, False]) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_submit_job_v2_image(internal_studio_init_mocker, machine, command, env, interruptible): teamspace = Teamspace("ts-abc", org="org-abc") job = Job("test-job", teamspace, _fetch_job=False) @@ -278,7 +278,7 @@ def test_download_logs_multi_machine_writes_merged_timeline(tmp_path, monkeypatc assert kwargs["follow"] is False -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_job_exposes_private_provisioning_metadata(internal_studio_init_mocker): teamspace = Teamspace("ts-abc", org="org-abc") job = Job("test-job", teamspace, _fetch_job=False) @@ -296,7 +296,7 @@ def test_job_exposes_private_provisioning_metadata(internal_studio_init_mocker): assert job.rank == 3 -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_job_exposes_run_attempt_metadata(internal_studio_init_mocker): teamspace = Teamspace("ts-abc", org="org-abc") job = Job("test-job", teamspace, _fetch_job=False) @@ -319,7 +319,7 @@ def test_job_exposes_run_attempt_metadata(internal_studio_init_mocker): assert getattr(unset._guaranteed_job, "parent_multi_machine_job_id", None) is None -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_job_exposes_mmt_level_run_attempt_metadata(internal_studio_init_mocker): teamspace = Teamspace("ts-abc", org="org-abc") job = Job("test-mmt", teamspace, _fetch_job=False, _num_machines=2) @@ -341,7 +341,7 @@ def test_job_exposes_mmt_level_run_attempt_metadata(internal_studio_init_mocker) assert job._guaranteed_job.spec.parent_job_id is None # standalone-only field -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_job_exposes_start_and_stop_times(internal_studio_init_mocker): teamspace = Teamspace("ts-abc", org="org-abc") job = Job("test-job", teamspace, _fetch_job=False) @@ -359,7 +359,7 @@ def test_job_exposes_start_and_stop_times(internal_studio_init_mocker): assert job.stopped_at is None -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_multi_machine_job_exposes_start_and_stop_times(internal_studio_init_mocker): teamspace = Teamspace("ts-abc", org="org-abc") job = Job("test-mmt", teamspace, _fetch_job=False) @@ -381,7 +381,7 @@ def test_multi_machine_job_exposes_start_and_stop_times(internal_studio_init_moc assert job.stopped_at == datetime(2026, 8, 2, 13, 0, tzinfo=timezone.utc) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_job_id_is_none_without_backing_model_or_api_request(internal_studio_init_mocker): teamspace = Teamspace("org-abc/ts-abc") job = Job("test-job", teamspace, _fetch_job=False) @@ -391,7 +391,7 @@ def test_job_id_is_none_without_backing_model_or_api_request(internal_studio_ini assert job.resource_id is None -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_submit_job_threads_placement_group_id(internal_studio_init_mocker): teamspace = Teamspace("ts-abc", org="org-abc") job = Job("test-job", teamspace, _fetch_job=False) @@ -409,7 +409,7 @@ def test_submit_job_threads_placement_group_id(internal_studio_init_mocker): assert submit_mock.call_args.kwargs["placement_group_id"] == "pg-1" -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_submit_job_threads_max_run_attempts(internal_studio_init_mocker): teamspace = Teamspace("ts-abc", org="org-abc") job = Job("test-job", teamspace, _fetch_job=False) @@ -427,7 +427,7 @@ def test_submit_job_threads_max_run_attempts(internal_studio_init_mocker): assert submit_mock.call_args.kwargs["max_run_attempts"] == 3 -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_submit_job_threads_keep_minutes(internal_studio_init_mocker): teamspace = Teamspace("ts-abc", org="org-abc") job = Job("test-job", teamspace, _fetch_job=False) @@ -445,7 +445,7 @@ def test_submit_job_threads_keep_minutes(internal_studio_init_mocker): assert submit_mock.call_args.kwargs["keep_minutes"] == 30 -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_job_exposes_keep_minutes(internal_studio_init_mocker): teamspace = Teamspace("ts-abc", org="org-abc") job = Job("test-job", teamspace, _fetch_job=False) @@ -458,7 +458,7 @@ def test_job_exposes_keep_minutes(internal_studio_init_mocker): assert job.keep_minutes == 30 -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_job_keep_minutes_is_none_when_unset(internal_studio_init_mocker): teamspace = Teamspace("ts-abc", org="org-abc") job = Job("test-job", teamspace, _fetch_job=False) @@ -467,7 +467,7 @@ def test_job_keep_minutes_is_none_when_unset(internal_studio_init_mocker): assert job.keep_minutes is None -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_submit_threads_max_run_attempts_for_multi_machine(internal_studio_init_mocker): teamspace = Teamspace("ts-abc", org="org-abc") job = Job("test-job", teamspace, _fetch_job=False) @@ -489,7 +489,7 @@ def test_submit_threads_max_run_attempts_for_multi_machine(internal_studio_init_ assert submit_mock.call_args.kwargs["max_run_attempts"] == 3 -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_submit_auto_derives_fault_tolerance_for_multi_machine_with_retries(internal_studio_init_mocker): teamspace = Teamspace("ts-abc", org="org-abc") job = Job("test-mmt", teamspace, _fetch_job=False) @@ -515,7 +515,7 @@ def test_submit_auto_derives_fault_tolerance_for_multi_machine_with_retries(inte assert body.fault_tolerance.strategy == "MULTI_MACHINE_JOB_FAULT_TOLERANCE_STRATEGY_RECREATE_ALL_NODES" -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_submit_omits_fault_tolerance_for_multi_machine_without_retries(internal_studio_init_mocker): teamspace = Teamspace("ts-abc", org="org-abc") job = Job("test-mmt", teamspace, _fetch_job=False) @@ -541,7 +541,7 @@ def test_submit_omits_fault_tolerance_for_multi_machine_without_retries(internal @pytest.mark.parametrize("machine", [Machine.L4, Machine.DATA_PREP_MAX]) @pytest.mark.parametrize("env", [None, {"key": "value"}]) @pytest.mark.parametrize("interruptible", [True, False]) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_submit_job_v2_studio(internal_studio_init_mocker, machine, env, interruptible): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") @@ -582,7 +582,7 @@ def test_submit_job_v2_studio(internal_studio_init_mocker, machine, env, interru ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_jobv2_run_arg_validation(internal_studio_init_mocker): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") @@ -600,7 +600,7 @@ def test_jobv2_run_arg_validation(internal_studio_init_mocker): Job.run(None, Machine.CPU) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_jobv2_run_entrypoint_validation(internal_studio_init_mocker): """Test entrypoint validation logic for image jobs.""" studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") @@ -664,7 +664,7 @@ def test_jobv2_run_entrypoint_validation(internal_studio_init_mocker): ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_submit_jobv2_error_cases(internal_studio_init_mocker): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") @@ -768,7 +768,7 @@ def test_submit_jobv2_error_cases(internal_studio_init_mocker): ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_job_by_name(internal_studio_init_mocker): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") job = Job("test-job", studio.teamspace, _fetch_job=False) @@ -781,7 +781,7 @@ def test_get_job_by_name(internal_studio_init_mocker): job._job_api.get_job_by_name.assert_called_once_with(name="test-job", teamspace_id="ts-abc001") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_job_by_id(internal_studio_init_mocker): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") job = Job("test-job", studio.teamspace, _fetch_job=False) @@ -797,7 +797,7 @@ def test_get_job_by_id(internal_studio_init_mocker): job._job_api.get_job.assert_called_once_with(job_id="test-job-id", teamspace_id="ts-abc001") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_job_by_name_first_and_then_by_id(internal_studio_init_mocker): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") job = Job("test-job", studio.teamspace, _fetch_job=False) @@ -815,7 +815,7 @@ def test_get_job_by_name_first_and_then_by_id(internal_studio_init_mocker): job._job_api.get_job.assert_called_once_with(job_id="test-job-id", teamspace_id="ts-abc001") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_job_by_name_on_init(job_api_get_job_by_name_mocker, internal_studio_init_mocker): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") job = Job("test-job", studio.teamspace) @@ -838,7 +838,7 @@ def test_get_job_by_name_on_init(job_api_get_job_by_name_mocker, internal_studio ("complete", Status.Stopping), ], ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_jobv2_status(job_api_get_job_by_name_mocker, internal_studio_init_mocker, internal_status, external_status): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") @@ -862,7 +862,7 @@ def test_jobv2_status(job_api_get_job_by_name_mocker, internal_studio_init_mocke ("", "unknown", Machine.from_str("unknown")), ], ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_jobv2_machine( internal_studio_init_mocker, internal_studio_api_mocker_get_machine, @@ -887,7 +887,7 @@ def test_jobv2_machine( get_job_mock.assert_called_once() -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_jobv2_stop(job_api_get_job_by_name_mocker, internal_studio_init_mocker): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") @@ -918,7 +918,7 @@ def get_job_side_effect(*args, **kwargs): ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_jobv2_delete(job_api_get_job_by_name_mocker, internal_studio_init_mocker): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") @@ -935,7 +935,7 @@ def test_jobv2_delete(job_api_get_job_by_name_mocker, internal_studio_init_mocke delete_job_mock.assert_called_once_with(job_id="test-job-id", teamspace_id="ts-abc001", cloudspace_id="st-abc") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_submit_jobv2_studio_resolve( internal_get_org_api_mocker, internal_teamspace_api_mocker, @@ -983,7 +983,7 @@ def test_submit_jobv2_studio_resolve( (None, "ubuntu", None), ], ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_job_artifacts_uri( internal_get_org_api_mocker, internal_teamspace_api_mocker, @@ -1005,7 +1005,7 @@ def test_job_artifacts_uri( assert job.artifacts_uri == expected_uri -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_job_artifact_helpers_address_the_drive_path( internal_get_org_api_mocker, internal_teamspace_api_mocker, @@ -1034,7 +1034,7 @@ def test_job_artifact_helpers_address_the_drive_path( download_folder.assert_called_once_with("jobs/test-job", "./out") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_job_without_artifacts_lists_nothing_and_refuses_to_download( internal_get_org_api_mocker, internal_teamspace_api_mocker, @@ -1052,7 +1052,7 @@ def test_job_without_artifacts_lists_nothing_and_refuses_to_download( job.download_artifacts("./out") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_multi_machine_job_artifacts_fan_out_over_machines( internal_get_org_api_mocker, internal_teamspace_api_mocker, @@ -1093,7 +1093,7 @@ def test_multi_machine_job_artifacts_fan_out_over_machines( ] -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_submit_job_v2_image_from_studio( internal_get_org_api_mocker, internal_teamspace_api_mocker, @@ -1152,7 +1152,7 @@ def test_submit_job_v2_image_from_studio( @mock.patch("lightning_sdk.studio._internal_status_to_external_status", return_value=Status.Running) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_run_job_with_cloud_provider( mock_in_studio, internal_get_org_api_mocker, @@ -1198,7 +1198,7 @@ def test_run_job_with_cloud_provider( ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_job_logs_v2( internal_job_logs_mocker, job_api_get_job_by_name_mocker, @@ -1215,7 +1215,7 @@ def test_job_logs_v2( assert job.logs == "âš¡ ~ echo Hello\nHello" -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_job_v2_dict_json(internal_studio_init_mocker, internal_studio_api_mocker_get_machine): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") @@ -1257,7 +1257,7 @@ def test_job_v2_dict_json(internal_studio_init_mocker, internal_studio_api_mocke @pytest.mark.parametrize("target_state", ["stopped", "completed", "failed"]) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_jobv2_wait(job_api_get_job_by_name_mocker, internal_studio_init_mocker, target_state): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") @@ -1288,7 +1288,7 @@ def get_job_side_effect(*args, **kwargs): assert get_job_mock.call_count == 11 -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_jobv2_wait_timeout(job_api_get_job_by_name_mocker, internal_studio_init_mocker): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") @@ -1310,7 +1310,7 @@ def get_job_side_effect(*args, **kwargs): @pytest.mark.asyncio() @pytest.mark.parametrize("target_state", ["stopped", "completed", "failed"]) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) async def test_jobv2_wait_async(job_api_get_job_by_name_mocker, internal_studio_init_mocker, target_state): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") job = Job("test-job", studio.teamspace) @@ -1340,7 +1340,7 @@ def get_job_side_effect(*args, **kwargs): @pytest.mark.asyncio() -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) async def test_jobv2_wait_async_timeout(job_api_get_job_by_name_mocker, internal_studio_init_mocker): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") job = Job("test-job", studio.teamspace) @@ -1360,7 +1360,7 @@ def get_job_side_effect(*args, **kwargs): @mock.patch("lightning_sdk.studio._internal_status_to_external_status", return_value=Status.Running) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_submit_job_from_running_studio( internal_get_org_api_mocker, internal_teamspace_api_mocker, @@ -1400,7 +1400,7 @@ def _stub_logs_api(job, entries=()): return stream_mock -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_job_logs_reads_logs_api(job_api_get_job_by_name_mocker, internal_studio_init_mocker): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") job = Job("test-job", studio.teamspace) @@ -1420,7 +1420,7 @@ def test_job_logs_reads_logs_api(job_api_get_job_by_name_mocker, internal_studio assert stream_mock.call_args.args == (job.teamspace.id,) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_job_logs_timestamps_and_filters(job_api_get_job_by_name_mocker, internal_studio_init_mocker): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") job = Job("test-job", studio.teamspace) @@ -1437,7 +1437,7 @@ def test_job_logs_timestamps_and_filters(job_api_get_job_by_name_mocker, interna assert kwargs["severity"] == "error" -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_job_logs_finished_raises_while_migrating(job_api_get_job_by_name_mocker, internal_studio_init_mocker): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") job = Job("test-job", studio.teamspace) @@ -1450,7 +1450,7 @@ def test_job_logs_finished_raises_while_migrating(job_api_get_job_by_name_mocker str(job.logs) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_job_logs_filtered_read_raises_while_migrating(job_api_get_job_by_name_mocker, internal_studio_init_mocker): """A filter the server never applied must not be reported as "no matches".""" studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") @@ -1463,7 +1463,7 @@ def test_job_logs_filtered_read_raises_while_migrating(job_api_get_job_by_name_m job.logs(query="nope") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_job_logs_filtered_read_matching_nothing_is_empty(job_api_get_job_by_name_mocker, internal_studio_init_mocker): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") job = Job("test-job", studio.teamspace) @@ -1478,7 +1478,7 @@ def test_job_logs_filtered_read_matching_nothing_is_empty(job_api_get_job_by_nam assert job.logs(since="2026-01-01T00:00:00Z") == "" -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_job_logs_finished_tail(job_api_get_job_by_name_mocker, internal_studio_init_mocker): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") job = Job("test-job", studio.teamspace) @@ -1492,7 +1492,7 @@ def test_job_logs_finished_tail(job_api_get_job_by_name_mocker, internal_studio_ assert stream_mock.call_args.kwargs["tail"] == 2 -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_job_logs_rank_on_single_machine_rejected(job_api_get_job_by_name_mocker, internal_studio_init_mocker): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") job = Job("test-job", studio.teamspace) @@ -1504,7 +1504,7 @@ def test_job_logs_rank_on_single_machine_rejected(job_api_get_job_by_name_mocker _ = job.logs(rank=1) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_job_logs_rank_matching_own_machine_reads_own_logs(job_api_get_job_by_name_mocker, internal_studio_init_mocker): """A machine of an MMT reading its own rank reads its own job, no sibling lookup.""" studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") @@ -1520,7 +1520,7 @@ def test_job_logs_rank_matching_own_machine_reads_own_logs(job_api_get_job_by_na job._mmt_job_api.list_mmt_subjobs.assert_not_called() -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_job_logs_rank_resolves_sibling_machine(job_api_get_job_by_name_mocker, internal_studio_init_mocker): """Reading another rank resolves that machine's job id client-side and reads it.""" studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") @@ -1544,7 +1544,7 @@ def _stream(teamspace_id, **kwargs): assert job.logs(rank=2) == "rank 2 line" -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_job_logs_running_snapshot_falls_back_to_live(job_api_get_job_by_name_mocker, internal_studio_init_mocker): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") job = Job("test-job", studio.teamspace) @@ -1559,7 +1559,7 @@ def test_job_logs_running_snapshot_falls_back_to_live(job_api_get_job_by_name_mo assert kwargs["idle_timeout"] == _RUNNING_LOGS_IDLE_TIMEOUT -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_job_logs_follow_streams_from_logs_api(job_api_get_job_by_name_mocker, internal_studio_init_mocker): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") job = Job("test-job", studio.teamspace) @@ -1575,7 +1575,7 @@ def test_job_logs_follow_streams_from_logs_api(job_api_get_job_by_name_mocker, i assert callable(kwargs["stop"]) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_job_logs_follow_on_finished_returns_saved_lines(job_api_get_job_by_name_mocker, internal_studio_init_mocker): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") job = Job("test-job", studio.teamspace) @@ -1590,7 +1590,7 @@ def test_job_logs_follow_on_finished_returns_saved_lines(job_api_get_job_by_name assert kwargs["follow"] is False -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_job_logs_follow_with_rank_streams_that_machine(job_api_get_job_by_name_mocker, internal_studio_init_mocker): """follow + rank resolves the machine client-side and streams its logs from the logs API.""" studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") diff --git a/python/tests/core/test_k8s_cluster.py b/python/tests/core/test_k8s_cluster.py index b7c8eaa8d..e56aae435 100644 --- a/python/tests/core/test_k8s_cluster.py +++ b/python/tests/core/test_k8s_cluster.py @@ -23,7 +23,7 @@ ] -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_billing_usage_with_empty_metrics(): k8s_api = K8sCluster("test") get_billing_usage_mock = mock.Mock(return_value=[]) @@ -38,7 +38,7 @@ def test_get_billing_usage_with_empty_metrics(): get_billing_usage_mock.assert_called_once_with(print_data=False) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_billing_usage_with_no_range(): k8s_api = K8sCluster("test") get_billing_usage_mock = mock.Mock(return_value=mock_cluster) @@ -55,7 +55,7 @@ def test_get_billing_usage_with_no_range(): get_billing_usage_mock.assert_called_once_with(print_data=False) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_billing_usage_with_only_start(): k8s_api = K8sCluster("test") get_billing_usage_mock = mock.Mock(return_value=[mock_cluster[1]]) @@ -74,7 +74,7 @@ def test_get_billing_usage_with_only_start(): get_billing_usage_mock.assert_called_once_with(print_data=False, start=to_iso_z(date)) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_billing_usage_with_only_end(): k8s_api = K8sCluster("test") get_billing_usage_mock = mock.Mock(return_value=[mock_cluster[0]]) @@ -93,7 +93,7 @@ def test_get_billing_usage_with_only_end(): get_billing_usage_mock.assert_called_once_with(print_data=False, end=to_iso_z(date)) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_billing_usage_with_date_range(): k8s_api = K8sCluster("test") get_billing_usage_mock = mock.Mock(return_value=mock_cluster) diff --git a/python/tests/core/test_llm.py b/python/tests/core/test_llm.py index 8dcf7d769..6a1f5d6d8 100644 --- a/python/tests/core/test_llm.py +++ b/python/tests/core/test_llm.py @@ -86,7 +86,7 @@ def mock_tools(): ] -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) def test_get_auth_info_missing_teamspace_raises(monkeypatch): # Clear all shared state LLMCLIENT._auth_info_cached = False @@ -121,7 +121,7 @@ def test_get_auth_info_missing_teamspace_raises(monkeypatch): LLM(name="openai/gpt-4o", teamspace=None) -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) def test_get_model_id_uses_cache(): # Reset shared class state to avoid side effects LLMCLIENT._llm_api_cache.clear() @@ -136,7 +136,7 @@ def test_get_model_id_uses_cache(): assert llm.context_length == 8192 -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) def test_different_model_names(monkeypatch, mock_org_model2): LLMCLIENT._llm_api_cache.clear() mock_api = MagicMock() @@ -148,7 +148,7 @@ def test_different_model_names(monkeypatch, mock_org_model2): assert llm._model_name == "my-model/v1/new-model" -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) def test_invalid_format(monkeypatch, mock_model_data): mock_api = MagicMock() mock_api.list_models.return_value = mock_model_data @@ -164,7 +164,7 @@ def test_invalid_format(monkeypatch, mock_model_data): LLM("gpt-4o") -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) def test_org_model(monkeypatch, mock_org_model): LLMCLIENT._llm_api_cache.clear() mock_api = MagicMock() @@ -175,7 +175,7 @@ def test_org_model(monkeypatch, mock_org_model): assert llm.provider == "org-name" -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) def test_invalid_org(monkeypatch): # there could be a case where the model provider is an org that exists, however, the user does not have access to it # then it would make sense to search for whatever they have availabe in public, teamspace and org @@ -195,7 +195,7 @@ def test_invalid_org(monkeypatch): LLM("org-123/dummy-model") -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) def test_user_model(monkeypatch, mock_user_model): LLMCLIENT._llm_api_cache.clear() mock_api = MagicMock() @@ -207,7 +207,7 @@ def test_user_model(monkeypatch, mock_user_model): assert llm.provider == "user-name" -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) def test_get_auth_info(monkeypatch): LLMCLIENT._llm_api_cache.clear() LLMCLIENT._auth_info_cached = False @@ -224,7 +224,7 @@ def test_get_auth_info(monkeypatch): mock_resolve.assert_called_with(teamspace="teamspace-name", org=None, user="my-user") -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) def test_ephemeral(monkeypatch, mock_public_model): LLMCLIENT._auth_info_cached = False LLMCLIENT._llm_api_cache.clear() @@ -256,7 +256,7 @@ def test_ephemeral(monkeypatch, mock_public_model): ) -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) def test_chat(monkeypatch, mock_public_model): LLMCLIENT._auth_info_cached = False LLMCLIENT._llm_api_cache.clear() @@ -458,7 +458,7 @@ def test_chat(monkeypatch, mock_public_model): assert kwargs_passed.get("parent_message_id") == "test-parent-msg-id" -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) def test_reasoning_effort(monkeypatch, mock_public_model): LLMCLIENT._auth_info_cached = False LLMCLIENT._llm_api_cache.clear() @@ -472,7 +472,7 @@ def test_reasoning_effort(monkeypatch, mock_public_model): llm.chat("Hello, how are you?", reasoning_effort="somewhat high") -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) def test_chat_full_response(monkeypatch): LLMCLIENT._auth_info_cached = False LLMCLIENT._llm_api_cache.clear() @@ -497,7 +497,7 @@ def test_chat_full_response(monkeypatch): assert response == mock_response -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) def test_chat_with_tools(monkeypatch, mock_tools): LLMCLIENT._auth_info_cached = False LLMCLIENT._llm_api_cache.clear() @@ -525,7 +525,7 @@ def test_chat_with_tools(monkeypatch, mock_tools): ) -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) def test_metadata(monkeypatch, mock_public_model): LLMCLIENT._auth_info_cached = False LLMCLIENT._llm_api_cache.clear() @@ -569,7 +569,7 @@ def test_metadata(monkeypatch, mock_public_model): assert metadata == metadata2 # should return the same object -@patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) +@patch("lightning_sdk.lightning_cloud.credentials.Auth", new=MagicMock()) def test_chat_backend(monkeypatch, mock_public_model): LLMCLIENT._auth_info_cached = False LLMCLIENT._llm_api_cache.clear() diff --git a/python/tests/core/test_mmt.py b/python/tests/core/test_mmt.py index 21d772073..5bf33cd4b 100644 --- a/python/tests/core/test_mmt.py +++ b/python/tests/core/test_mmt.py @@ -62,7 +62,7 @@ def test_mmt_run_emits_deprecation_warning() -> None: @pytest.mark.parametrize("command", [None, "echo hello"]) @pytest.mark.parametrize("env", [None, {"key": "value"}]) @pytest.mark.parametrize("interruptible", [True, False]) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_submit_mmt_v2_image(internal_studio_init_mocker, machine, command, env, interruptible): teamspace = Teamspace("ts-abc", org="org-abc") job = MMT("test-job", teamspace, _fetch_job=False) @@ -106,7 +106,7 @@ def test_submit_mmt_v2_image(internal_studio_init_mocker, machine, command, env, ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_submit_mmt_threads_placement_group_id(internal_studio_init_mocker): teamspace = Teamspace("ts-abc", org="org-abc") job = MMT("test-job", teamspace, _fetch_job=False) @@ -125,7 +125,7 @@ def test_submit_mmt_threads_placement_group_id(internal_studio_init_mocker): assert submit_mock.call_args.kwargs["placement_group_id"] == "pg-1" -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_submit_mmt_threads_keep_minutes(internal_studio_init_mocker): teamspace = Teamspace("ts-abc", org="org-abc") job = MMT("test-job", teamspace, _fetch_job=False) @@ -144,7 +144,7 @@ def test_submit_mmt_threads_keep_minutes(internal_studio_init_mocker): assert submit_mock.call_args.kwargs["keep_minutes"] == 15 -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_mmt_exposes_keep_minutes(mmt_api_get_job_by_name_mocker, internal_studio_init_mocker): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") job = MMT("test-job", studio.teamspace) @@ -153,7 +153,7 @@ def test_mmt_exposes_keep_minutes(mmt_api_get_job_by_name_mocker, internal_studi assert job.keep_minutes == 15 -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_submit_mmt_threads_max_run_attempts(internal_studio_init_mocker): teamspace = Teamspace("ts-abc", org="org-abc") job = MMT("test-job", teamspace, _fetch_job=False) @@ -172,7 +172,7 @@ def test_submit_mmt_threads_max_run_attempts(internal_studio_init_mocker): assert submit_mock.call_args.kwargs["max_run_attempts"] == 4 -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_mmt_exposes_placement_group_id(mmt_api_get_job_by_name_mocker, internal_studio_init_mocker): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") job = MMT("test-job", studio.teamspace) @@ -190,7 +190,7 @@ def test_mmt_exposes_resource_id(mmt_api_get_job_by_name_mocker, internal_studio assert job.resource_id == "mmt-123" -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_mmt_id_is_none_without_backing_model_or_api_request(internal_studio_init_mocker): teamspace = Teamspace("org-abc/ts-abc") job = MMT("test-job", teamspace, _fetch_job=False) @@ -199,7 +199,7 @@ def test_mmt_id_is_none_without_backing_model_or_api_request(internal_studio_ini assert job.id is None -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_mmt_members_have_stable_rank_identity(mmt_api_get_job_by_name_mocker, internal_studio_init_mocker): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") job = MMT("test-job", studio.teamspace) @@ -230,7 +230,7 @@ def test_mmt_members_have_stable_rank_identity(mmt_api_get_job_by_name_mocker, i @pytest.mark.parametrize("machine", [Machine.L4, Machine.DATA_PREP_MAX]) @pytest.mark.parametrize("env", [None, {"key": "value"}]) @pytest.mark.parametrize("interruptible", [True, False]) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_submit_mmt_v2_studio(internal_studio_init_mocker, machine, env, interruptible): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") @@ -272,7 +272,7 @@ def test_submit_mmt_v2_studio(internal_studio_init_mocker, machine, env, interru ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_mmt_run_arg_validation(internal_studio_init_mocker): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") @@ -290,7 +290,7 @@ def test_mmt_run_arg_validation(internal_studio_init_mocker): MMT.run(None, 5, Machine.CPU) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_mmt_run_entrypoint_validation(internal_studio_init_mocker): """Test entrypoint validation logic for MMT image jobs.""" studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") @@ -352,7 +352,7 @@ def test_mmt_run_entrypoint_validation(internal_studio_init_mocker): ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_submit_mmtv2_error_cases(internal_studio_init_mocker): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") @@ -397,7 +397,7 @@ def test_submit_mmtv2_error_cases(internal_studio_init_mocker): ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_mmt_by_name(internal_studio_init_mocker): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") job = MMT("test-job", studio.teamspace, _fetch_job=False) @@ -410,7 +410,7 @@ def test_get_mmt_by_name(internal_studio_init_mocker): job._job_api.get_job_by_name.assert_called_once_with(name="test-job", teamspace_id="ts-abc001") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_mmt_by_id(internal_studio_init_mocker): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") job = MMT("test-job", studio.teamspace, _fetch_job=False) @@ -426,7 +426,7 @@ def test_get_mmt_by_id(internal_studio_init_mocker): job._job_api.get_job.assert_called_once_with(job_id="test-job-id", teamspace_id="ts-abc001") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_mmt_by_name_first_and_then_by_id(internal_studio_init_mocker): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") job = MMT("test-job", studio.teamspace, _fetch_job=False) @@ -444,7 +444,7 @@ def test_get_mmt_by_name_first_and_then_by_id(internal_studio_init_mocker): job._job_api.get_job.assert_called_once_with(job_id="test-job-id", teamspace_id="ts-abc001") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_mmt_by_name_on_init(mmt_api_get_job_by_name_mocker, internal_studio_init_mocker): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") job = MMT("test-job", studio.teamspace) @@ -466,7 +466,7 @@ def test_get_mmt_by_name_on_init(mmt_api_get_job_by_name_mocker, internal_studio ("MultiMachineJob_STATE_COMPLETED", Status.Completed), ], ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_mmtv2_status(mmt_api_get_job_by_name_mocker, internal_studio_init_mocker, internal_status, external_status): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") @@ -490,7 +490,7 @@ def test_mmtv2_status(mmt_api_get_job_by_name_mocker, internal_studio_init_mocke ("", "unknown", Machine.from_str("unknown")), ], ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_mmtv2_machine( internal_studio_init_mocker, internal_studio_api_mocker_get_machine, @@ -515,7 +515,7 @@ def test_mmtv2_machine( get_job_mock.assert_called_once() -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_mmtv2_stop(mmt_api_get_job_by_name_mocker, internal_studio_init_mocker): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") @@ -556,7 +556,7 @@ def get_job_side_effect(*args, **kwargs): ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_mmtv2_delete(mmt_api_get_job_by_name_mocker, internal_studio_init_mocker): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") @@ -573,7 +573,7 @@ def test_mmtv2_delete(mmt_api_get_job_by_name_mocker, internal_studio_init_mocke delete_job_mock.assert_called_once_with(job_id="test-job-id", teamspace_id="ts-abc001", cloudspace_id=None) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_mmtv2_logs_merges_ranks(mmt_api_get_job_by_name_mocker, internal_studio_init_mocker, monkeypatch): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") job = MMT("test-job", studio.teamspace) @@ -599,7 +599,7 @@ def test_mmtv2_logs_merges_ranks(mmt_api_get_job_by_name_mocker, internal_studio assert stream_mock.call_args.kwargs["follow"] is False -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_mmtv2_logs_follow_and_filters(mmt_api_get_job_by_name_mocker, internal_studio_init_mocker, monkeypatch): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") job = MMT("test-job", studio.teamspace) @@ -617,7 +617,7 @@ def test_mmtv2_logs_follow_and_filters(mmt_api_get_job_by_name_mocker, internal_ assert kwargs["idle_timeout"] is None -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_mmtv2_logs_rank_reads_that_machine(mmt_api_get_job_by_name_mocker, internal_studio_init_mocker, monkeypatch): """The logs API has no rank parameter; a ranked read resolves the machine client-side.""" studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") @@ -639,7 +639,7 @@ def test_mmtv2_logs_rank_reads_that_machine(mmt_api_get_job_by_name_mocker, inte assert kwargs["follow"] is False -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_mmtv2_logs_rank_out_of_range(mmt_api_get_job_by_name_mocker, internal_studio_init_mocker, monkeypatch): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") job = MMT("test-job", studio.teamspace) @@ -653,7 +653,7 @@ def test_mmtv2_logs_rank_out_of_range(mmt_api_get_job_by_name_mocker, internal_s job.logs(rank=1) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_mmtv2_logs_unavailable_while_pending(mmt_api_get_job_by_name_mocker, internal_studio_init_mocker, monkeypatch): studio = Studio(name="st-abc", teamspace="ts-abc", org="org-abc") job = MMT("test-job", studio.teamspace) diff --git a/python/tests/core/test_models.py b/python/tests/core/test_models.py index ee9f0ac86..58013ec1d 100644 --- a/python/tests/core/test_models.py +++ b/python/tests/core/test_models.py @@ -26,7 +26,7 @@ def __init__(self, id: str) -> None: @mock.patch("lightning_sdk.models._get_authed_user") @mock.patch("lightning_sdk.models.UserApi") @mock.patch("lightning_sdk.models.OrgApi") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_teamspace_org_owner( mock_org_api, mock_user_api, mock_get_authed_user, mock_teamspace_api, mock_resolve_org ): @@ -61,7 +61,7 @@ def test_get_teamspace_org_owner( @mock.patch("lightning_sdk.models._get_authed_user") @mock.patch("lightning_sdk.models.UserApi") @mock.patch("lightning_sdk.models.OrgApi") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_teamspace_authed_owner( mock_org_api, mock_user_api, mock_get_authed_user, mock_teamspace_api, mock_authed_user_api ): @@ -93,7 +93,7 @@ def test_get_teamspace_authed_owner( @mock.patch("lightning_sdk.models._get_authed_user") @mock.patch("lightning_sdk.models.UserApi") @mock.patch("lightning_sdk.models.OrgApi") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_teamspace_other_user_owner( mock_org_api, mock_user_api, mock_get_authed_user, mock_teamspace_api, mock_authed_user_api ): @@ -124,7 +124,7 @@ def test_get_teamspace_other_user_owner( @mock.patch("lightning_sdk.models._get_teamspace") @mock.patch("lightning_sdk.models.TeamspaceApi") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_download_model_errors(mock_teamspace_api, mock_get_teamspace): from lightning_sdk.teamspace import Teamspace @@ -152,7 +152,7 @@ def test_download_model_errors(mock_teamspace_api, mock_get_teamspace): @mock.patch("lightning_sdk.api.teamspace_api._download_model_files") @mock.patch("lightning_sdk.teamspace.TeamspaceApi") @mock.patch("lightning_sdk.organization.OrgApi") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_download_model_in_studio_with_org( mock_org_api, mock_teamspace_api, @@ -189,7 +189,7 @@ def test_download_model_in_studio_with_org( @mock.patch("lightning_sdk.api.teamspace_api._download_model_files") @mock.patch("lightning_sdk.teamspace.TeamspaceApi") @mock.patch("lightning_sdk.user.UserApi") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_download_model_in_studio_with_user( mock_user_api, mock_teamspace_api, @@ -225,7 +225,7 @@ def test_download_model_in_studio_with_user( @mock.patch("lightning_sdk.models._get_teamspace") @mock.patch("lightning_sdk.teamspace.TeamspaceApi") @mock.patch("lightning_sdk.organization.OrgApi") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_upload_model_in_studio_with_org( mock_org_api, mock_teamspace_api, mock_get_teamspace, mock_parse_org_teamspace_model_version ): @@ -253,7 +253,7 @@ def test_upload_model_in_studio_with_org( @mock.patch("lightning_sdk.models._get_teamspace") @mock.patch("lightning_sdk.teamspace.TeamspaceApi") @mock.patch("lightning_sdk.organization.OrgApi") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_upload_model_in_studio_with_org_and_experiment( mock_org_api, mock_teamspace_api, mock_get_teamspace, mock_parse_org_teamspace_model_version ): @@ -283,7 +283,7 @@ def test_upload_model_in_studio_with_org_and_experiment( @mock.patch("lightning_sdk.models._get_teamspace") @mock.patch("lightning_sdk.teamspace.TeamspaceApi") @mock.patch("lightning_sdk.user.UserApi") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_upload_model_in_studio_with_user( mock_user_api, mock_teamspace_api, mock_get_teamspace, mock_parse_org_teamspace_model_version ): @@ -311,7 +311,7 @@ def test_upload_model_in_studio_with_user( @mock.patch("lightning_sdk.models._get_teamspace") @mock.patch("lightning_sdk.teamspace.TeamspaceApi") @mock.patch("lightning_sdk.organization.OrgApi") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_delete_model_in_studio_with_org( mock_org_api, mock_teamspace_api, mock_get_teamspace, mock_parse_org_teamspace_model_version ): @@ -332,7 +332,7 @@ def test_delete_model_in_studio_with_org( @mock.patch("lightning_sdk.models._get_teamspace") @mock.patch("lightning_sdk.teamspace.TeamspaceApi") @mock.patch("lightning_sdk.user.UserApi") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_delete_model_in_studio_with_user( mock_user_api, mock_teamspace_api, mock_get_teamspace, mock_parse_org_teamspace_model_version ): @@ -352,7 +352,7 @@ def test_delete_model_in_studio_with_user( @mock.patch("lightning_sdk.models._get_teamspace") @mock.patch("lightning_sdk.teamspace.TeamspaceApi") @mock.patch("lightning_sdk.organization.OrgApi") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_delete_model_version_in_studio_with_org( mock_org_api, mock_teamspace_api, mock_get_teamspace, mock_parse_org_teamspace_model_version ): @@ -373,7 +373,7 @@ def test_delete_model_version_in_studio_with_org( @mock.patch("lightning_sdk.models._get_teamspace") @mock.patch("lightning_sdk.teamspace.TeamspaceApi") @mock.patch("lightning_sdk.user.UserApi") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_delete_model_version_in_studio_with_user( mock_user_api, mock_teamspace_api, mock_get_teamspace, mock_parse_org_teamspace_model_version ): @@ -393,7 +393,7 @@ def test_delete_model_version_in_studio_with_user( @mock.patch("lightning_sdk.models._get_teamspace") @mock.patch("lightning_sdk.teamspace.TeamspaceApi") @mock.patch("lightning_sdk.organization.OrgApi") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_list_model_versions_in_studio_with_org( mock_org_api, mock_teamspace_api, mock_get_teamspace, mock_parse_org_teamspace_model_version ): @@ -414,7 +414,7 @@ def test_list_model_versions_in_studio_with_org( @mock.patch("lightning_sdk.models._get_teamspace") @mock.patch("lightning_sdk.teamspace.TeamspaceApi") @mock.patch("lightning_sdk.user.UserApi") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_list_model_versions_in_studio_with_user( mock_user_api, mock_teamspace_api, mock_get_teamspace, mock_parse_org_teamspace_model_version ): @@ -431,7 +431,7 @@ def test_list_model_versions_in_studio_with_user( @mock.patch("lightning_sdk.models._get_teamspace") @mock.patch("lightning_sdk.api.teamspace_api._download_model_files") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_download_dir_home_path_shorthand(mock_download_model_files, mock_get_teamspace): from lightning_sdk.teamspace import Teamspace diff --git a/python/tests/core/test_org.py b/python/tests/core/test_org.py index 7e06deecb..2efc3862a 100644 --- a/python/tests/core/test_org.py +++ b/python/tests/core/test_org.py @@ -5,7 +5,7 @@ from lightning_sdk.teamspace import Teamspace -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_org_init_from_name(internal_get_org_api_mocker): org = Organization("my-org-name") @@ -14,7 +14,7 @@ def test_org_init_from_name(internal_get_org_api_mocker): @mock.patch.dict(os.environ, {"LIGHTNING_ORG": "my-org-name"}) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_org_init_from_env_var(internal_get_org_api_mocker): org = Organization() @@ -22,7 +22,7 @@ def test_org_init_from_env_var(internal_get_org_api_mocker): assert org.id == "my-org-name" -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_org_teamspaces(internal_get_org_api_mocker, internal_teamspace_api_list_mocker): org = Organization("org-abc") @@ -34,7 +34,7 @@ def test_org_teamspaces(internal_get_org_api_mocker, internal_teamspace_api_list assert ts.owner == org -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_equality(internal_get_org_api_mocker): assert Organization("my-orgname") == Organization("my-orgname") assert Organization("my-orgname") != Organization("your-orgname") @@ -44,24 +44,24 @@ class SubOrg(Organization): pass -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_inequality_org_subclass(internal_get_org_api_mocker): assert Organization("my-orgname") != SubOrg("my-orgname") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_repr(internal_get_org_api_mocker): org = Organization("my-org-name") assert repr(org) == "Organization(name=my-org-name)" -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_str(internal_get_org_api_mocker): org = Organization("my-org-name") assert str(org) == "Organization(name=my-org-name)" -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_org_default_cloud_account(internal_get_org_api_mocker): org = Organization("my-org-name") assert org.default_cloud_account == "my-preferred-cluster" @@ -73,7 +73,7 @@ def test_org_default_cloud_account(internal_get_org_api_mocker): @mock.patch("lightning_sdk.teamspace.Teamspace.__init__", return_value=None) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_org_create_teamspace(mock_teamspace_init, internal_get_org_api_mocker): org = Organization("my-org-name") diff --git a/python/tests/core/test_permissions_integration.py b/python/tests/core/test_permissions_integration.py index 2c4227bfb..3152d13b6 100644 --- a/python/tests/core/test_permissions_integration.py +++ b/python/tests/core/test_permissions_integration.py @@ -182,7 +182,7 @@ def test_teamspace_multi_machine_jobs_property_succeeds_when_enabled(mock_teamsp # Studio class permission tests -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) @mock.patch("lightning_sdk.user.UserApi") @mock.patch("lightning_sdk.studio._resolve_teamspace") @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi") @@ -298,7 +298,7 @@ def test_pipeline_init_raises_error_when_pipelines_disabled(mock_teamspace_api, # Deployment class permission tests -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) @mock.patch("lightning_sdk.lightning_cloud.login.Auth") @mock.patch("lightning_sdk.user.UserApi") @mock.patch("lightning_sdk.deployment._resolve_teamspace") @@ -440,7 +440,7 @@ def test_litcontainer_upload_raises_error_when_containers_disabled( @mock.patch("lightning_sdk.api.utils.LightningClient") @mock.patch("lightning_sdk.api.lit_container_api.docker.from_env") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) @mock.patch("lightning_sdk.user.UserApi") @mock.patch("lightning_sdk.lit_container._resolve_teamspace") @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi") diff --git a/python/tests/core/test_sandbox.py b/python/tests/core/test_sandbox.py index 6501ecce5..6441390ff 100644 --- a/python/tests/core/test_sandbox.py +++ b/python/tests/core/test_sandbox.py @@ -28,7 +28,7 @@ def _v1(**kwargs) -> V1Sandbox: return V1Sandbox(**defaults) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_resolve_sandbox_api_rejects_config_and_api_together(): cfg = SandboxConfig(api_key="k", base_url="https://x") fake_api = mock.MagicMock() @@ -36,7 +36,7 @@ def test_resolve_sandbox_api_rejects_config_and_api_together(): _resolve_sandbox_api(sandbox_api=fake_api, config=cfg) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_create_sandbox_waits_until_running(): api = SandboxApi({"api_key": "unit-key", "base_url": "https://unit.test", "organization_id": "org-1"}) mock_sb = mock.MagicMock() @@ -56,7 +56,7 @@ def test_create_sandbox_waits_until_running(): mock_sb.sandboxes_service_get_sandbox.assert_called_once() -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_create_sandbox_raises_on_terminal_status(): api = SandboxApi({"api_key": "unit-key", "base_url": "https://unit.test", "organization_id": "org-1"}) mock_sb = mock.MagicMock() @@ -68,7 +68,7 @@ def test_create_sandbox_raises_on_terminal_status(): create_sandbox(name="n", sandbox_api=api) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_create_sandbox_omits_project_id_without_api_response(): api = SandboxApi({"api_key": "unit-key", "base_url": "https://unit.test", "organization_id": "org-1"}) mock_sb = mock.MagicMock() @@ -85,7 +85,7 @@ def test_create_sandbox_omits_project_id_without_api_response(): assert body.project_id is None -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_create_sandbox_restore_omits_project_id_prefetch(): api = SandboxApi({"api_key": "unit-key", "base_url": "https://unit.test", "organization_id": "org-1"}) mock_sb = mock.MagicMock() @@ -101,7 +101,7 @@ def test_create_sandbox_restore_omits_project_id_prefetch(): assert body.snapshot_id == "snap-1" -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_create_sandbox_sends_explicit_project_id(): api = SandboxApi({"api_key": "unit-key", "base_url": "https://unit.test", "organization_id": "org-1"}) mock_sb = mock.MagicMock() @@ -114,7 +114,7 @@ def test_create_sandbox_sends_explicit_project_id(): assert body.project_id == "proj-1" -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_create_sandbox_sends_image_and_secret_ref(): api = SandboxApi({"api_key": "unit-key", "base_url": "https://unit.test", "organization_id": "org-1"}) mock_sb = mock.MagicMock() @@ -134,21 +134,21 @@ def test_create_sandbox_sends_image_and_secret_ref(): assert body.runtime is None -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_create_sandbox_rejects_image_with_runtime(): api = SandboxApi({"api_key": "unit-key", "base_url": "https://unit.test", "organization_id": "org-1"}) with pytest.raises(ValueError, match="mutually exclusive"): create_sandbox(name="n", sandbox_api=api, runtime="python", image="ghcr.io/org/img:latest") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_create_sandbox_rejects_secret_ref_without_image(): api = SandboxApi({"api_key": "unit-key", "base_url": "https://unit.test", "organization_id": "org-1"}) with pytest.raises(ValueError, match="only valid together with 'image'"): create_sandbox(name="n", sandbox_api=api, image_secret_ref="my-secret") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_create_sandbox_docker_maps_runtime_to_variant(): api = SandboxApi({"api_key": "unit-key", "base_url": "https://unit.test", "organization_id": "org-1"}) mock_sb = mock.MagicMock() @@ -161,7 +161,7 @@ def test_create_sandbox_docker_maps_runtime_to_variant(): assert body.runtime == "node24-docker" -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_create_sandbox_docker_defaults_runtime(): api = SandboxApi({"api_key": "unit-key", "base_url": "https://unit.test", "organization_id": "org-1"}) mock_sb = mock.MagicMock() @@ -174,7 +174,7 @@ def test_create_sandbox_docker_defaults_runtime(): assert body.runtime == "node24-docker" -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_create_sandbox_docker_does_not_double_suffix(): api = SandboxApi({"api_key": "unit-key", "base_url": "https://unit.test", "organization_id": "org-1"}) mock_sb = mock.MagicMock() @@ -187,21 +187,21 @@ def test_create_sandbox_docker_does_not_double_suffix(): assert body.runtime == "python313-docker" -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_create_sandbox_docker_rejects_image(): api = SandboxApi({"api_key": "unit-key", "base_url": "https://unit.test", "organization_id": "org-1"}) with pytest.raises(ValueError, match="cannot be combined with 'image'"): create_sandbox(name="n", sandbox_api=api, image="ghcr.io/org/img:latest", docker=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_create_sandbox_docker_rejects_snapshot_id(): api = SandboxApi({"api_key": "unit-key", "base_url": "https://unit.test", "organization_id": "org-1"}) with pytest.raises(ValueError, match="cannot be combined with 'snapshot_id'"): create_sandbox(name="n", sandbox_api=api, snapshot_id="snap-1", docker=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_sandbox_instance_snapshot_omits_project_id_for_org_scoped_sandbox(): from lightning_sdk.lightning_cloud.openapi import V1SandboxSnapshot @@ -217,7 +217,7 @@ def test_sandbox_instance_snapshot_omits_project_id_for_org_scoped_sandbox(): assert body.project_id is None -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_sandbox_instance_snapshot_maps_project_id_required_error(): api = SandboxApi({"api_key": "unit-key", "base_url": "https://unit.test"}) sb_svc = mock.MagicMock() @@ -233,7 +233,7 @@ def test_sandbox_instance_snapshot_maps_project_id_required_error(): sb.snapshot(wait=False) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_sandbox_instance_snapshot_sends_project_id_from_sandbox_row(): from lightning_sdk.lightning_cloud.openapi import V1SandboxSnapshot @@ -251,7 +251,7 @@ def test_sandbox_instance_snapshot_sends_project_id_from_sandbox_row(): assert info.id == "snap-1" -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_sandbox_instance_stop_sends_project_id_from_sandbox_row(): api = SandboxApi({"api_key": "unit-key", "base_url": "https://unit.test"}) sb_svc = mock.MagicMock() @@ -266,7 +266,7 @@ def test_sandbox_instance_stop_sends_project_id_from_sandbox_row(): assert body.project_id == "proj-on-sandbox" -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_sandbox_client_snapshot_helpers(): from lightning_sdk.lightning_cloud.openapi import V1ListSandboxSnapshotsResponse, V1SandboxSnapshot @@ -286,7 +286,7 @@ def test_sandbox_client_snapshot_helpers(): sb_svc.sandboxes_service_delete_sandbox_snapshot.assert_called_once_with("snap-1") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_sandbox_instance_run_command_string_and_args(): mock_api = mock.MagicMock() mock_api.run_command.return_value = CommandResult(cmd_id="c1", output="hi", exit_code=0) @@ -300,7 +300,7 @@ def test_sandbox_instance_run_command_string_and_args(): ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_sandbox_instance_run_command_auto_splits_simple_command(): mock_api = mock.MagicMock() mock_api.run_command.return_value = CommandResult(cmd_id="", output="", exit_code=0) @@ -314,7 +314,7 @@ def test_sandbox_instance_run_command_auto_splits_simple_command(): assert kwargs["args"] == ["-a"] -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_sandbox_instance_run_command_run_command_opts(): mock_api = mock.MagicMock() mock_api.run_command.return_value = CommandResult(cmd_id="d1", output="", exit_code=0) @@ -345,7 +345,7 @@ def test_sandbox_instance_run_command_run_command_opts(): ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_sandbox_instance_run_command_non_detached_returns_finished_command(): mock_api = mock.MagicMock() mock_api.run_command.return_value = CommandResult(cmd_id="c1", output="hi", exit_code=0) @@ -360,7 +360,7 @@ def test_sandbox_instance_run_command_non_detached_returns_finished_command(): assert cmd.output == "hi" -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_command_wait_polls_until_done_for_detached(): mock_api = mock.MagicMock() mock_api.run_command.return_value = CommandResult(cmd_id="c-bg", output="", exit_code=0) @@ -382,7 +382,7 @@ def test_command_wait_polls_until_done_for_detached(): assert cmd.output == "done\n" -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_command_wait_returns_immediately_when_already_finished(): mock_api = mock.MagicMock() mock_api.run_command.return_value = CommandResult(cmd_id="c1", output="hi", exit_code=0) @@ -395,7 +395,7 @@ def test_command_wait_returns_immediately_when_already_finished(): mock_api.get_command.assert_not_called() -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_command_kill_calls_kill_command(): mock_api = mock.MagicMock() mock_api.run_command.return_value = CommandResult(cmd_id="c-k", output="", exit_code=0) @@ -407,7 +407,7 @@ def test_command_kill_calls_kill_command(): mock_api.kill_command.assert_called_once_with("sb-k", "c-k", "org-k") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_sandbox_instance_wait_for_command_polls_until_done(): mock_api = mock.MagicMock() mock_api.get_command.side_effect = [ @@ -425,7 +425,7 @@ def test_sandbox_instance_wait_for_command_polls_until_done(): mock_api.get_command.assert_called_with("sb-w", "c-1", "org-w") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_sandbox_instance_wait_for_command_times_out(): mock_api = mock.MagicMock() mock_api.get_command.return_value = CommandStatus(output="", exit_code=0, running=True) @@ -440,7 +440,7 @@ def test_sandbox_instance_wait_for_command_times_out(): sb.wait_for_command("c-1", timeout=1.0, poll_interval=0.0) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_sandbox_instance_read_file_returns_none_on_404(): mock_api = mock.MagicMock() sb_svc = mock.MagicMock() @@ -452,7 +452,7 @@ def test_sandbox_instance_read_file_returns_none_on_404(): assert sb.read_file("/missing") is None -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_sandbox_instance_stop_swallows_404(): mock_api = mock.MagicMock() sb_svc = mock.MagicMock() @@ -463,7 +463,7 @@ def test_sandbox_instance_stop_swallows_404(): sb.stop() -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_sandbox_instance_list_builds_result(): sdk = Sandbox(SandboxConfig(api_key="unit-key", base_url="https://unit.test")) sb_svc = mock.MagicMock() @@ -488,7 +488,7 @@ def test_sandbox_instance_list_builds_result(): sb_svc.sandboxes_service_list_sandboxes.assert_called_once_with(page_token="pt", limit=10) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_list_sandboxes_omits_organization_id_when_not_configured(): sdk = Sandbox(SandboxConfig(api_key="unit-key", base_url="https://unit.test")) sb_svc = mock.MagicMock() @@ -500,7 +500,7 @@ def test_list_sandboxes_omits_organization_id_when_not_configured(): sb_svc.sandboxes_service_list_sandboxes.assert_called_once_with() -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_sandbox_instance_write_files_uses_write_file(): mock_api = mock.MagicMock() sb = SandboxInstance(_v1(id="sb-6"), sandbox_api=mock_api) @@ -518,7 +518,7 @@ def test_sandbox_instance_write_files_uses_write_file(): sb.write_file.assert_any_call("/b", "B") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_sandbox_instance_port_urls_and_get_port_url(): url = "https://8080-sb-7-s.cloudspaces.litng.ai" sb = SandboxInstance(_v1(id="sb-7", ports=["8080"], port_urls={"8080": url})) @@ -528,7 +528,7 @@ def test_sandbox_instance_port_urls_and_get_port_url(): assert sb.get_port_url("8080") == url -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_sandbox_instance_get_port_url_raises_for_unexposed_port(): sb = SandboxInstance(_v1(id="sb-8")) @@ -537,7 +537,7 @@ def test_sandbox_instance_get_port_url_raises_for_unexposed_port(): sb.get_port_url(9000) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_sandbox_entry_create_forwards_to_create_sandbox(): mock_inst = mock.MagicMock() with mock.patch("lightning_sdk.sandbox.sandbox.create_sandbox", return_value=mock_inst) as m_create: @@ -550,7 +550,7 @@ def test_sandbox_entry_create_forwards_to_create_sandbox(): assert m_create.call_args.kwargs["sandbox_api"] is sdk.api -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_sandbox_create_classmethod_forwards(): mock_inst = mock.MagicMock() cfg = SandboxConfig(api_key="k", base_url="https://unit.test") @@ -567,7 +567,7 @@ def test_sandbox_create_classmethod_forwards(): assert received.config_get("base_url") == "https://unit.test" -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_sandbox_create_without_config_uses_configured_globals(): """Sandbox.configure() defaults must reach Sandbox.create() when no config= is passed.""" import lightning_sdk.sandbox.base as base @@ -589,7 +589,7 @@ def test_sandbox_create_without_config_uses_configured_globals(): base._default_api().reset() -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_sandbox_client_without_config_uses_configured_globals(): """Sandbox() (the client) must honor Sandbox.configure() defaults, like Sandbox.create().""" import lightning_sdk.sandbox.base as base @@ -610,7 +610,7 @@ def test_sandbox_client_without_config_uses_configured_globals(): base._default_api().reset() -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_sandbox_client_with_config_uses_isolated_client(): """An explicit config gets its own client, not the shared global one.""" import lightning_sdk.sandbox.base as base @@ -620,7 +620,7 @@ def test_sandbox_client_with_config_uses_isolated_client(): assert client.api.config_get("api_key") == "explicit-key" -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_sandbox_create_resolves_teamspace_to_project_id(): mock_inst = mock.MagicMock() cfg = SandboxConfig(api_key="k", base_url="https://unit.test") @@ -635,7 +635,7 @@ def test_sandbox_create_resolves_teamspace_to_project_id(): assert m_create.call_args.kwargs["project_id"] == "proj-1" -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_resolve_teamspace_id_accepts_owner_teamspace_name(): resolved = mock.MagicMock() resolved.id = "proj-1" @@ -648,7 +648,7 @@ def test_resolve_teamspace_id_accepts_owner_teamspace_name(): def test_resolve_teamspace_id_falls_back_to_memberships(mocker): from lightning_sdk.lightning_cloud.openapi import V1ListMembershipsResponse, V1Membership - mocker.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) + mocker.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) m_resolve = mocker.patch("lightning_sdk.utils.resolve._resolve_teamspace", side_effect=Exception("fail")) membership = V1Membership( @@ -671,7 +671,7 @@ def test_resolve_teamspace_id_falls_back_to_memberships(mocker): m_resolve.assert_any_call("teamspace", org=None, user="owner") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_sandbox_entry_get_and_list_use_sdk_api(): sdk = Sandbox(SandboxConfig(api_key="k", base_url="https://unit.test")) with mock.patch.object(sdk.api, "get_sandbox", return_value=_v1(id="sb-x")) as m_get, mock.patch.object( @@ -687,7 +687,7 @@ def test_sandbox_entry_get_and_list_use_sdk_api(): m_list.assert_called_once_with(page_token=None, limit=3, project_id=None) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_sandbox_list_resolves_teamspace_to_project_id(): sdk = Sandbox(SandboxConfig(api_key="k", base_url="https://unit.test")) with mock.patch.object( @@ -701,7 +701,7 @@ def test_sandbox_list_resolves_teamspace_to_project_id(): m_list.assert_called_once_with(page_token=None, limit=3, project_id="proj-1") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_configure_updates_globals_and_resets_client(): import lightning_sdk.sandbox.base as base diff --git a/python/tests/core/test_sandbox_config.py b/python/tests/core/test_sandbox_config.py index 49cd413aa..886478e8e 100644 --- a/python/tests/core/test_sandbox_config.py +++ b/python/tests/core/test_sandbox_config.py @@ -7,26 +7,26 @@ from lightning_sdk.sandbox.config import SandboxConfig, reject_legacy_org_id_config -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_sandbox_config_to_api_dict(): cfg = SandboxConfig(api_key="k", base_url="https://x") assert cfg.to_api_dict() == {"api_key": "k", "base_url": "https://x"} -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_sandbox_config_rejects_organization_id_kwarg(): with pytest.raises(TypeError): SandboxConfig(api_key="k", organization_id="org-uuid") # type: ignore[call-arg] -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_from_env_rejects_lightning_org_id(monkeypatch): monkeypatch.setenv("LIGHTNING_ORG_ID", "org-1") with pytest.raises(ValueError, match="LIGHTNING_ORG_ID is no longer supported"): SandboxConfig.from_env() -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_configure_rejects_organization_id_kwarg(): from lightning_sdk.sandbox.base import configure @@ -34,7 +34,7 @@ def test_configure_rejects_organization_id_kwarg(): configure(organization_id="org-1") # type: ignore[call-arg] -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_sandbox_config_api_without_api_key_uses_lightning_auth(): with mock.patch("lightning_sdk.api.sandbox_api.Auth") as auth_cls: auth_cls.return_value.authenticate.return_value = "Basic auth" @@ -47,7 +47,7 @@ def test_sandbox_config_api_without_api_key_uses_lightning_auth(): auth_cls.return_value.authenticate.assert_called_once() -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_sandbox_config_api_passes_through_config(): api = SandboxConfig(api_key="k", base_url="https://x").api() assert api.config_get("api_key") == "k" @@ -55,12 +55,12 @@ def test_sandbox_config_api_passes_through_config(): assert api.config_get("organization_id") is None -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_reject_legacy_org_id_config_is_noop_without_env(): reject_legacy_org_id_config() -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_importing_sandbox_module_tolerates_legacy_org_id_env(monkeypatch): """Importing must not read the environment. LIGHTNING_ORG_ID is rejected, so reading it at import time made every `lightning` CLI command fail in an environment that @@ -78,7 +78,7 @@ def test_importing_sandbox_module_tolerates_legacy_org_id_env(monkeypatch): base._default_api() -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_global_config_loads_env_on_first_use(monkeypatch): monkeypatch.setenv("LIGHTNING_SANDBOX_API_KEY", "key-from-env") monkeypatch.setenv("LIGHTNING_CLOUD_URL", "https://env.unit") diff --git a/python/tests/core/test_teamspace.py b/python/tests/core/test_teamspace.py index e1c33caa8..4fdf36b92 100644 --- a/python/tests/core/test_teamspace.py +++ b/python/tests/core/test_teamspace.py @@ -52,7 +52,7 @@ def __init__(self, id: str) -> None: @pytest.mark.parametrize("org", ["org-abc", None, -1]) @mock.patch.dict(os.environ, clear=True) @mock.patch("lightning_sdk.utils.config._DEFAULT_CONFIG_FILE_PATH", "/nonexistent/config.yaml") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_teamspace_init( internal_teamspace_api_list_mocker, internal_user_api_mocker, internal_get_org_api_mocker, user, org ): @@ -82,7 +82,7 @@ def test_teamspace_init( @pytest.mark.parametrize("user", ["user-abc", None, -1]) @pytest.mark.parametrize("org", ["org-abc", None, -1]) @mock.patch("lightning_sdk.utils.config._DEFAULT_CONFIG_FILE_PATH", "/nonexistent/config.yaml") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_teamspace_init_env( internal_teamspace_api_list_mocker, internal_user_api_mocker, internal_get_org_api_mocker, user, org ): @@ -117,7 +117,7 @@ def test_teamspace_init_env( @mock.patch.dict(os.environ, clear=True) @mock.patch("lightning_sdk.lightning_cloud.login.Auth.authenticate", return_value="my-auth-header") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) @mock.patch( "lightning_sdk.api.teamspace_api.AuthApi", new=mock.MagicMock( @@ -141,7 +141,7 @@ def test_teamspace_list_clusters_studios_user( @mock.patch.dict(os.environ, clear=True) @mock.patch("lightning_sdk.lightning_cloud.login.Auth.authenticate", return_value="my-auth-header") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) @mock.patch( "lightning_sdk.api.teamspace_api.AuthApi", new=mock.MagicMock( @@ -171,7 +171,7 @@ def test_teamspace_list_clusters_studios_org( ], ) @mock.patch("lightning_sdk.lightning_cloud.login.Auth.authenticate", return_value="my-auth-header") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_repr( _mock_authenticate, internal_teamspace_api_list_mocker, @@ -192,7 +192,7 @@ def test_repr( ], ) @mock.patch("lightning_sdk.lightning_cloud.login.Auth.authenticate", return_value="my-auth-header") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_str( _mock_authenticate, internal_teamspace_api_list_mocker, @@ -205,7 +205,7 @@ def test_str( assert str(ts) == expected -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_teamspace_error_user_and_org(): with pytest.raises( ValueError, match="User and org are mutually exclusive. Please only specify the one who owns the teamspace." @@ -216,7 +216,7 @@ def test_teamspace_error_user_and_org(): @mock.patch("lightning_sdk.teamspace.TeamspaceApi") @mock.patch("lightning_sdk.teamspace._resolve_user") @mock.patch("lightning_sdk.teamspace._resolve_org") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_teamspace_slug_resolves_as_org(mock_resolve_org, mock_resolve_user, mock_teamspace_api): mock_org = mock_resolve_org.return_value mock_org.id = "org-1" @@ -233,7 +233,7 @@ def test_teamspace_slug_resolves_as_org(mock_resolve_org, mock_resolve_user, moc @mock.patch("lightning_sdk.teamspace.TeamspaceApi") @mock.patch("lightning_sdk.teamspace._resolve_user") @mock.patch("lightning_sdk.teamspace._resolve_org") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_teamspace_slug_falls_back_to_user_when_owner_is_not_an_org( mock_resolve_org, mock_resolve_user, mock_teamspace_api ): @@ -248,13 +248,13 @@ def test_teamspace_slug_falls_back_to_user_when_owner_is_not_an_org( assert ts.owner is mock_user -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_teamspace_slug_conflicts_with_explicit_org_raises(): with pytest.raises(ValueError, match="specified both as part of 'name'"): Teamspace(name="my-org/my-ts", org="another-org") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_teamspace_slug_conflicts_with_explicit_user_raises(): with pytest.raises(ValueError, match="specified both as part of 'name'"): Teamspace(name="my-org/my-ts", user="another-user") @@ -263,7 +263,7 @@ def test_teamspace_slug_conflicts_with_explicit_user_raises(): @mock.patch("lightning_sdk.teamspace.TeamspaceApi") @mock.patch("lightning_sdk.teamspace._resolve_user") @mock.patch("lightning_sdk.teamspace._resolve_org") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_teamspace_org_kwarg_emits_deprecation_warning(mock_resolve_org, mock_resolve_user, mock_teamspace_api): mock_resolve_user.return_value = None mock_teamspace_api().get_teamspace.return_value = mock.MagicMock(name="my-ts", id="ts-1") @@ -274,7 +274,7 @@ def test_teamspace_org_kwarg_emits_deprecation_warning(mock_resolve_org, mock_re @mock.patch("lightning_sdk.teamspace.TeamspaceApi") @mock.patch("lightning_sdk.teamspace._resolve_user") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_teamspace_user_kwarg_emits_deprecation_warning(mock_resolve_user, mock_teamspace_api): mock_teamspace_api().get_teamspace.return_value = mock.MagicMock(name="my-ts", id="ts-1") @@ -285,7 +285,7 @@ def test_teamspace_user_kwarg_emits_deprecation_warning(mock_resolve_user, mock_ @mock.patch("lightning_sdk.teamspace.TeamspaceApi") @mock.patch("lightning_sdk.teamspace._resolve_user") @mock.patch("lightning_sdk.teamspace._resolve_org") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_teamspace_slug_without_org_or_user_does_not_warn(mock_resolve_org, mock_resolve_user, mock_teamspace_api): mock_org = mock_resolve_org.return_value mock_org.id = "org-1" @@ -297,7 +297,7 @@ def test_teamspace_slug_without_org_or_user_does_not_warn(mock_resolve_org, mock @mock.patch("lightning_sdk.lightning_cloud.login.Auth.authenticate", return_value="my-auth-header") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_create_agent( _mock_authenticate, internal_teamspace_api_list_mocker, @@ -316,7 +316,7 @@ def test_create_agent( @mock.patch.dict(os.environ, {"LIGHTNING_CLUSTER_ID": "test-cluster-id"}) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_upload_model_single_file( internal_teamspace_api_list_mocker, internal_user_api_mocker, @@ -371,7 +371,7 @@ def test_upload_model_single_file( @mock.patch.dict(os.environ, {"LIGHTNING_CLUSTER_ID": "test-cluster-id"}) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_upload_file_normalizes_remote_path_to_posix( internal_teamspace_api_list_mocker, internal_user_api_mocker, @@ -399,7 +399,7 @@ def test_upload_file_normalizes_remote_path_to_posix( @mock.patch.dict(os.environ, {"LIGHTNING_CLUSTER_ID": "test-cluster-id"}) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_upload_model_single_file_experiment( internal_teamspace_api_list_mocker, internal_user_api_mocker, @@ -458,7 +458,7 @@ def test_upload_model_single_file_experiment( @mock.patch.dict(os.environ, {"LIGHTNING_CLUSTER_ID": "test-cluster-id"}) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_upload_model_multiple_files( internal_teamspace_api_list_mocker, internal_user_api_mocker, @@ -516,7 +516,7 @@ def test_upload_model_multiple_files( @mock.patch.dict(os.environ, {"LIGHTNING_CLUSTER_ID": "test-cluster-id"}) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_upload_model_auto_resumes_incomplete_upload( internal_teamspace_api_list_mocker, internal_user_api_mocker, @@ -560,7 +560,7 @@ def test_upload_model_auto_resumes_incomplete_upload( @mock.patch.dict(os.environ, {"LIGHTNING_CLUSTER_ID": "test-cluster-id"}) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_upload_model_auto_skip_completed( internal_teamspace_api_list_mocker, internal_user_api_mocker, @@ -605,7 +605,7 @@ def test_upload_model_auto_skip_completed( @pytest.mark.parametrize("folder", ["download_dir", None]) @mock.patch("lightning_sdk.api.teamspace_api._download_model_files") @mock.patch("lightning_sdk.api.teamspace_api._get_model_version") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_download_model( get_model_version_mock, download_model_files_mock, @@ -646,7 +646,7 @@ def test_download_model( @mock.patch("lightning_sdk.api.teamspace_api._download_model_files") @mock.patch("lightning_sdk.api.teamspace_api._get_model_version") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_download_model_version( get_model_version_mock, download_model_files_mock, @@ -678,7 +678,7 @@ def test_download_model_version( @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.list_mmts") @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.list_jobs") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_list_jobs( list_jobs_mock, list_mmts_mock, @@ -712,7 +712,7 @@ def test_list_jobs( @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.list_mmts") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_list_mmts( list_mmts_mock, internal_get_org_api_mocker, @@ -740,7 +740,7 @@ def test_list_mmts( @mock.patch("lightning_sdk.teamspace.TeamspaceApi") @mock.patch("lightning_sdk.teamspace._resolve_user") @mock.patch("lightning_sdk.teamspace._resolve_org") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_list_machines(mock_resolve_org, mock_resolve_user, mock_teamspace_api): mock_teamspace_api().list_machines.return_value = [ V1ClusterAccelerator( @@ -769,7 +769,7 @@ def test_list_machines(mock_resolve_org, mock_resolve_user, mock_teamspace_api): @mock.patch("lightning_sdk.teamspace.TeamspaceApi") @mock.patch("lightning_sdk.teamspace._resolve_user") @mock.patch("lightning_sdk.teamspace._resolve_org") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_list_machines_global_cloud_account( mock_resolve_org, mock_resolve_user, mock_teamspace_api, mock_cloud_account_api ): @@ -807,7 +807,7 @@ def test_list_machines_global_cloud_account( @mock.patch("lightning_sdk.teamspace.TeamspaceApi") @mock.patch("lightning_sdk.teamspace._resolve_user") @mock.patch("lightning_sdk.teamspace._resolve_org") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_list_machines_env_cloud_account(mock_resolve_org, mock_resolve_user, mock_teamspace_api): mock_teamspace_api().list_machines.return_value = [ V1ClusterAccelerator( @@ -836,7 +836,7 @@ def test_list_machines_env_cloud_account(mock_resolve_org, mock_resolve_user, mo @mock.patch("lightning_sdk.teamspace.CloudAccountApi") @mock.patch("lightning_sdk.teamspace._resolve_user") @mock.patch("lightning_sdk.teamspace._resolve_org") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_list_machines_with_machine_filter( mock_resolve_org, mock_resolve_user, mock_cloud_account_api, mock_teamspace_api ): @@ -871,7 +871,7 @@ def test_list_machines_with_machine_filter( @mock.patch("lightning_sdk.teamspace.CloudAccountApi") @mock.patch("lightning_sdk.teamspace._resolve_user") @mock.patch("lightning_sdk.teamspace._resolve_org") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_list_machines_with_invalid_machine( mock_resolve_org, mock_resolve_user, mock_cloud_account_api, mock_teamspace_api ): @@ -907,7 +907,7 @@ def test_list_machines_with_invalid_machine( @mock.patch("lightning_sdk.teamspace.TeamspaceApi") @mock.patch("lightning_sdk.teamspace._resolve_user") @mock.patch("lightning_sdk.teamspace._resolve_org") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_teamspace_default_cloud_account_resolution( mock_resolve_org, mock_resolve_user, @@ -932,7 +932,7 @@ def test_teamspace_default_cloud_account_resolution( assert teamspace.default_cloud_account == expected_result -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_teamspace_secrets_property( internal_teamspace_api_list_mocker, internal_user_api_mocker, @@ -948,7 +948,7 @@ def test_teamspace_secrets_property( mock_get.assert_called_once_with("ts-abc002") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_teamspace_set_secret( internal_teamspace_api_list_mocker, internal_user_api_mocker, @@ -962,7 +962,7 @@ def test_teamspace_set_secret( @pytest.mark.parametrize("secret_type", [SecretType.HF_TOKEN, "hf_token"]) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_teamspace_set_secret_hf_token( internal_teamspace_api_list_mocker, internal_user_api_mocker, @@ -976,7 +976,7 @@ def test_teamspace_set_secret_hf_token( mock_set.assert_called_once_with("ts-abc002", "HF_TOKEN", "hf_xxx", secret_type=secret_type) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_teamspace_set_secret_invalid_name( internal_teamspace_api_list_mocker, internal_user_api_mocker, @@ -990,7 +990,7 @@ def test_teamspace_set_secret_invalid_name( ts.set_secret("123_INVALID", "secret_value") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_teamspace_delete_secret( internal_teamspace_api_list_mocker, internal_user_api_mocker, @@ -1003,7 +1003,7 @@ def test_teamspace_delete_secret( mock_delete.assert_called_once_with("ts-abc002", "OLD_SECRET") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_teamspace_delete_secret_invalid_name( internal_teamspace_api_list_mocker, internal_user_api_mocker, @@ -1020,7 +1020,7 @@ def test_teamspace_delete_secret_invalid_name( @mock.patch("lightning_sdk.api.utils.LightningClient") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_new_folder_agnostic( mock_cloud_account_client, internal_user_api_mocker, @@ -1080,7 +1080,7 @@ def test_new_folder_agnostic( @mock.patch("lightning_sdk.api.utils.LightningClient") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_new_folder_byoc( mock_cloud_account_client, internal_user_api_mocker, @@ -1149,7 +1149,7 @@ def test_new_folder_byoc( @pytest.mark.parametrize("writable", [True, False]) @mock.patch("lightning_sdk.api.utils.LightningClient") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_new_connection_efs(mock_cloud_account_client, internal_user_api_mocker, writable): mock_cloud_account_client().projects_service_list_memberships.return_value = V1ListMembershipsResponse( [ @@ -1234,7 +1234,7 @@ def _clear_teamspace_permission_cache(): @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi") @mock.patch("lightning_sdk.teamspace._resolve_user") @pytest.mark.project_permission_test() -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_teamspace_studios_property_when_allowed(mock_resolve_user, mock_teamspace_api, mock_teamspace_api_module): """Test that Teamspace.studios succeeds when Studios permission is allowed.""" mock_user = mock.Mock(spec=User) @@ -1274,7 +1274,7 @@ def test_teamspace_studios_property_when_allowed(mock_resolve_user, mock_teamspa @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi") @mock.patch("lightning_sdk.teamspace._resolve_user") @pytest.mark.project_permission_test() -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_teamspace_studios_property_when_denied(mock_resolve_user, mock_teamspace_api, mock_teamspace_api_module): """Test that Teamspace.studios raises PermissionError when Studios permission is denied.""" mock_user = mock.Mock(spec=User) @@ -1308,7 +1308,7 @@ def test_teamspace_studios_property_when_denied(mock_resolve_user, mock_teamspac @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi") @mock.patch("lightning_sdk.teamspace._resolve_user") @pytest.mark.project_permission_test() -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_teamspace_jobs_property_when_allowed(mock_resolve_user, mock_teamspace_api, mock_teamspace_api_module): """Test that Teamspace.jobs succeeds when Jobs permission is allowed.""" mock_user = mock.Mock(spec=User) @@ -1344,7 +1344,7 @@ def test_teamspace_jobs_property_when_allowed(mock_resolve_user, mock_teamspace_ @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi") @mock.patch("lightning_sdk.teamspace._resolve_user") @pytest.mark.project_permission_test() -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_teamspace_jobs_property_when_denied(mock_resolve_user, mock_teamspace_api, mock_teamspace_api_module): """Test that Teamspace.jobs raises PermissionError when Jobs permission is denied.""" mock_user = mock.Mock(spec=User) @@ -1378,7 +1378,7 @@ def test_teamspace_jobs_property_when_denied(mock_resolve_user, mock_teamspace_a @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi") @mock.patch("lightning_sdk.teamspace._resolve_user") @pytest.mark.project_permission_test() -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_teamspace_multi_machine_jobs_property_when_allowed( mock_resolve_user, mock_teamspace_api, mock_teamspace_api_module ): @@ -1416,7 +1416,7 @@ def test_teamspace_multi_machine_jobs_property_when_allowed( @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi") @mock.patch("lightning_sdk.teamspace._resolve_user") @pytest.mark.project_permission_test() -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_teamspace_multi_machine_jobs_property_when_denied( mock_resolve_user, mock_teamspace_api, mock_teamspace_api_module ): @@ -1497,7 +1497,7 @@ def iter_content(self, chunk_size=None): os.environ, {"LIGHTNING_CLOUD_URL": "https://platform.lightning.ai", "LIGHTNING_AUTH_TOKEN": "test-token"} ) @mock.patch("lightning_sdk.api.utils.LightningClient") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_download_dataset_version( mock_lightning_client, internal_teamspace_api_list_mocker, @@ -1566,7 +1566,7 @@ def fake_get(url, headers=None, stream=None, **kwargs): @mock.patch.dict(os.environ, {"LIGHTNING_CLOUD_URL": "https://platform.lightning.ai"}) @mock.patch("lightning_sdk.api.utils.LightningClient") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_download_dataset_version_no_token_no_cluster( mock_lightning_client, internal_teamspace_api_list_mocker, @@ -1623,7 +1623,7 @@ def test_download_dataset_version_no_token_no_cluster( @mock.patch.dict(os.environ, {"LIGHTNING_CLOUD_URL": "https://platform.lightning.ai"}) @mock.patch("lightning_sdk.api.utils.LightningClient") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_resolve_dataset_id_and_version( mock_lightning_client, internal_teamspace_api_list_mocker, @@ -1717,7 +1717,7 @@ def fake_get(url, headers=None, stream=None, **kwargs): @mock.patch("lightning_sdk.api.dataset._complete_dataset_upload") @mock.patch("lightning_sdk.api.dataset._upload_dataset_files") @mock.patch("lightning_sdk.api.utils.LightningClient") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_upload_dataset_new( mock_lightning_client, mock_upload_files, @@ -1927,7 +1927,7 @@ def __call__(self): @mock.patch("lightning_sdk.api.dataset._complete_dataset_upload") @mock.patch("lightning_sdk.api.dataset._upload_dataset_files") @mock.patch("lightning_sdk.api.utils.LightningClient") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_upload_dataset_existing( mock_lightning_client, mock_upload_files, @@ -2021,7 +2021,7 @@ def test_upload_dataset_existing( @mock.patch("lightning_sdk.api.dataset._complete_dataset_upload") @mock.patch("lightning_sdk.api.dataset._upload_dataset_files") @mock.patch("lightning_sdk.api.utils.LightningClient") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_upload_dataset_with_explicit_version( mock_lightning_client, mock_upload_files, @@ -2095,7 +2095,7 @@ def test_upload_dataset_with_explicit_version( @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.resolve_tag_ids", return_value=["tag-1", "tag-2"]) @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.list_mmts") @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.list_jobs") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_list_jobs_by_tag_resolves_names_once_and_filters_both_job_kinds( list_jobs_mock, list_mmts_mock, @@ -2121,7 +2121,7 @@ def test_list_jobs_by_tag_resolves_names_once_and_filters_both_job_kinds( @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.resolve_tag_ids") @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.list_mmts", return_value=[]) @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.list_jobs", return_value=[]) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_list_jobs_without_tags_skips_the_tag_lookup( list_jobs_mock, list_mmts_mock, @@ -2142,7 +2142,7 @@ def test_list_jobs_without_tags_skips_the_tag_lookup( @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.list_mmts", return_value=[]) @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.list_jobs", return_value=[]) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_jobs_property_is_list_jobs_without_a_filter( list_jobs_mock, list_mmts_mock, @@ -2162,7 +2162,7 @@ def test_jobs_property_is_list_jobs_without_a_filter( ) @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.list_mmts") @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.list_jobs") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_list_jobs_with_an_unknown_tag_fails_before_listing( list_jobs_mock, list_mmts_mock, @@ -2181,7 +2181,7 @@ def test_list_jobs_with_an_unknown_tag_fails_before_listing( @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.list_workload_tags") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_teamspace_tags_are_the_tag_names( list_workload_tags_mock, internal_get_org_api_mocker, @@ -2199,7 +2199,7 @@ def test_teamspace_tags_are_the_tag_names( @mock.patch("lightning_sdk.api.teamspace_api.TeamspaceApi.list_files") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_teamspace_list_files_converts_tree_entries( list_files_mock, internal_get_org_api_mocker, diff --git a/python/tests/core/test_user.py b/python/tests/core/test_user.py index 71609a996..662eb70b4 100644 --- a/python/tests/core/test_user.py +++ b/python/tests/core/test_user.py @@ -7,7 +7,7 @@ from lightning_sdk.user import User -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_user_init_from_name(internal_user_api_mocker): user = User("my-user-name") @@ -16,7 +16,7 @@ def test_user_init_from_name(internal_user_api_mocker): @mock.patch.dict(os.environ, {"LIGHTNING_USERNAME": "my-user-name"}) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_user_init_from_env_var(internal_user_api_mocker): user = User() @@ -24,7 +24,7 @@ def test_user_init_from_env_var(internal_user_api_mocker): assert user.id == "my-user-name" -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_user_teamspaces(internal_user_api_mocker, internal_teamspace_api_list_mocker): user = User("user-abc") @@ -36,7 +36,7 @@ def test_user_teamspaces(internal_user_api_mocker, internal_teamspace_api_list_m assert ts.owner == user -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_equality(internal_user_api_mocker): assert User("my-username") == User("my-username") assert User("my-username") != User("your-username") @@ -46,24 +46,24 @@ class SubUser(User): pass -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_inequality_user_subclass(internal_user_api_mocker): assert User("my-username") != SubUser("my-username") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_repr(internal_user_api_mocker): user = User("my-user-name") assert repr(user) == "User(name=my-user-name)" -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_str(internal_user_api_mocker): user = User("my-user-name") assert str(user) == "User(name=my-user-name)" -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_user_secrets_property(internal_user_api_mocker): user = User("my-user-name") @@ -76,7 +76,7 @@ def test_user_secrets_property(internal_user_api_mocker): mock_get.assert_called_once() -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_user_set_secret(internal_user_api_mocker): user = User("my-user-name") @@ -86,7 +86,7 @@ def test_user_set_secret(internal_user_api_mocker): mock_set.assert_called_once_with("NEW_SECRET", "secret_value") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_user_set_secret_invalid_name(internal_user_api_mocker): user = User("my-user-name") @@ -97,7 +97,7 @@ def test_user_set_secret_invalid_name(internal_user_api_mocker): user.set_secret("123_INVALID", "secret_value") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_user_delete_secret(internal_user_api_mocker): user = User("my-user-name") @@ -107,7 +107,7 @@ def test_user_delete_secret(internal_user_api_mocker): mock_delete.assert_called_once_with("OLD_SECRET") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_user_delete_secret_invalid_name(internal_user_api_mocker): user = User("my-user-name") @@ -122,7 +122,7 @@ def test_user_delete_secret_invalid_name(internal_user_api_mocker): @mock.patch("lightning_sdk.teamspace.Teamspace.__init__", return_value=None) @mock.patch("lightning_sdk.user._get_authed_user") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_user_create_teamspace(mock_get_authed_user, mock_teamspace_init, internal_user_api_mocker): user = User("user-abc") mock_get_authed_user.return_value.id = user.id @@ -135,7 +135,7 @@ def test_user_create_teamspace(mock_get_authed_user, mock_teamspace_init, intern @mock.patch("lightning_sdk.user._get_authed_user") -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_user_create_teamspace_not_authed_user(mock_get_authed_user, internal_user_api_mocker): user = User("user-abc") mock_get_authed_user.return_value.id = "different-user-id" diff --git a/python/tests/lightning_cloud/test_cli_login.py b/python/tests/lightning_cloud/test_cli_login.py new file mode 100644 index 000000000..6b9b9e33a --- /dev/null +++ b/python/tests/lightning_cloud/test_cli_login.py @@ -0,0 +1,114 @@ +import json +import os +from datetime import datetime, timedelta, timezone +from urllib.parse import parse_qs, urlparse + +import pytest + +from lightning_sdk.lightning_cloud import cli_login, env +from lightning_sdk.lightning_cloud.login import Auth + + +class _Response: + def __init__(self, status_code, body): + self.status_code = status_code + self._body = body + + def json(self): + return self._body + + +@pytest.fixture() +def platform(monkeypatch, tmp_path): + login_path = tmp_path / "login.json" + monkeypatch.setattr(env, "LIGHTNING_LOGIN_PATH", str(login_path)) + monkeypatch.setattr(env, "LIGHTNING_CLOUD_URL", "http://control-plane.test") + monkeypatch.setattr(Auth, "secrets_file", tmp_path / "credentials.json") + monkeypatch.setattr(cli_login.getpass, "getpass", lambda prompt: "the-code") + calls = [] + + def post(url, data=None, headers=None, timeout=None): + calls.append((url, dict(data or {}))) + if url.endswith("/oidc/revoke"): + return _Response(200, {}) + return _Response( + 200, + { + "access_token": "lit_ua_access", + "refresh_token": "lit_ur_refresh", + "expires_in": 3600, + "refresh_expires_at": (datetime.now(timezone.utc) + timedelta(days=90)) + .isoformat() + .replace("+00:00", "Z"), + "lightning_grant": {"orgs": "all", "permissions": ["read", "write"]}, + }, + ) + + monkeypatch.setattr(cli_login.requests, "post", post) + return login_path, calls + + +def test_login_stores_the_pair_and_current_login_reads_it(platform, capsys): + login_path, calls = platform + + login = cli_login.login( + no_browser=True, + insecure_storage=True, + orgs=["acme", "beta"], + project="team", + ) + + assert login.access_token == "lit_ua_access" + assert login.refresh_token == "lit_ur_refresh" + assert login.grant == {"orgs": "all", "permissions": ["read", "write"]} + current = cli_login.current_login() + assert current.storage == "file" + assert current.grant == {"orgs": "all", "permissions": ["read", "write"]} + + saved = json.loads(login_path.read_text()) + assert saved["storage"] == "file" + assert saved["access_token"] == "lit_ua_access" + assert (login_path.stat().st_mode & 0o777) == 0o600 + + url = next(word for word in capsys.readouterr().err.split() if "/oidc/authorize" in word) + query = {key: values[0] for key, values in parse_qs(urlparse(url).query).items()} + assert query["orgs_hint"] == "acme,beta" + assert query["project_hint"] == "team" + assert query["redirect_uri"] == "http://control-plane.test/cli-login" + + form = calls[0][1] + assert calls[0][0] == "http://control-plane.test/oidc/token" + assert form["client_id"] == "lightning-cli" + assert form["grant_type"] == "authorization_code" + assert form["code"] == "the-code" + assert form["code_verifier"] + assert form["redirect_uri"] == "http://control-plane.test/cli-login" + + +def test_token_returns_the_access_token_and_logout_revokes_it(platform): + login_path, calls = platform + cli_login.login(no_browser=True, insecure_storage=True) + + assert cli_login.token() == "lit_ua_access" + assert len(calls) == 1 + + cli_login.logout() + + assert calls[-1][0] == "http://control-plane.test/oidc/revoke" + assert calls[-1][1] == {"client_id": "lightning-cli", "token": "lit_ur_refresh"} + assert not login_path.exists() + + +def test_login_refuses_without_a_keychain(monkeypatch, tmp_path): + monkeypatch.setattr(env, "LIGHTNING_LOGIN_PATH", str(tmp_path / "login.json")) + monkeypatch.setattr(cli_login, "_keychain_available", lambda: False) + + with pytest.raises(cli_login.LoginError, match="insecure-storage"): + cli_login.login() + + +def test_file_login_is_readable_only_by_the_user(platform): + login_path, _ = platform + cli_login.login(no_browser=True, insecure_storage=True) + + assert os.stat(login_path).st_mode & 0o777 == 0o600 diff --git a/python/tests/lightning_cloud/test_credentials.py b/python/tests/lightning_cloud/test_credentials.py new file mode 100644 index 000000000..0f698453f --- /dev/null +++ b/python/tests/lightning_cloud/test_credentials.py @@ -0,0 +1,77 @@ +import json + +import pytest + +from lightning_sdk.lightning_cloud import env +from lightning_sdk.lightning_cloud.credentials import CredentialSource +from lightning_sdk.lightning_cloud.login import Auth +from lightning_sdk.lightning_cloud.rest_client import GridRestClient, create_swagger_client + + +class _Api: + def __init__(self, request): + self.request = request + + +@pytest.fixture() +def isolated_credentials(monkeypatch, tmp_path): + for name in ( + "LIGHTNING_AUTH_TOKEN_FILE", + "LIGHTNING_AUTH_TOKEN", + "LIGHTNING_API_KEY", + "LIGHTNING_USER_ID", + ): + monkeypatch.delenv(name, raising=False) + credentials = tmp_path / "credentials.json" + monkeypatch.setattr(Auth, "secrets_file", credentials) + monkeypatch.setattr( + Auth, + "_run_server", + lambda self: (_ for _ in ()).throw(AssertionError("browser login is not expected")), + ) + return credentials + + +def test_token_file_is_reread_between_requests(isolated_credentials, monkeypatch, tmp_path) -> None: + token_file = tmp_path / "token" + token_file.write_text("sess_one\n") + monkeypatch.setenv(env.LIGHTNING_AUTH_TOKEN_FILE, str(token_file)) + seen = [] + + def request(method, url, headers=None, **kwargs): + seen.append(headers["Authorization"]) + + client = GridRestClient(api_client=_Api(request), with_auth=True) + client.api_client.request("GET", "http://example") + token_file.write_text("sess_two") + client.api_client.request("GET", "http://example") + + assert seen == ["Bearer sess_one", "Bearer sess_two"] + + +def test_saved_api_key_is_bearer_and_never_basic(isolated_credentials, monkeypatch) -> None: + isolated_credentials.write_text(json.dumps({"user_id": "user-1", "api_key": "key-1"})) + monkeypatch.setenv("LIGHTNING_USER_ID", "user-1") + monkeypatch.setenv("LIGHTNING_API_KEY", "key-1") + + assert CredentialSource().header() == "Bearer key-1" + assert Auth().authenticate() == "Bearer key-1" + assert Auth().auth_header == "Bearer key-1" + assert Auth().get_auth_header("api_key") == "Bearer key-1" + + +def test_a_legacy_session_login_still_sends_bearer(isolated_credentials) -> None: + isolated_credentials.write_text(json.dumps({"user_id": "user-1", "auth_token": "sess_interim"})) + seen = [] + + def request(method, url, headers=None, **kwargs): + seen.append(headers["Authorization"]) + + GridRestClient(api_client=_Api(request), with_auth=True).api_client.request("GET", "http://example") + + assert seen == ["Bearer sess_interim"] + + +def test_the_user_agent_names_the_installed_sdk() -> None: + assert env.USER_AGENT.startswith("lightning-sdk/") + assert create_swagger_client().user_agent == env.USER_AGENT diff --git a/python/tests/lightning_cloud/test_login.py b/python/tests/lightning_cloud/test_login.py index 660fcae23..cba7da46e 100644 --- a/python/tests/lightning_cloud/test_login.py +++ b/python/tests/lightning_cloud/test_login.py @@ -13,12 +13,18 @@ def test_auth_override_no_longer_accepts_guest(): assert "guest" not in get_args(AuthOverride) -def test_missing_credentials_error_does_not_recommend_guest_login(monkeypatch): - monkeypatch.delenv("LIGHTNING_USER_ID", raising=False) - monkeypatch.delenv("LIGHTNING_API_KEY", raising=False) - monkeypatch.delenv("LIGHTNING_AUTH_TOKEN", raising=False) - - with pytest.raises(ValueError, match="No authentication credentials available") as exc_info: +def test_missing_credentials_error_does_not_recommend_guest_login(monkeypatch, tmp_path): + for name in ("LIGHTNING_USER_ID", "LIGHTNING_API_KEY", "LIGHTNING_AUTH_TOKEN"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("LIGHTNING_AUTH_TOKEN_FILE", str(tmp_path / "missing-token")) + monkeypatch.setattr(Auth, "secrets_file", tmp_path / "credentials.json") + monkeypatch.setattr( + Auth, + "_run_server", + lambda self: (_ for _ in ()).throw(AssertionError("browser login is not expected")), + ) + + with pytest.raises(ValueError, match="No Lightning credentials are available") as exc_info: Auth().get_auth_header() assert "guest_login" not in str(exc_info.value) diff --git a/python/tests/utils/test_resolve.py b/python/tests/utils/test_resolve.py index b16203120..eb4f5485c 100644 --- a/python/tests/utils/test_resolve.py +++ b/python/tests/utils/test_resolve.py @@ -71,7 +71,7 @@ def test_resolve_org_name_empty_env_var(provided): @pytest.mark.parametrize("provided", [None, "org_name", -1]) @mock.patch.dict(os.environ, clear=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_resolve_org(internal_get_org_api_mocker, provided): # can't instantiate outside without proper mocking if provided == -1: @@ -89,7 +89,7 @@ def test_resolve_org(internal_get_org_api_mocker, provided): @mock.patch.dict(os.environ, clear=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_resolve_org_with_user(internal_get_org_api_mocker): name = "user-name" with pytest.raises(ValueError, match=f"Organization '{name}' does not exist or you are not a member of it."): @@ -126,7 +126,7 @@ def test_resolve_user_name_env_var(provided): @pytest.mark.parametrize("provided", [None, "user_name", -1]) @mock.patch.dict(os.environ, clear=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_resolve_user(internal_user_api_mocker, provided): # can't instantiate outside without proper mocking if provided == -1: @@ -178,7 +178,7 @@ def test_resolve_teamspace_name_env_var(provided): ("ts-def", "org-abc", None, {"name": "ts-def", "org": {"name": "org-abc"}}), ], ) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_resolve_teamspace_combinations( internal_user_api_mocker, internal_get_org_api_mocker, @@ -403,7 +403,7 @@ def test_parse_model_name_and_version(): @mock.patch.dict(os.environ, {"LIGHTNING_CLOUD_URL": "lightning.ai:443"}, clear=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_studio_url(internal_studio_init_mocker): studio = Studio("st-abc", "ts-abc", org="org-abc") @@ -413,7 +413,7 @@ def test_get_studio_url(internal_studio_init_mocker): @mock.patch.dict(os.environ, {"LIGHTNING_CLOUD_URL": "lightning.ai:443"}, clear=True) -@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_get_studio_url_turn_on(internal_studio_init_mocker): studio = Studio("st-abc", "ts-abc", org="org-abc") From 678347c88096350fc20114644970b8fe41dd856a Mon Sep 17 00:00:00 2001 From: fnwoke-lightning Date: Wed, 4 Nov 2026 13:47:23 -0500 Subject: [PATCH 02/27] feat(cli): ship the PKCE login and auth commands Add lightning login options, auth status, auth token and logout, and move the remaining SDK credential paths to the per-request source. --- python/lightning_sdk/api/filesystem_api.py | 9 +- python/lightning_sdk/api/license_api.py | 4 +- python/lightning_sdk/api/logs_api.py | 2 +- python/lightning_sdk/api/sandbox_api.py | 21 ++-- python/lightning_sdk/api/utils.py | 3 +- python/lightning_sdk/cli/auth/__init__.py | 4 + python/lightning_sdk/cli/auth/status.py | 115 ++++++++++++++++++ python/lightning_sdk/cli/auth/token.py | 26 ++++ .../lightning_sdk/cli/resource_completion.py | 8 +- python/lightning_sdk/cli/utils/logging.py | 9 +- python/tests/api/test_license_api.py | 4 +- python/tests/api/test_utils.py | 44 ++++--- python/tests/cli/auth/test_status_token.py | 109 +++++++++++++++++ python/tests/cli/cp/test_completion.py | 4 +- python/tests/cli/utils/test_logging.py | 3 +- python/tests/core/test_sandbox_config.py | 10 +- python/tests/test_license_integration.py | 23 +++- 17 files changed, 332 insertions(+), 66 deletions(-) create mode 100644 python/lightning_sdk/cli/auth/status.py create mode 100644 python/lightning_sdk/cli/auth/token.py create mode 100644 python/tests/cli/auth/test_status_token.py diff --git a/python/lightning_sdk/api/filesystem_api.py b/python/lightning_sdk/api/filesystem_api.py index 9cc195f62..c42431a58 100644 --- a/python/lightning_sdk/api/filesystem_api.py +++ b/python/lightning_sdk/api/filesystem_api.py @@ -31,7 +31,6 @@ class FilesystemApi: def __init__(self) -> None: self._client = cached_lightning_client() - self._auth_headers = _authenticate_and_get_auth_headers() @property def client(self) -> LightningClient: @@ -75,7 +74,7 @@ def fetch_page(query_params: Dict[str, str]) -> Dict: r = requests.get( f"{self._client.api_client.configuration.host}/v1/projects/{teamspace_id}/artifacts/trees/{path}", params=query_params, - headers=self._auth_headers, + headers=_authenticate_and_get_auth_headers(), ) if r.status_code != 200: raise _RemoteApiError(f"Failed to list files: {r.status_code}", status_code=r.status_code) @@ -164,7 +163,7 @@ def download() -> None: try: r = requests.get( url, - headers=self._auth_headers, + headers=_authenticate_and_get_auth_headers(), stream=True, allow_redirects=True, timeout=(_DOWNLOAD_CONNECT_TIMEOUT_SECONDS, _DOWNLOAD_READ_TIMEOUT_SECONDS), @@ -288,7 +287,7 @@ def delete_file(self, teamspace_id: str, remote_path: str) -> None: r = requests.delete( f"{self._client.api_client.configuration.host}" f"/v1/projects/{teamspace_id}/artifacts/blobs/{remote_path.strip('/')}", - headers=self._auth_headers, + headers=_authenticate_and_get_auth_headers(), ) if r.status_code not in (200, 204): raise _RemoteApiError( @@ -311,7 +310,7 @@ def delete_folder(self, teamspace_id: str, remote_path: str) -> None: r = requests.delete( f"{self._client.api_client.configuration.host}" f"/v1/projects/{teamspace_id}/artifacts/trees/{remote_path.strip('/')}", - headers=self._auth_headers, + headers=_authenticate_and_get_auth_headers(), ) if r.status_code not in (200, 204): raise _RemoteApiError( diff --git a/python/lightning_sdk/api/license_api.py b/python/lightning_sdk/api/license_api.py index ed4ca08a2..898b77b30 100644 --- a/python/lightning_sdk/api/license_api.py +++ b/python/lightning_sdk/api/license_api.py @@ -9,8 +9,8 @@ class LicenseApi: def __init__(self, login_token: str) -> None: self._cloud_url = _cloud_url() self._auth = Auth() - self._auth.token_login(login_token, save_token=True) - self._client = self._auth.create_api_client() + self._auth.token_login(login_token, save_token=False) + self._client = self._auth.create_api_client(override="auth_token") self._api = ProductLicenseServiceApi(self._client) def validate_license(self, license_key: str, product_id: str) -> bool: diff --git a/python/lightning_sdk/api/logs_api.py b/python/lightning_sdk/api/logs_api.py index fa6b99043..7b8de8055 100644 --- a/python/lightning_sdk/api/logs_api.py +++ b/python/lightning_sdk/api/logs_api.py @@ -506,7 +506,6 @@ def follow( "Following logs requires the 'websocket-client' package (pip install websocket-client)." ) from ex - auth_header = Auth().authenticate() url = _websocket_url(follow_url) # Poll `recv()` rather than blocking on it, so `stop` and `idle_timeout` get a chance to # run. The poll has to stay below the server's heartbeat interval: a heartbeat resets the @@ -521,6 +520,7 @@ def follow( while True: ws = None try: + auth_header = Auth().authenticate() # The handshake gets a longer budget than the deliberately short read poll. ws = websocket.create_connection( url, header=[f"Authorization: {auth_header}"], timeout=_CONNECT_TIMEOUT diff --git a/python/lightning_sdk/api/sandbox_api.py b/python/lightning_sdk/api/sandbox_api.py index 81a9cd264..03c611e99 100644 --- a/python/lightning_sdk/api/sandbox_api.py +++ b/python/lightning_sdk/api/sandbox_api.py @@ -4,7 +4,7 @@ from typing import Any from lightning_sdk.lightning_cloud import env as lightning_env -from lightning_sdk.lightning_cloud.login import Auth +from lightning_sdk.lightning_cloud.credentials import CredentialSource from lightning_sdk.lightning_cloud.openapi import JobsServiceApi, SandboxesServiceApi from lightning_sdk.lightning_cloud.openapi.models import ( SandboxesServiceCreateSandboxDirectoryBody, @@ -17,7 +17,7 @@ V1SandboxSnapshot, ) from lightning_sdk.lightning_cloud.openapi.rest import ApiException -from lightning_sdk.lightning_cloud.rest_client import LightningClient +from lightning_sdk.lightning_cloud.rest_client import LightningClient, credential_request_wrapper @dataclass @@ -228,10 +228,9 @@ def _apply_auth_and_host(self) -> None: def _ensure_auth(self) -> None: if self._auth_configured: return - auth_header = Auth().authenticate() - if not auth_header: - raise RuntimeError("Missing credentials: set api_key via configure() or use lightning login / env vars.") - self._client.api_client.set_default_header("Authorization", auth_header) + self._client.api_client.request = credential_request_wrapper( + self._client.api_client.request, CredentialSource() + ) self._auth_configured = True @property @@ -243,11 +242,13 @@ def host(self) -> str: return str(self._client.api_client.configuration.host).rstrip("/") def auth_header(self) -> str: - """Return the configured Authorization header, authenticating lazily when needed.""" + """Return the Authorization header the next request would send.""" self._ensure_auth() - headers = getattr(self._client.api_client, "default_headers", {}) - auth = headers.get("Authorization", "") if isinstance(headers, dict) else "" - return str(auth or "") + if self._config.get("api_key"): + headers = getattr(self._client.api_client, "default_headers", {}) + auth = headers.get("Authorization", "") if isinstance(headers, dict) else "" + return str(auth or "") + return CredentialSource().header() def sandboxes(self) -> SandboxesServiceApi: self._ensure_auth() diff --git a/python/lightning_sdk/api/utils.py b/python/lightning_sdk/api/utils.py index a748aa0dd..61e97efad 100644 --- a/python/lightning_sdk/api/utils.py +++ b/python/lightning_sdk/api/utils.py @@ -254,7 +254,6 @@ def __init__( if self.filesize > self.chunk_size * _MAX_UPLOAD_PARTS: self.chunk_size = math.ceil(self.filesize / _MAX_UPLOAD_PARTS) self.max_workers = int(os.environ.get("LIGHTNING_MULTI_PART_MAX_WORKERS", _MAX_WORKERS)) - self._auth_headers = _authenticate_and_get_auth_headers() def __call__(self) -> None: """Execute the upload, dispatching to single-part or multipart.""" @@ -299,7 +298,7 @@ def _post_blob_request(self, suffix: str, blob: Dict[str, Any], action: str) -> body: Dict[str, Any] = {"blobs": [blob]} if self.cluster_id: body["cluster_id"] = self.cluster_id - r = requests.post(url, json=body, headers=self._auth_headers, timeout=30) + r = requests.post(url, json=body, headers=_authenticate_and_get_auth_headers(), timeout=30) if r.status_code >= 500 or r.status_code == 429: raise HTTPError( f"Transient error trying to {action} '{self.remote_path}'. Status code: {r.status_code}", response=r diff --git a/python/lightning_sdk/cli/auth/__init__.py b/python/lightning_sdk/cli/auth/__init__.py index adb626f49..06b18183b 100644 --- a/python/lightning_sdk/cli/auth/__init__.py +++ b/python/lightning_sdk/cli/auth/__init__.py @@ -7,8 +7,12 @@ def register_commands(group: click.Group) -> None: """Register auth commands with the given group.""" from lightning_sdk.cli.auth.role import role from lightning_sdk.cli.auth.roles import roles + from lightning_sdk.cli.auth.status import status + from lightning_sdk.cli.auth.token import token from lightning_sdk.cli.auth.whoami import whoami group.add_command(whoami) + group.add_command(status) + group.add_command(token) group.add_command(roles) group.add_command(role) diff --git a/python/lightning_sdk/cli/auth/status.py b/python/lightning_sdk/cli/auth/status.py new file mode 100644 index 000000000..c0f3a1698 --- /dev/null +++ b/python/lightning_sdk/cli/auth/status.py @@ -0,0 +1,115 @@ +"""Show the saved CLI login.""" + +from datetime import datetime +from typing import Any, Dict, Optional + +import rich_click as click +from rich.table import Table + +from lightning_sdk.api.utils import cached_lightning_client +from lightning_sdk.cli.utils.logging import LightningCommand +from lightning_sdk.cli.utils.richt_print import rich_to_str +from lightning_sdk.lightning_cloud import cli_login +from lightning_sdk.lightning_cloud.login import Auth + +_NOT_LOGGED_IN = "You are not logged in. Run `lightning login` first." + + +def _stamp(epoch: float) -> str: + return datetime.fromtimestamp(epoch).astimezone().strftime("%Y-%m-%d %H:%M:%S %Z") + + +def _user(client: Any) -> Dict[str, Optional[str]]: + try: + user = client.auth_service_get_user() + except Exception: + return {} + return { + "id": getattr(user, "id", None), + "username": getattr(user, "username", None), + "email": getattr(user, "email", None), + } + + +def _org_names(client: Any) -> Dict[str, str]: + try: + orgs = client.organizations_service_list_organizations().organizations or [] + except Exception: + return {} + return {org.id: org.name for org in orgs if getattr(org, "id", None)} + + +def _teamspace_names(client: Any) -> Dict[str, str]: + try: + memberships = client.projects_service_list_memberships(filter_by_user_id=True).memberships or [] + except Exception: + return {} + return {member.project_id: member.name for member in memberships if getattr(member, "project_id", None)} + + +def _grant_lines(grant: Dict[str, Any], orgs: Dict[str, str], teamspaces: Dict[str, str]) -> tuple[str, str]: + scopes = grant.get("orgs") + if scopes == "all": + return "all", "all" + if not isinstance(scopes, list) or not scopes: + return "none", "none" + + shown_orgs: list[str] = [] + shown_teamspaces: list[str] = [] + for scope in scopes: + if not isinstance(scope, dict): + continue + org_id = str(scope.get("org", "")) + shown_orgs.append(orgs.get(org_id, org_id)) + for project in scope.get("projects") or []: + project_id = str(project) + shown_teamspaces.append(teamspaces.get(project_id, project_id)) + return ", ".join(shown_orgs) or "none", ", ".join(shown_teamspaces) or "none" + + +def _legacy_session() -> bool: + auth = Auth() + return bool(auth.load() and auth.auth_token) + + +@click.command("status", cls=LightningCommand) +def status() -> None: + """Show who is logged in, the login's reach, its permissions and when it expires.""" + login = cli_login.current_login() + client = cached_lightning_client() + user = _user(client) + + if login is None: + if not _legacy_session(): + raise click.ClickException(_NOT_LOGGED_IN) + table = Table(show_header=False, box=None) + table.add_column("Field", style="bold") + table.add_column("Value") + table.add_row("User", user.get("username") or user.get("email") or user.get("id") or "unknown") + table.add_row("Login", "session token saved by an older SDK") + table.add_row("Expiry", "managed by the server") + table.add_row("Hint", "Run `lightning login` to use the new login.") + click.echo(rich_to_str(table), color=True) + return + + if login.grant.get("orgs") == "all": + orgs, teamspaces = "all", "all" + else: + orgs, teamspaces = _grant_lines(login.grant, _org_names(client), _teamspace_names(client)) + permissions = ", ".join(login.grant.get("permissions") or []) or "none" + + table = Table(show_header=False, box=None) + table.add_column("Field", style="bold") + table.add_column("Value") + table.add_row("User", user.get("username") or user.get("email") or user.get("id") or "unknown") + if user.get("email"): + table.add_row("Email", user["email"]) + if user.get("id"): + table.add_row("User ID", user["id"]) + table.add_row("Storage", login.storage) + table.add_row("Orgs", orgs) + table.add_row("Teamspaces", teamspaces) + table.add_row("Permissions", permissions) + table.add_row("Access expires", _stamp(login.access_expires_at)) + table.add_row("Login expires", _stamp(login.refresh_expires_at)) + click.echo(rich_to_str(table), color=True) diff --git a/python/lightning_sdk/cli/auth/token.py b/python/lightning_sdk/cli/auth/token.py new file mode 100644 index 000000000..afa4d239e --- /dev/null +++ b/python/lightning_sdk/cli/auth/token.py @@ -0,0 +1,26 @@ +"""Print an access token for curl and other callers.""" + +import rich_click as click + +from lightning_sdk.cli.utils.logging import LightningCommand +from lightning_sdk.lightning_cloud import cli_login +from lightning_sdk.lightning_cloud.login import Auth + + +@click.command("token", cls=LightningCommand) +def token() -> None: + """Print an access token with at least 30 minutes left, refreshing it first when needed.""" + try: + click.echo(cli_login.token()) + return + except cli_login.LoginError as err: + failure = err + + try: + header = Auth().auth_header + except (ValueError, ConnectionError): + header = None + if header and header.startswith("Bearer "): + click.echo(header.removeprefix("Bearer ")) + return + raise click.ClickException(str(failure)) from failure diff --git a/python/lightning_sdk/cli/resource_completion.py b/python/lightning_sdk/cli/resource_completion.py index ac613cb25..48ad03f13 100644 --- a/python/lightning_sdk/cli/resource_completion.py +++ b/python/lightning_sdk/cli/resource_completion.py @@ -13,12 +13,10 @@ def has_credentials() -> bool: """Check for existing credentials without starting the browser login flow.""" - auth = Auth() - if auth.api_key or auth.auth_token: - return True - if not auth.load(): + try: + return bool(Auth().auth_header) + except Exception: return False - return bool(auth.api_key or auth.auth_token) def accessible_teamspaces() -> dict[str, dict[str, str]]: diff --git a/python/lightning_sdk/cli/utils/logging.py b/python/lightning_sdk/cli/utils/logging.py index 5158be6fe..d39716bc5 100644 --- a/python/lightning_sdk/cli/utils/logging.py +++ b/python/lightning_sdk/cli/utils/logging.py @@ -33,15 +33,8 @@ def _auth_header_without_browser() -> Optional[str]: - auth = Auth() - if not (auth.api_key or auth.auth_token): - try: - if not auth.load(): - return None - except Exception: - return None try: - return auth.auth_header + return Auth().auth_header except Exception: return None diff --git a/python/tests/api/test_license_api.py b/python/tests/api/test_license_api.py index ce2a7f2f5..14eb6879d 100644 --- a/python/tests/api/test_license_api.py +++ b/python/tests/api/test_license_api.py @@ -40,8 +40,8 @@ def test_license_api_initialization(mock_cloud_url, mock_api_class, mock_auth_cl license_api = LicenseApi(test_token) mock_auth_class.assert_called_once() - mock_auth_instance.token_login.assert_called_once_with(test_token, save_token=True) - mock_auth_instance.create_api_client.assert_called_once() + mock_auth_instance.token_login.assert_called_once_with(test_token, save_token=False) + mock_auth_instance.create_api_client.assert_called_once_with(override="auth_token") mock_api_class.assert_called_once_with(mock_client) assert license_api._api == mock_api_instance diff --git a/python/tests/api/test_utils.py b/python/tests/api/test_utils.py index cc0857118..b1befbea2 100644 --- a/python/tests/api/test_utils.py +++ b/python/tests/api/test_utils.py @@ -790,21 +790,27 @@ def test_resolve_path_mappings(): assert path_mappings[1].connection_path == "" +@pytest.fixture() +def upload_auth(monkeypatch): + headers = {"Authorization": "Bearer test-token"} + monkeypatch.setattr( + "lightning_sdk.api.utils._authenticate_and_get_auth_headers", + lambda: headers, + ) + return headers + + def _make_single_part_blob_uploader(tmp_path, progress_bar=False, extra_headers=None, content=b"hello world"): file_path = tmp_path / "test_file.txt" file_path.write_bytes(content) - with mock.patch( - "lightning_sdk.api.utils._authenticate_and_get_auth_headers", - return_value={"Authorization": "Bearer test-token"}, - ): - return utils._BlobUploader( - client=Mock(), - endpoint_base=_TEST_ENDPOINT_BASE, - file_path=str(file_path), - remote_path="test_file.txt", - progress_bar=progress_bar, - extra_headers=extra_headers, - ) + return utils._BlobUploader( + client=Mock(), + endpoint_base=_TEST_ENDPOINT_BASE, + file_path=str(file_path), + remote_path="test_file.txt", + progress_bar=progress_bar, + extra_headers=extra_headers, + ) def _configure_single_part_create(mock_requests, url="http://example.com/upload", headers=None): @@ -812,7 +818,7 @@ def _configure_single_part_create(mock_requests, url="http://example.com/upload" @mock.patch("lightning_sdk.api.utils.requests") -def test_single_part_uploader_basic(mock_requests, tmp_path): +def test_single_part_uploader_basic(mock_requests, tmp_path, upload_auth): uploader = _make_single_part_blob_uploader(tmp_path) _configure_single_part_create(mock_requests) @@ -836,7 +842,7 @@ def test_single_part_uploader_basic(mock_requests, tmp_path): @mock.patch("lightning_sdk.api.utils.requests") -def test_single_part_uploader_with_progress_bar(mock_requests, tmp_path): +def test_single_part_uploader_with_progress_bar(mock_requests, tmp_path, upload_auth): uploader = _make_single_part_blob_uploader(tmp_path, progress_bar=True) _configure_single_part_create(mock_requests) @@ -852,7 +858,7 @@ def test_single_part_uploader_with_progress_bar(mock_requests, tmp_path): @mock.patch("time.sleep", return_value=None) @mock.patch("lightning_sdk.api.utils.requests") -def test_single_part_uploader_retries_on_http_error(mock_requests, _, tmp_path): +def test_single_part_uploader_retries_on_http_error(mock_requests, _, tmp_path, upload_auth): uploader = _make_single_part_blob_uploader(tmp_path) _configure_single_part_create(mock_requests) @@ -870,7 +876,7 @@ def test_single_part_uploader_retries_on_http_error(mock_requests, _, tmp_path): @mock.patch("time.sleep", return_value=None) @mock.patch("lightning_sdk.api.utils.requests") -def test_single_part_uploader_retries_on_transient_status(mock_requests, _, tmp_path): +def test_single_part_uploader_retries_on_transient_status(mock_requests, _, tmp_path, upload_auth): uploader = _make_single_part_blob_uploader(tmp_path) _configure_single_part_create(mock_requests) @@ -887,7 +893,7 @@ def test_single_part_uploader_retries_on_transient_status(mock_requests, _, tmp_ @mock.patch("time.sleep", return_value=None) @mock.patch("lightning_sdk.api.utils.requests") -def test_single_part_uploader_retries_auth_errors_with_fresh_url(mock_requests, _, tmp_path): +def test_single_part_uploader_retries_auth_errors_with_fresh_url(mock_requests, _, tmp_path, upload_auth): uploader = _make_single_part_blob_uploader(tmp_path) _configure_single_part_create(mock_requests) @@ -911,7 +917,9 @@ def test_single_part_uploader_retries_auth_errors_with_fresh_url(mock_requests, @pytest.mark.parametrize("content", [b"hello world", b""]) @mock.patch("time.sleep", return_value=None) @mock.patch("lightning_sdk.api.utils.requests") -def test_single_part_uploader_does_not_retry_client_error(mock_requests, _, tmp_path, content, status_code): +def test_single_part_uploader_does_not_retry_client_error( + mock_requests, _, tmp_path, content, status_code, upload_auth +): uploader = _make_single_part_blob_uploader(tmp_path, content=content) _configure_single_part_create(mock_requests) diff --git a/python/tests/cli/auth/test_status_token.py b/python/tests/cli/auth/test_status_token.py new file mode 100644 index 000000000..ae2dd477a --- /dev/null +++ b/python/tests/cli/auth/test_status_token.py @@ -0,0 +1,109 @@ +import time +from types import SimpleNamespace +from unittest.mock import MagicMock, PropertyMock + +from click.testing import CliRunner + +from lightning_sdk.cli.entrypoint import main_cli +from lightning_sdk.lightning_cloud.cli_login import LoginError, LoginStatus + + +def _status(grant=None) -> LoginStatus: + now = time.time() + return LoginStatus( + url="https://platform.lightning.ai", + storage="keychain", + access_expires_at=now + 3600, + refresh_expires_at=now + 90 * 24 * 3600, + grant=grant or {"orgs": "all", "permissions": ["read", "write"]}, + ) + + +def _client() -> MagicMock: + client = MagicMock() + client.auth_service_get_user.return_value = SimpleNamespace(id="user-1", username="ada", email="ada@example.com") + client.organizations_service_list_organizations.return_value = SimpleNamespace( + organizations=[SimpleNamespace(id="org-1", name="acme"), SimpleNamespace(id="org-2", name="beta")] + ) + client.projects_service_list_memberships.return_value = SimpleNamespace( + memberships=[SimpleNamespace(project_id="proj-1", name="research")] + ) + return client + + +def test_auth_status_shows_the_login_without_a_token(monkeypatch) -> None: + monkeypatch.setattr("lightning_sdk.cli.auth.status.cli_login.current_login", lambda: _status()) + monkeypatch.setattr("lightning_sdk.cli.auth.status.cached_lightning_client", _client) + + result = CliRunner().invoke(main_cli, ["auth", "status"]) + + assert result.exit_code == 0, result.output + assert "ada" in result.output + assert "ada@example.com" in result.output + assert "keychain" in result.output + assert "all" in result.output + assert "read, write" in result.output + assert "lit_ua_" not in result.output + + +def test_auth_status_resolves_a_narrowed_grant(monkeypatch) -> None: + grant = {"orgs": [{"org": "org-1", "projects": ["proj-1"]}], "permissions": ["read"]} + monkeypatch.setattr("lightning_sdk.cli.auth.status.cli_login.current_login", lambda: _status(grant)) + monkeypatch.setattr("lightning_sdk.cli.auth.status.cached_lightning_client", _client) + + result = CliRunner().invoke(main_cli, ["auth", "status"]) + + assert result.exit_code == 0, result.output + assert "acme" in result.output + assert "research" in result.output + assert "read" in result.output + + +def test_auth_status_instructs_login_without_a_login(monkeypatch) -> None: + monkeypatch.setattr("lightning_sdk.cli.auth.status.cli_login.current_login", lambda: None) + auth = MagicMock() + auth.load.return_value = False + monkeypatch.setattr("lightning_sdk.cli.auth.status.Auth", lambda: auth) + + result = CliRunner().invoke(main_cli, ["auth", "status"]) + + assert result.exit_code == 1 + assert "Run `lightning login` first" in result.output + + +def test_auth_token_prints_only_the_token(monkeypatch) -> None: + monkeypatch.setattr("lightning_sdk.cli.auth.token.cli_login.token", lambda: "lit_ua_access") + + result = CliRunner().invoke(main_cli, ["auth", "token"]) + + assert result.exit_code == 0, result.output + assert result.output == "lit_ua_access\n" + + +def test_auth_token_falls_back_to_a_saved_session(monkeypatch) -> None: + monkeypatch.setattr( + "lightning_sdk.cli.auth.token.cli_login.token", + MagicMock(side_effect=LoginError("You are not logged in. Run `lightning login`.")), + ) + auth = MagicMock(auth_header="Bearer sess_saved") + monkeypatch.setattr("lightning_sdk.cli.auth.token.Auth", lambda: auth) + + result = CliRunner().invoke(main_cli, ["auth", "token"]) + + assert result.exit_code == 0, result.output + assert result.output == "sess_saved\n" + + +def test_auth_token_reports_a_missing_login(monkeypatch) -> None: + monkeypatch.setattr( + "lightning_sdk.cli.auth.token.cli_login.token", + MagicMock(side_effect=LoginError("You are not logged in. Run `lightning login`.")), + ) + auth = MagicMock() + type(auth).auth_header = PropertyMock(side_effect=ValueError("missing")) + monkeypatch.setattr("lightning_sdk.cli.auth.token.Auth", lambda: auth) + + result = CliRunner().invoke(main_cli, ["auth", "token"]) + + assert result.exit_code == 1 + assert "lightning login" in result.output diff --git a/python/tests/cli/cp/test_completion.py b/python/tests/cli/cp/test_completion.py index f696f9b80..6f74f2277 100644 --- a/python/tests/cli/cp/test_completion.py +++ b/python/tests/cli/cp/test_completion.py @@ -120,9 +120,7 @@ def test_remote_completion_silently_ignores_api_failures(_accessible_teamspaces) @patch("lightning_sdk.cli.cp.completion._accessible_teamspaces") @patch("lightning_sdk.cli.resource_completion.Auth") def test_remote_completion_does_not_start_login_without_credentials(auth, accessible_teamspaces): - auth.return_value.api_key = None - auth.return_value.auth_token = None - auth.return_value.load.return_value = False + auth.return_value.auth_header = None items = _complete_argument("source", "lit://") diff --git a/python/tests/cli/utils/test_logging.py b/python/tests/cli/utils/test_logging.py index 155668346..bde8cb72c 100644 --- a/python/tests/cli/utils/test_logging.py +++ b/python/tests/cli/utils/test_logging.py @@ -152,8 +152,7 @@ def test_log_command_empty_message(self, mock_client_class, mock_auth_header): @mock.patch("lightning_sdk.cli.utils.logging.LightningClient") def test_log_command_skips_without_credentials(self, mock_client_class, mock_auth_class): """Test command logging does not trigger authentication when no credentials exist.""" - auth = mock.Mock(api_key=None, auth_token=None) - auth.load.return_value = False + auth = mock.Mock(auth_header=None) mock_auth_class.return_value = auth _log_command(message="Test message", duration=100) diff --git a/python/tests/core/test_sandbox_config.py b/python/tests/core/test_sandbox_config.py index 886478e8e..8ce15ce09 100644 --- a/python/tests/core/test_sandbox_config.py +++ b/python/tests/core/test_sandbox_config.py @@ -36,15 +36,15 @@ def test_configure_rejects_organization_id_kwarg(): @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_sandbox_config_api_without_api_key_uses_lightning_auth(): - with mock.patch("lightning_sdk.api.sandbox_api.Auth") as auth_cls: - auth_cls.return_value.authenticate.return_value = "Basic auth" + with mock.patch("lightning_sdk.api.sandbox_api.CredentialSource") as source_cls: + source_cls.return_value.header.return_value = "Bearer auth" api = SandboxConfig(base_url="https://x").api() - auth_cls.return_value.authenticate.assert_not_called() - api.sandboxes() + source_cls.return_value.header.assert_not_called() + assert api.auth_header() == "Bearer auth" assert api.config_get("api_key") is None assert api.config_get("base_url") == "https://x" - auth_cls.return_value.authenticate.assert_called_once() + source_cls.return_value.header.assert_called_once() @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) diff --git a/python/tests/test_license_integration.py b/python/tests/test_license_integration.py index 070dc7d4d..100cd00bb 100644 --- a/python/tests/test_license_integration.py +++ b/python/tests/test_license_integration.py @@ -40,7 +40,8 @@ def test_full_license_validation_workflow(mock_cloud_url, mock_api_class, mock_a # Assert assert result is True - mock_auth_instance.token_login.assert_called_once_with(license_key, save_token=True) + mock_auth_instance.token_login.assert_called_once_with(license_key, save_token=False) + mock_auth_instance.create_api_client.assert_called_once_with(override="auth_token") mock_api_instance.product_license_service_validate_license.assert_called_once() @@ -71,7 +72,8 @@ def test_full_license_validation_workflow_failure(mock_cloud_url, mock_api_class # Assert assert result is False - mock_auth_instance.token_login.assert_called_once_with(license_key, save_token=True) + mock_auth_instance.token_login.assert_called_once_with(license_key, save_token=False) + mock_auth_instance.create_api_client.assert_called_once_with(override="auth_token") mock_api_instance.product_license_service_validate_license.assert_called_once() @@ -101,5 +103,20 @@ def test_full_license_validation_workflow_exception(mock_cloud_url, mock_api_cla license_obj.validate() assert str(exc_info.value) == f"Invalid license key {license_key} for product {product_name}" - mock_auth_instance.token_login.assert_called_once_with(license_key, save_token=True) + mock_auth_instance.token_login.assert_called_once_with(license_key, save_token=False) + mock_auth_instance.create_api_client.assert_called_once_with(override="auth_token") mock_api_instance.product_license_service_validate_license.assert_called_once() + + +def test_license_validation_leaves_the_saved_login_unchanged(): + from lightning_sdk.api.license_api import LicenseApi + + fake_auth = Mock() + + with patch("lightning_sdk.api.license_api.Auth", return_value=fake_auth), patch( + "lightning_sdk.api.license_api.ProductLicenseServiceApi" + ): + LicenseApi("license-key") + + fake_auth.token_login.assert_called_once_with("license-key", save_token=False) + fake_auth.save.assert_not_called() From 05f68ba12312a3943a25122e442598277b4b1c49 Mon Sep 17 00:00:00 2001 From: fnwoke-lightning Date: Wed, 4 Nov 2026 18:01:38 -0500 Subject: [PATCH 03/27] fix(lightning_cloud): keep the interim sess_ login as the default A release from main must keep logging in on production, which serves only the interim login until relaunch. A new login is the interim one unless LIGHTNING_CLI_OIDC_LOGIN is set, read when the login starts. - Restore AuthServer with get_auth_url(port) and login_with_browser(auth). It sends state on the /sign-in redirectTo, adds caller sign-in params, refuses a callback whose state doesn't match, saves the sess_ as auth_token with Auth.save, and redirects with no credential in the URL. - Auth.save accepts and ignores token again, writes a 0600 file. - credentials: CredentialKind, Credential, resolve(), refresh(token) for the token a request sent, start_login() dispatching on the switch with LoginUnavailable for options the interim login can't honour, and logout() that revokes each saved kind and reports which revokes failed. - A login that saves revokes the login it replaced, after the save: /oidc/revoke for a new login, /v1/auth/logout for a sess_. - Both browser logins stop after five minutes with no callback. - Reuse is detected only from the exact error_description. - The source's own login goes through Auth.authenticate, so main's browser_authentication(False) still keeps resource commands from opening a browser. - Restore llm_api.py, cli/entrypoint.py, cli/legacy/entrypoint.py and cli/utils/auth.py to main. Their changes belong to sdk#235. Co-authored-by: Cursor --- python/lightning_sdk/api/llm_api.py | 34 +- python/lightning_sdk/cli/entrypoint.py | 70 +-- python/lightning_sdk/cli/legacy/entrypoint.py | 13 +- python/lightning_sdk/cli/utils/auth.py | 381 ++++++++++++++- .../lightning_cloud/cli_login.py | 123 +++-- .../lightning_cloud/credentials.py | 217 ++++++--- python/lightning_sdk/lightning_cloud/login.py | 113 ++++- .../lightning_cloud/rest_client.py | 5 +- python/tests/cli/test_entrypoint.py | 157 ++++--- python/tests/cli/utils/test_auth.py | 436 ++++++++++++++++-- .../tests/lightning_cloud/test_cli_login.py | 4 +- .../tests/lightning_cloud/test_credentials.py | 9 +- 12 files changed, 1265 insertions(+), 297 deletions(-) diff --git a/python/lightning_sdk/api/llm_api.py b/python/lightning_sdk/api/llm_api.py index 501ae6281..712bf1933 100644 --- a/python/lightning_sdk/api/llm_api.py +++ b/python/lightning_sdk/api/llm_api.py @@ -4,15 +4,16 @@ import json import os import threading +import time import warnings from typing import Any, AsyncGenerator, Dict, Generator, List, Optional, Union +from urllib.parse import urlencode from rich.prompt import Confirm from lightning_sdk.api.utils import cached_lightning_client -from lightning_sdk.lightning_cloud import cli_login -from lightning_sdk.lightning_cloud.credentials import browser_login -from lightning_sdk.lightning_cloud.login import Auth +from lightning_sdk.lightning_cloud import env +from lightning_sdk.lightning_cloud.login import Auth, AuthServer from lightning_sdk.lightning_cloud.openapi.models import ( StreamResultOfV1ConversationResponseChunk, V1ConversationResponseChunk, @@ -21,6 +22,19 @@ V1ResponseChoiceDelta, ) +LITAI_CODE = os.environ.get("LITAI_CODE", "x334uv8t7v") + + +class _AuthServer(AuthServer): + def __init__(self, model: str, *args: Any, **kwargs: Any) -> None: + self._model = model + super().__init__(*args, **kwargs) + + def get_auth_url(self, port: int) -> str: + redirect_uri = f"http://localhost:{port}/login-complete" + params = urlencode({"redirectTo": redirect_uri, "okbhrt": LITAI_CODE, "litaimodel": self._model}) + return f"{env.LIGHTNING_CLOUD_URL}/sign-in?{params}" + class _AuthLitAI(Auth): def __init__(self, model: str, shall_confirm: bool = False) -> None: @@ -41,19 +55,15 @@ def _run_server(self) -> None: "Login cancelled. Please login at https://platform.lightning.ai/sign-up. " "Or run `lightning login` to login." ) - cli_login.login() + print("Opening browser for authentication...") + print("Please come back to the terminal after logging in.") + time.sleep(1) + _AuthServer(self._model).login_with_browser(self) def authenticate(model: str, shall_confirm: bool = True) -> None: auth = _AuthLitAI(model, shall_confirm) - with browser_login(False): - try: - auth.authenticate() - except ValueError: - pass - else: - return - auth._run_server() + auth.authenticate() class LLMApi: diff --git a/python/lightning_sdk/cli/entrypoint.py b/python/lightning_sdk/cli/entrypoint.py index 313636915..3cb4d8ee7 100644 --- a/python/lightning_sdk/cli/entrypoint.py +++ b/python/lightning_sdk/cli/entrypoint.py @@ -2,10 +2,7 @@ import os import sys -from contextlib import suppress -from typing import Optional -import requests import rich_click as click from lightning_sdk import __version__ @@ -48,9 +45,13 @@ from lightning_sdk.cli.lockdown import apply_lockdown from lightning_sdk.cli.ls import ls from lightning_sdk.cli.rm import rm -from lightning_sdk.cli.utils.auth import browser_authentication +from lightning_sdk.cli.utils.auth import ( + browser_authentication, + login_with_pasted_code, + logout_session, + oidc_login_enabled, +) from lightning_sdk.cli.utils.logging import CommandLoggingGroup, logging_excepthook -from lightning_sdk.lightning_cloud import cli_login from lightning_sdk.lightning_cloud.login import Auth from lightning_sdk.utils.resolve import _get_authed_user, in_studio @@ -103,22 +104,13 @@ def main_cli(ctx: click.Context) -> None: @click.option( "--no-browser", is_flag=True, + hidden=True, help="Sign in without a browser on this machine: approve on another device, then paste the code it shows.", ) -@click.option( - "--org", - "orgs", - metavar="A,B", - help="Comma-separated organizations to ask for when the login is approved.", -) -@click.option("--project", help="Teamspace to ask for when the login is approved.") -@click.option( - "--insecure-storage", - is_flag=True, - help="Store the login in a file only you can read when no OS keychain is available.", -) -def login(no_browser: bool, orgs: Optional[str], project: Optional[str], insecure_storage: bool) -> None: +def login(no_browser: bool) -> None: """Sign in to Lightning AI.""" + if no_browser and not oidc_login_enabled(): + raise click.UsageError("`lightning login --no-browser` is not available yet.") auth = Auth() if auth.api_key or auth.auth_token: click.echo('"lightning login" is not required within a Studio or when credentials are set in the environment') @@ -128,18 +120,21 @@ def login(no_browser: bool, orgs: Optional[str], project: Optional[str], insecur # this is unexpected, as we automatically auth within a Studio raise RuntimeError("Unable to login within a Studio. Did you change your shell setup?") from None - org_list = [org.strip() for org in orgs.split(",") if org.strip()] if orgs else None + # always start a fresh session: a saved one may have expired + auth.clear() try: - cli_login.login( - no_browser=no_browser, - insecure_storage=insecure_storage, - orgs=org_list, - project=project, - ) - except cli_login.LoginError: - raise + if no_browser: + login_with_pasted_code(auth) + else: + with browser_authentication(True): + auth.authenticate() except ConnectionError: raise RuntimeError(f"Unable to connect to {_cloud_url()}. Please check your internet connection.") from None + except ValueError: + # the browser flow ended without the callback saving a token (see the warning it logs) + raise RuntimeError( + "Login failed: the browser sign-in finished without returning a token. Try running `lightning login` again" + ) from None try: click.echo(f'Successfully logged in as "{_get_authed_user().name}".') @@ -148,32 +143,13 @@ def login(no_browser: bool, orgs: Optional[str], project: Optional[str], insecur click.echo("Successfully logged in.") -def _saved_session_token() -> Optional[str]: - """A `sess_` token saved by an older SDK, so logout can revoke it too.""" - auth = Auth() - try: - if auth.load() and isinstance(auth.auth_token, str) and auth.auth_token.startswith("sess_"): - return auth.auth_token - except Exception: - pass - return None - - @main_cli.command def logout() -> None: """Sign out of the current session.""" - legacy = _saved_session_token() try: - cli_login.logout() + logout_session(Auth()) except Exception as e: raise RuntimeError(f"Logout failed: {e}") from e - if legacy: - with suppress(requests.RequestException): - requests.post( - f"{_cloud_url().rstrip('/')}/v1/auth/logout", - headers={"Authorization": f"Bearer {legacy}"}, - timeout=10, - ) click.echo("Successfully logged out.") diff --git a/python/lightning_sdk/cli/legacy/entrypoint.py b/python/lightning_sdk/cli/legacy/entrypoint.py index 6e9cdc6c4..f36108b72 100644 --- a/python/lightning_sdk/cli/legacy/entrypoint.py +++ b/python/lightning_sdk/cli/legacy/entrypoint.py @@ -30,7 +30,7 @@ from lightning_sdk.cli.utils.auth import browser_authentication from lightning_sdk.cli.utils.coloring import CustomHelpFormatter from lightning_sdk.constants import _LIGHTNING_DEBUG -from lightning_sdk.lightning_cloud import cli_login +from lightning_sdk.lightning_cloud.login import Auth def _notify_exception(exception_type: Type[BaseException], value: BaseException, tb: Optional[TracebackType]) -> None: @@ -73,10 +73,12 @@ def main_cli(ctx: click.Context) -> None: @main_cli.command def login() -> None: """Login to Lightning AI Studios.""" + auth = Auth() + auth.clear() + try: - cli_login.login() - except cli_login.LoginError: - raise + with browser_authentication(True): + auth.authenticate() except ConnectionError: raise RuntimeError(f"Unable to connect to {_cloud_url()}. Please check your internet connection.") from None @@ -84,7 +86,8 @@ def login() -> None: @main_cli.command def logout() -> None: """Logout from Lightning AI Studios.""" - cli_login.logout() + auth = Auth() + auth.clear() # additional commands diff --git a/python/lightning_sdk/cli/utils/auth.py b/python/lightning_sdk/cli/utils/auth.py index 8555a328d..5de4c68fd 100644 --- a/python/lightning_sdk/cli/utils/auth.py +++ b/python/lightning_sdk/cli/utils/auth.py @@ -1,27 +1,384 @@ -"""Authentication helpers shared by the CLI commands.""" +"""Non-interactive authentication helpers for CLI commands.""" -from contextlib import contextmanager -from typing import Generator +import base64 +import hashlib +import hmac +import json +import os +import secrets +from contextlib import contextmanager, suppress +from contextvars import ContextVar +from pathlib import Path +from typing import Generator, Optional, Tuple +from urllib.parse import urlencode +import requests import rich_click as click -from lightning_sdk.lightning_cloud.credentials import browser_login -from lightning_sdk.lightning_cloud.login import Auth +from lightning_sdk.lightning_cloud import env +from lightning_sdk.lightning_cloud.login import Auth, AuthServer, logger + +_BROWSER_AUTH_ALLOWED = ContextVar("lightning_browser_auth_allowed", default=True) +_ORIGINAL_AUTH_SAVE = Auth.save + +_CLIENT_ID = "lightning-cli" +_PASTE_ATTEMPTS = 3 + +# The authorization-code login needs GRID_AUTH_CLI_LOGIN_ENABLED on the control plane, which production +# doesn't set yet. Until it does, `lightning login` uses the session-token login, and this opts in to test it. +OIDC_LOGIN_ENV = "LIGHTNING_CLI_OIDC_LOGIN" + + +def oidc_login_enabled() -> bool: + return env.truthy(os.getenv(OIDC_LOGIN_ENV)) + + +def _generate_state() -> str: + return base64.urlsafe_b64encode(os.urandom(16)).rstrip(b"=").decode("ascii") + + +def _origin() -> str: + """Platform origin without a trailing slash, so paths are not `//cli-login`.""" + return env.LIGHTNING_CLOUD_URL.rstrip("/") + + +def _pkce_pair() -> Tuple[str, str]: + """Return a PKCE `(verifier, S256 challenge)` pair.""" + verifier = secrets.token_urlsafe(64) + challenge = base64.urlsafe_b64encode(hashlib.sha256(verifier.encode("ascii")).digest()).rstrip(b"=") + return verifier, challenge.decode("ascii") + + +def _authorize_url(redirect_uri: str, state: str, challenge: str) -> str: + query = urlencode( + { + "response_type": "code", + "client_id": _CLIENT_ID, + "redirect_uri": redirect_uri, + "scope": "openid", + "state": state, + "code_challenge": challenge, + "code_challenge_method": "S256", + } + ) + return f"{_origin()}/oidc/authorize?{query}" + + +class _CodeRejectedError(Exception): + """The token endpoint refused the code; `str(e)` is its OAuth error.""" + + +def _exchange_code(code: str, verifier: str, redirect_uri: str) -> str: + """Trade an approved code and its PKCE verifier for a CLI session token.""" + try: + response = requests.post( + f"{_origin()}/oidc/token", + data={ + "grant_type": "authorization_code", + "client_id": _CLIENT_ID, + "code": code, + "code_verifier": verifier, + "redirect_uri": redirect_uri, + }, + timeout=10, + ) + except requests.ConnectionError as e: + raise ConnectionError(str(e)) from e + body = _json_body(response) + if response.ok and body.get("access_token"): + return body["access_token"] + raise _CodeRejectedError(body.get("error") or f"HTTP {response.status_code}") + + +def _save_with_session_token( + self: Auth, + token: str = "", + user_id: str = "", + api_key: str = "", + auth_token: str = "", +) -> None: + """Persist session tokens (`token`) as `auth_token` so the SDK/CLI sends `Bearer sess_...`.""" + if not auth_token and token and (not api_key or token.startswith("sess_")): + auth_token = token + if token.startswith("sess_"): + api_key = "" + _ORIGINAL_AUTH_SAVE( + self, + token=token, + user_id=user_id, + api_key=api_key, + auth_token=auth_token, + ) + with suppress(OSError): + os.chmod(self.secrets_file, 0o600) + + +def _login_with_authorization_code(self: AuthServer, auth: Auth) -> None: + """Sign in with an authorization code and PKCE, delivered to a callback on 127.0.0.1.""" + import webbrowser + + import uvicorn + from fastapi import FastAPI + from starlette.background import BackgroundTask + from starlette.responses import RedirectResponse + + from lightning_sdk.lightning_cloud.utils.network import find_free_network_port + + expected_state = _generate_state() + verifier, challenge = _pkce_pair() + app = FastAPI() + port = find_free_network_port() + redirect_uri = f"http://127.0.0.1:{port}/callback" + url = _authorize_url(redirect_uri, expected_state, challenge) + + try: + requests.head(url, timeout=10) + except requests.ConnectionError as e: + raise requests.ConnectionError( + f"No internet connection available. Please connect to a stable internet connection \n{e}" + ) from e + except requests.RequestException as e: + raise requests.RequestException( + f"An error occurred with the request. Please report this issue to Lightning Team \n{e}" + ) from e + + logger.info(f"login started for lightning.ai, opening {url}") + if not webbrowser.open(url): + if os.environ.get("LIGHTNING_DEPLOYMENT_ID"): + raise RuntimeError( + "Failed to authenticate to Lightning. Ensure that you have selected 'Include SDK credentials' in the " + "'Environment' section of the deployment settings." + ) + click.echo("No browser could be opened on this machine.") + login_with_pasted_code(auth) + return + + @app.get("/callback") + def save_token(code: str = "", state: str = "", error: str = "") -> RedirectResponse: + failure = "" + if not state or not hmac.compare_digest(state, expected_state): + failure = "state mismatch" + elif error or not code: + failure = error or "no code in the callback" + else: + try: + auth.save(auth_token=_exchange_code(code, verifier, redirect_uri)) + except (_CodeRejectedError, ConnectionError) as e: + failure = str(e) + + if failure: + logger.warning(f"Authentication Failed ({failure}). Run `lightning login` again, or update the CLI.") + target = f"{_origin()}/cli-login-failed" + else: + logger.info("Authentication Successful") + target = f"{_origin()}/cli-login-successful" + + return RedirectResponse( + url=target, + background=BackgroundTask(stop_server), + ) + + def stop_server() -> None: + server.should_exit = True + + server = uvicorn.Server(config=uvicorn.Config(app, host="127.0.0.1", port=port, log_level="error")) + server.run() + + +def _login_with_session_token(self: AuthServer, auth: Auth) -> None: + """Sign in through the sign-in page, which hands a session token to a callback on localhost.""" + import webbrowser + + import uvicorn + from fastapi import FastAPI, Query + from starlette.background import BackgroundTask + from starlette.responses import RedirectResponse + + from lightning_sdk.lightning_cloud.utils.network import find_free_network_port + + expected_state = _generate_state() + app = FastAPI() + port = find_free_network_port() + base_url = self.get_auth_url(port) + redirect_uri = f"http://localhost:{port}/login-complete" + if urlencode({"redirectTo": redirect_uri}) in base_url: + stateful_redirect = f"{redirect_uri}?{urlencode({'state': expected_state})}" + url = base_url.replace(urlencode({"redirectTo": redirect_uri}), urlencode({"redirectTo": stateful_redirect}), 1) + else: + url = base_url + + try: + requests.head(url, timeout=10) + except requests.ConnectionError as e: + raise requests.ConnectionError( + f"No internet connection available. Please connect to a stable internet connection \n{e}" + ) from e + except requests.RequestException as e: + raise requests.RequestException( + f"An error occurred with the request. Please report this issue to Lightning Team \n{e}" + ) from e + + logger.info(f"login started for lightning.ai, opening {url}") + if not webbrowser.open(url): + deployment_id = os.environ.get("LIGHTNING_DEPLOYMENT_ID") + if deployment_id: + raise RuntimeError( + "Failed to authenticate to Lightning. Ensure that you have selected 'Include SDK credentials' in the " + "'Environment' section of the deployment settings." + ) + raise RuntimeError( + "Failed to authenticate to Lightning. When running without access to a browser, " + "'LIGHTNING_API_KEY' should be exported." + ) + + @app.get("/login-complete") + async def save_token( + token: str = "", + key: str = "", + state: str = "", + user_id: str = Query("", alias="userID"), + ) -> RedirectResponse: + state_valid = bool(state) and hmac.compare_digest(state, expected_state) + if token and state_valid: + auth.save(token=token, user_id=user_id, api_key=key) + logger.info("Authentication Successful") + target = f"{env.LIGHTNING_CLOUD_URL}/cli-login-successful" + else: + reason = "no token in the callback" if not token else "state mismatch" + logger.warning( + f"Authentication Failed ({reason}). If the browser was already signed in, the sign-in page " + "could not mint a CLI token. Run `lightning login` again, or update the CLI." + ) + target = f"{env.LIGHTNING_CLOUD_URL}/cli-login-failed" + + return RedirectResponse( + url=target, + background=BackgroundTask(stop_server), + ) + + def stop_server() -> None: + server.should_exit = True + + server = uvicorn.Server(config=uvicorn.Config(app, port=port, log_level="error")) + server.run() + + +def _login_with_browser(self: AuthServer, auth: Auth) -> None: + if oidc_login_enabled(): + _login_with_authorization_code(self, auth) + else: + _login_with_session_token(self, auth) + + +Auth.save = _save_with_session_token +AuthServer.login_with_browser = _login_with_browser + + +def login_with_pasted_code(auth: Auth) -> None: + """Sign in on a machine without a browser: approve on any device, then paste the code shown there.""" + redirect_uri = f"{_origin()}/cli-login" + verifier, challenge = _pkce_pair() + click.echo( + "To sign in, open this address in a browser on any device and approve the request:\n\n" + f" {_authorize_url(redirect_uri, _generate_state(), challenge)}\n\n" + "The page then shows a code. Only paste a code from a page you opened this way; " + "never give it to anyone, as it gives full access to your Lightning account." + ) + for attempt in range(_PASTE_ATTEMPTS): + code = click.prompt("Code").strip() + try: + auth.save(auth_token=_exchange_code(code, verifier, redirect_uri)) + return + except _CodeRejectedError as e: + if str(e) != "invalid_grant": + raise RuntimeError(f"Login failed: {e}.") from None + if attempt < _PASTE_ATTEMPTS - 1: + click.echo("That code was not accepted. It works once and expires after a minute. Try again.") + raise RuntimeError("Login failed: the code was not accepted. Run `lightning login --no-browser` again.") + + +def _json_body(response: requests.Response) -> dict: + try: + body = response.json() + except ValueError: + return {} + return body if isinstance(body, dict) else {} + + +def logout_session(auth: Auth) -> None: + """Revoke the active CLI session token on the server and clear local credentials.""" + token = getattr(auth, "auth_token", None) + if not isinstance(token, str) or not token: + secrets_file = getattr(auth, "secrets_file", None) + if isinstance(secrets_file, Path) and secrets_file.exists(): + try: + with secrets_file.open() as creds: + data = json.load(creds) + loaded = data.get("auth_token") + if isinstance(loaded, str) and loaded: + token = loaded + except Exception: + token = None + if isinstance(token, str) and token: + with suppress(requests.RequestException): + requests.post( + f"{_origin()}/v1/auth/logout", + headers={"Authorization": f"Bearer {token}"}, + timeout=10, + ) + auth.clear() @contextmanager def browser_authentication(allowed: bool) -> Generator[None, None, None]: """Temporarily control whether missing credentials may start browser authentication.""" - with browser_login(allowed): + token = _BROWSER_AUTH_ALLOWED.set(allowed) + authenticate = Auth.authenticate + + def guarded_authenticate(auth: Auth) -> Optional[str]: + if _BROWSER_AUTH_ALLOWED.get(): + return authenticate(auth) + + has_own_run_server = "_run_server" in auth.__dict__ + run_server = auth.__dict__.get("_run_server") + + def blocked_run_server() -> None: + raise ValueError("No Lightning credentials are available. Run `lightning login` first.") + + auth.__dict__["_run_server"] = blocked_run_server + try: + return authenticate(auth) + finally: + if has_own_run_server: + auth.__dict__["_run_server"] = run_server + else: + auth.__dict__.pop("_run_server", None) + + Auth.authenticate = guarded_authenticate + try: yield + finally: + Auth.authenticate = authenticate + _BROWSER_AUTH_ALLOWED.reset(token) def require_auth_header() -> str: """Return available credentials without falling back to browser authentication.""" + auth = Auth() try: - header = Auth().auth_header - except (ValueError, ConnectionError) as err: - raise click.UsageError(str(err)) from err - if not header: - raise click.UsageError("No Lightning credentials are available. Run `lightning login` first.") - return header + header = auth.auth_header + except ValueError: + header = None + + if header: + return header + + if auth.load(): + try: + header = auth.auth_header + except ValueError: + header = None + if header: + return header + + raise click.UsageError("No Lightning credentials are available. Run `lightning login` first.") diff --git a/python/lightning_sdk/lightning_cloud/cli_login.py b/python/lightning_sdk/lightning_cloud/cli_login.py index fa03ec7d4..d681ef11b 100644 --- a/python/lightning_sdk/lightning_cloud/cli_login.py +++ b/python/lightning_sdk/lightning_cloud/cli_login.py @@ -18,7 +18,7 @@ import time import webbrowser from datetime import datetime -from typing import Any, Dict, Iterator, Optional, Sequence +from typing import Any, Callable, Dict, Iterator, Optional, Sequence, Tuple from urllib.parse import parse_qs, urlencode, urlsplit import requests @@ -27,16 +27,25 @@ from lightning_sdk.lightning_cloud.login import Auth CLIENT_ID = "lightning-cli" +# Production serves this login only from relaunch, so until then a new login +# is the interim one unless this is set. +OIDC_LOGIN_ENV = "LIGHTNING_CLI_OIDC_LOGIN" _KEYRING_SERVICE = CLIENT_ID _REFRESH_MARGIN = 300.0 _PASTE_TRIES = 3 +_CALLBACK_TIMEOUT = 300.0 _NOT_LOGGED_IN = "You are not logged in. Run `lightning login`." _LOGIN_ENDED = "Your login has expired or was revoked. Run `lightning login`." +_REUSED = "refresh token was reused; the login is revoked" _LOGIN_REUSED = ( "This login was revoked because a copy of it was used elsewhere. " "Run `lightning login` and consider rotating any keys this machine held." ) +_BROWSER_TIMED_OUT = ( + "The browser didn't finish the login within 5 minutes, so nothing was saved. " + "Run `lightning login` again, or `lightning login --no-browser` to paste a code instead." +) _NO_KEYCHAIN = ( "No OS keychain is available to store the login. Run " "`lightning login --insecure-storage` to keep it in a file only you can read." @@ -81,6 +90,11 @@ class Login(LoginStatus): refresh_token: str = dataclasses.field(repr=False) +def oidc_login_enabled() -> bool: + """Whether a new login uses this login rather than the interim one, read now.""" + return env.truthy(os.environ.get(OIDC_LOGIN_ENV)) + + def login( no_browser: bool = False, insecure_storage: bool = False, @@ -130,21 +144,6 @@ def current_login() -> Optional[LoginStatus]: ) -def logout() -> None: - """Revoke the saved login on the server, then forget it on this machine.""" - with _locked(): - login = _load(_origin(), cached=False) - revoked = login is None or _revoke(login.refresh_token) - if login is not None: - _forget() - Auth.secrets_file.unlink(missing_ok=True) - if not revoked: - raise LoginError( - "Logged out on this machine, but Lightning couldn't be reached to " - "revoke the login. It may stay active until it expires." - ) - - def _request_token(rejected: Optional[str] = None) -> Optional[str]: """The saved login's access token for a request, or None without a login. @@ -198,18 +197,27 @@ def complete(self, code: str, redirect_uri: str) -> Login: "storage": self.storage, }) new = _login_from(body, self.storage, self.origin) - try: - with _locked(): + + def save() -> None: + try: _save(new) - except LoginError: - _revoke(new.refresh_token) - raise - Auth.secrets_file.unlink(missing_ok=True) + except LoginError: + _revoke(new.refresh_token) + raise + Auth.secrets_file.unlink(missing_ok=True) + + _replace_saved(save) return new def _browser_login(attempt: _Attempt) -> Optional[Login]: - with _callback_server(attempt) as (redirect_uri, wait): + redirect_uri = "" + + def complete(query: Dict[str, str]) -> Login: + return _callback_login(attempt, query, redirect_uri) + + with _callback_server("/callback", complete, _BROWSER_TIMED_OUT) as (port, wait): + redirect_uri = f"http://127.0.0.1:{port}/callback" url = attempt.authorize_url(redirect_uri) if not webbrowser.open(url): return None @@ -221,46 +229,56 @@ def _browser_login(attempt: _Attempt) -> Optional[Login]: @contextlib.contextmanager -def _callback_server(attempt: _Attempt) -> Iterator[Any]: +def _callback_server( + path: str, complete: Callable[[Dict[str, str]], Any], timed_out: str +) -> Iterator[Tuple[int, Callable[[], Any]]]: + """Serve one browser callback on 127.0.0.1, which ``complete`` turns into a login. + + ``wait()`` returns what ``complete`` returned, raises what it raised, or + raises ``timed_out`` when no callback came within ``_CALLBACK_TIMEOUT``. + """ import http.server import socketserver + origin = _origin() outcome = [] class Callback(http.server.BaseHTTPRequestHandler): def do_GET(self) -> None: parts = urlsplit(self.path) - if parts.path != "/callback": + if parts.path != path: self.send_error(404) return query = {key: values[0] for key, values in parse_qs(parts.query).items()} try: - outcome.append(_callback_login(attempt, query, redirect_uri)) + outcome.append(complete(query)) page = "cli-login-successful" except Exception as err: outcome.append(err) page = "cli-login-failed" self.send_response(302) - self.send_header("Location", f"{attempt.origin}/{page}") + self.send_header("Location", f"{origin}/{page}") self.end_headers() def log_message(self, format: str, *args: Any) -> None: - # The default logs the request line, which carries the code. + # The default logs the request line, which carries the code or token. pass with socketserver.TCPServer(("127.0.0.1", 0), Callback) as server: # A short timeout lets Ctrl-C interrupt the wait. server.timeout = 0.5 - redirect_uri = f"http://127.0.0.1:{server.server_address[1]}/callback" + deadline = time.time() + _CALLBACK_TIMEOUT - def wait() -> Login: + def wait() -> Any: while not outcome: + if time.time() >= deadline: + raise LoginError(timed_out) server.handle_request() if isinstance(outcome[0], Exception): raise outcome[0] return outcome[0] - yield redirect_uri, wait + yield server.server_address[1], wait def _callback_login(attempt: _Attempt, query: Dict[str, str], redirect_uri: str) -> Login: @@ -321,7 +339,7 @@ def _fresh_login(url: str, min_remaining: float, rejected: Optional[str] = None) raise except _Refused as err: _forget() - message = _LOGIN_REUSED if "reused" in err.description else _LOGIN_ENDED + message = _LOGIN_REUSED if err.description == _REUSED else _LOGIN_ENDED raise LoginError(message) from None fresh = _login_from(body, login.storage, url) _save(fresh) @@ -600,6 +618,47 @@ def _revoke(token: str) -> bool: return False +def _revoke_session(token: str) -> bool: + try: + response = requests.post( + f"{_origin()}/v1/auth/logout", + headers={"Authorization": f"Bearer {token}", "User-Agent": env.USER_AGENT}, + timeout=(env.LIGHTNING_CLOUD_CONNECT_TIMEOUT, env.LIGHTNING_CLOUD_READ_TIMEOUT), + ) + except requests.RequestException: + return False + # A 401 says the session already doesn't work, which is all a revoke is for. + return response.status_code in (200, 401) + + +def _saved_session() -> Optional[str]: + """The interim ``sess_`` login saved as ``auth_token`` in credentials.json.""" + try: + data = json.loads(Auth.secrets_file.read_text()) + except (OSError, ValueError): + return None + token = data.get("auth_token") if isinstance(data, dict) else None + if not isinstance(token, str) or not token.strip(): + return None + return token.strip() + + +def _replace_saved(save: Callable[[], Optional[str]]) -> None: + """Save a new login with ``save``, then revoke the logins it replaced. + + ``save`` returns the session token it saved, if it saved one. A ``save`` + that raises leaves the old logins saved and live. + """ + with _locked(): + replaced_login = _load(_origin(), cached=False) + replaced_session = _saved_session() + kept_session = save() + if replaced_login is not None: + _revoke(replaced_login.refresh_token) + if replaced_session is not None and replaced_session != kept_session: + _revoke_session(replaced_session) + + def _login_from(body: Dict[str, Any], storage: str, url: str) -> Login: try: return Login( diff --git a/python/lightning_sdk/lightning_cloud/credentials.py b/python/lightning_sdk/lightning_cloud/credentials.py index 7ac1090e8..71e08c7c4 100644 --- a/python/lightning_sdk/lightning_cloud/credentials.py +++ b/python/lightning_sdk/lightning_cloud/credentials.py @@ -4,12 +4,14 @@ import contextlib import contextvars +import dataclasses +import enum import os import pathlib -from typing import Iterator, Optional +from typing import Any, Callable, Dict, Iterator, Optional, Sequence, Tuple from lightning_sdk.lightning_cloud import cli_login, env -from lightning_sdk.lightning_cloud.login import Auth, Keys +from lightning_sdk.lightning_cloud.login import Auth, AuthServer, Keys _MISSING_CREDENTIALS = ( "We couldn't find any credentials linked to your account. " @@ -22,6 +24,23 @@ ) +class CredentialKind(str, enum.Enum): + TOKEN_FILE = "token_file" + ENV_TOKEN = "env_token" + LOGIN = "login" + SESSION = "session" + + +@dataclasses.dataclass(frozen=True) +class Credential: + kind: CredentialKind + token: str = dataclasses.field(repr=False) + + +class LoginUnavailable(ValueError): + """A login option the interim login can't honour.""" + + @contextlib.contextmanager def browser_login(allowed: bool) -> Iterator[None]: """Allow or forbid opening a browser to log in when no credential is found.""" @@ -41,77 +60,147 @@ def _present(value: Optional[str]) -> Optional[str]: return text +def _from_token_file() -> Optional[str]: + path = os.environ.get(env.LIGHTNING_AUTH_TOKEN_FILE) + if not path: + return None + try: + raw = pathlib.Path(path).read_text() + except FileNotFoundError: + return None + return _present(raw) + + +def _from_env_token() -> Optional[str]: + return _present(os.environ.get(Keys.AUTH_TOKEN.value)) + + +# The order is the precedence. LIGHTNING_API_KEY and a saved api_key have no +# row, because the control plane accepts a token, never a legacy key. +_READERS = ( + (CredentialKind.TOKEN_FILE, _from_token_file), + (CredentialKind.ENV_TOKEN, _from_env_token), + (CredentialKind.LOGIN, cli_login._request_token), + (CredentialKind.SESSION, cli_login._saved_session), +) + + +def _login_through_auth() -> None: + # An SDK CLI may turn browser login off by patching Auth.authenticate, so + # the source's own login goes through it rather than calling start_login. + Auth().authenticate() + + class CredentialSource: - def __init__(self) -> None: - self._may_open_browser = True - self._login_token: Optional[str] = None + """Works out the credential each request sends, when it is sent.""" + + def __init__(self, login: Optional[Callable[[], object]] = None) -> None: + self._login = login or _login_through_auth + # None until the source has tried its one login, then what to say + # when a request still finds no credential. + self._refusal: Optional[str] = None + + def resolve(self) -> Optional[Credential]: + """The credential a request would send now. Never opens a browser.""" + for kind, read in _READERS: + token = read() + if token: + return Credential(kind, token) + return None def header(self) -> str: - token = self._token() - if token is None and self._may_open_browser: + """``Bearer `` for a request, logging in once when nothing is saved.""" + credential = self._resolve() + if credential is None and self._refusal is None: if not _browser_login_allowed.get(): raise ValueError(_LOGIN_FIRST) - self._may_open_browser = False - Auth()._run_server() - token = self._token() - if not token: - raise ValueError(_MISSING_CREDENTIALS) - return f"Bearer {token}" - - def refresh(self) -> bool: - """Refresh the login after the server refused its access token. - - Only the saved login can refresh, and only when its token was the one - last handed out. - """ - rejected = self._login_token - if rejected is None: + self._refusal = _MISSING_CREDENTIALS + try: + self._login() + except ValueError as err: + self._refusal = str(err) + raise + credential = self._resolve() + if credential is None: + raise ValueError(self._refusal) + return f"Bearer {credential.token}" + + def refresh(self, token: str) -> bool: + """After the server refused ``token``, whether a different one is now available.""" + current = self._resolve() + if current is None: return False - fresh = self._saved_login_token(rejected) - return fresh is not None and fresh != rejected - - def _token(self) -> Optional[str]: - self._login_token = None - for read in ( - self._from_token_file, - self._from_env_token, - self._from_saved_login, - ): - token = read() - if token: - return token - return None + if current.kind is CredentialKind.LOGIN and current.token == token: + fresh = self._noting_refusal(lambda: cli_login._request_token(token)) + return fresh is not None and fresh != token + return current.token != token - def _from_token_file(self) -> Optional[str]: - path = os.environ.get(env.LIGHTNING_AUTH_TOKEN_FILE) - if not path: - return None - try: - raw = pathlib.Path(path).read_text() - except FileNotFoundError: - return None - return _present(raw) - - def _from_env_token(self) -> Optional[str]: - return _present(os.environ.get(Keys.AUTH_TOKEN.value)) - - def _from_saved_login(self) -> Optional[str]: - self._login_token = self._saved_login_token() - if self._login_token: - return self._login_token - auth = Auth() - if not auth.load(): - return None - token = _present(auth.auth_token) - if token: - return token - return _present(auth.api_key) - - def _saved_login_token(self, rejected: Optional[str] = None) -> Optional[str]: + def _resolve(self) -> Optional[Credential]: + return self._noting_refusal(self.resolve) + + def _noting_refusal(self, read: Callable[[], Any]) -> Any: try: - return cli_login._request_token(rejected) + return read() except cli_login.LoginError: # Every process sharing a refused login would otherwise open a # browser. Saying "run `lightning login`" once is enough. - self._may_open_browser = False + self._refusal = _MISSING_CREDENTIALS raise + + +def start_login( + no_browser: bool = False, + insecure_storage: bool = False, + orgs: Optional[Sequence[str]] = None, + project: Optional[str] = None, + sign_in_params: Optional[Dict[str, str]] = None, +) -> Optional[cli_login.Login]: + """Log in with the login ``LIGHTNING_CLI_OIDC_LOGIN`` selects when this is called. + + The new login returns the ``Login`` and ignores ``sign_in_params``. The + interim login returns None, and raises ``LoginUnavailable`` for any other + option, because it has no way to honour one. + """ + if cli_login.oidc_login_enabled(): + return cli_login.login( + no_browser=no_browser, + insecure_storage=insecure_storage, + orgs=orgs, + project=project, + ) + unavailable = [ + flag + for flag, given in ( + ("--no-browser", no_browser), + ("--insecure-storage", insecure_storage), + ("--org", orgs), + ("--project", project), + ) + if given + ] + if unavailable: + raise LoginUnavailable( + f"{', '.join(unavailable)} needs the new login. " + f"Set {cli_login.OIDC_LOGIN_ENV}=1 to use it." + ) + AuthServer(sign_in_params).login_with_browser(Auth()) + return None + + +def logout() -> Tuple[CredentialKind, ...]: + """Forget every saved login on this machine and revoke each on the server. + + Returns the kinds whose revoke failed, which may still work until they expire. + """ + with cli_login._locked(): + login = cli_login._load(cli_login._origin(), cached=False) + session = cli_login._saved_session() + if login is not None: + cli_login._forget() + Auth.secrets_file.unlink(missing_ok=True) + failed = [] + if login is not None and not cli_login._revoke(login.refresh_token): + failed.append(CredentialKind.LOGIN) + if session is not None and not cli_login._revoke_session(session): + failed.append(CredentialKind.SESSION) + return tuple(failed) diff --git a/python/lightning_sdk/lightning_cloud/login.py b/python/lightning_sdk/lightning_cloud/login.py index 2fea997bf..ee77cca76 100644 --- a/python/lightning_sdk/lightning_cloud/login.py +++ b/python/lightning_sdk/lightning_cloud/login.py @@ -3,12 +3,17 @@ # Make changes in the internal upstream repository instead. from enum import Enum +import hmac import json import logging import os import pathlib +import secrets from dataclasses import dataclass -from typing import Optional, Literal +from typing import Dict, Optional, Literal +from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit + +import webbrowser import requests @@ -70,18 +75,18 @@ def load(self) -> bool: def save( self, - user_id: Optional[str] = None, - api_key: Optional[str] = None, - auth_token: Optional[str] = None, + token: str = "", + user_id: str = "", + api_key: str = "", + auth_token: str = "", # username: str = "", ) -> None: - """Replace the saved login with the given fields.""" - given = { + """Replace the saved login with the given fields. ``token`` is ignored.""" + stored = { Keys.USER_ID.suffix: user_id, Keys.API_KEY.suffix: api_key, Keys.AUTH_TOKEN.suffix: auth_token, } - stored = {key: value for key, value in given.items() if value} self.secrets_file.parent.mkdir(exist_ok=True, parents=True) payload = json.dumps(stored).encode() @@ -163,10 +168,10 @@ def auth_header(self) -> Optional[str]: return self.get_auth_header() def _run_server(self) -> None: - """Log in through the browser, or by pasting a code when none opens.""" - from lightning_sdk.lightning_cloud.cli_login import login + """Start the login ``LIGHTNING_CLI_OIDC_LOGIN`` selects.""" + from lightning_sdk.lightning_cloud.credentials import start_login - login() + start_login() def authenticate(self) -> Optional[str]: """Performs end to end authentication flow. @@ -174,12 +179,12 @@ def authenticate(self) -> Optional[str]: Returns ---------- authorization header to use when authentication completes. With no - credential, this opens a browser to log in, unless browser login is - turned off. + credential, this logs in through ``_run_server``, unless browser login + is turned off. """ from lightning_sdk.lightning_cloud.credentials import CredentialSource - return CredentialSource().header() + return CredentialSource(login=self._run_server).header() def token_login(self, token_key: str, save_token: bool = True) -> Optional[str]: """Performs token-based authentication. @@ -211,6 +216,7 @@ def token_login(self, token_key: str, save_token: bool = True) -> Optional[str]: config = Configuration() config.host = env.LIGHTNING_CLOUD_URL api_client = ApiClient(configuration=config) + api_client.user_agent = env.USER_AGENT auth_api = AuthServiceApi(api_client) logger.debug(f"Attempting token login to {config.host}") @@ -287,6 +293,7 @@ def refresh_token(self, duration: int = 43200) -> Optional[str]: config.api_key['Authorization'] = self.auth_token api_client = ApiClient(configuration=config) + api_client.user_agent = env.USER_AGENT auth_api = AuthServiceApi(api_client) logger.debug(f"Attempting to refresh JWT token with duration {duration} seconds") @@ -346,9 +353,89 @@ def create_api_client(self, override: Optional[AuthOverride] = None) -> 'ApiClie # Create the API client client = ApiClient(configuration=config) + client.user_agent = env.USER_AGENT # Set the Authorization header directly in default_headers if auth_header: client.set_default_header('Authorization', auth_header) return client + + +_INTERIM_TIMED_OUT = ( + "The browser didn't finish the login within 5 minutes, so nothing was saved. " + "Run `lightning login` again." +) + + +class AuthServer: + """The interim browser login, which saves a ``sess_`` token as ``auth_token``.""" + + def __init__(self, sign_in_params: Optional[Dict[str, str]] = None) -> None: + self.sign_in_params = dict(sign_in_params or {}) + + def get_auth_url(self, port: int) -> str: + redirect_uri = f"http://localhost:{port}/login-complete" + params = urlencode(dict(redirectTo=redirect_uri, **self.sign_in_params)) + return f"{env.LIGHTNING_CLOUD_URL}/sign-in?{params}" + + def login_with_browser(self, auth: Auth) -> None: + from lightning_sdk.lightning_cloud import cli_login + + state = secrets.token_urlsafe(32) + + def complete(query: Dict[str, str]) -> None: + if not hmac.compare_digest(query.get("state", "").encode(), state.encode()): + raise cli_login.LoginError( + "The login callback didn't match this login, so nothing was saved. " + "Run `lightning login` again." + ) + token = query.get("token", "") + if not token: + raise cli_login.LoginError( + "The sign-in page sent no token, so nothing was saved. " + "Run `lightning login` again, or update the SDK." + ) + + def save() -> str: + auth.save(auth_token=token, user_id=query.get("userID", "")) + cli_login._forget() + return token + + cli_login._replace_saved(save) + logger.info("Authentication Successful") + + with cli_login._callback_server("/login-complete", complete, _INTERIM_TIMED_OUT) as (port, wait): + redirect_uri = f"http://localhost:{port}/login-complete" + url = _with_state(self.get_auth_url(port), redirect_uri, state) + try: + requests.head( + url, + headers={"User-Agent": env.USER_AGENT}, + timeout=env.LIGHTNING_CLOUD_CONNECT_TIMEOUT, + ) + except requests.RequestException as err: + raise cli_login.LoginError( + f"Couldn't reach {env.LIGHTNING_CLOUD_URL}. Check your connection and try again." + ) from err + logger.info(f"login started for lightning.ai, opening {url}") + if not webbrowser.open(url): + if os.environ.get("LIGHTNING_DEPLOYMENT_ID"): + raise cli_login.LoginError(cli_login._DEPLOYMENT) + raise cli_login.LoginError( + "No browser could be opened to log in. Run `lightning login` where a " + "browser is available, or set LIGHTNING_AUTH_TOKEN." + ) + wait() + + +def _with_state(url: str, redirect_uri: str, state: str) -> str: + # A subclass may build its own URL from get_auth_url(port), which knows no + # state, so the state goes onto whichever redirectTo names this callback. + parts = urlsplit(url) + stateful = f"{redirect_uri}?{urlencode(dict(state=state))}" + query = [ + (key, stateful if key == "redirectTo" and value == redirect_uri else value) + for key, value in parse_qsl(parts.query, keep_blank_values=True) + ] + return urlunsplit(parts._replace(query=urlencode(query))) diff --git a/python/lightning_sdk/lightning_cloud/rest_client.py b/python/lightning_sdk/lightning_cloud/rest_client.py index 4c12c9fae..78568b316 100644 --- a/python/lightning_sdk/lightning_cloud/rest_client.py +++ b/python/lightning_sdk/lightning_cloud/rest_client.py @@ -292,7 +292,10 @@ def wrap(*args, **kwargs): except ApiException as err: if err.status != 401: raise - if supplied or not source.refresh(): + if supplied: + raise ConnectionError(_LOGIN_MESSAGE) + sent = headers["Authorization"][len("Bearer "):] + if not source.refresh(sent): raise ConnectionError(_LOGIN_MESSAGE) headers["Authorization"] = source.header() try: diff --git a/python/tests/cli/test_entrypoint.py b/python/tests/cli/test_entrypoint.py index 260a90f69..ba8e6c302 100644 --- a/python/tests/cli/test_entrypoint.py +++ b/python/tests/cli/test_entrypoint.py @@ -9,8 +9,8 @@ from lightning_sdk.cli.entrypoint import login, main_cli from lightning_sdk.cli.sandbox.commands import run_sandbox_command +from lightning_sdk.cli.utils.auth import OIDC_LOGIN_ENV from lightning_sdk.cli.vm.ssh import ssh_vm -from lightning_sdk.lightning_cloud import cli_login from lightning_sdk.lightning_cloud.login import Auth from tests.cli.help import assert_help_contains, command_text, mock_command_logging, run_cli @@ -20,9 +20,10 @@ def _clear_auth(monkeypatch, tmp_path) -> None: monkeypatch.setattr(Auth, "secrets_file", tmp_path / "credentials.json") - monkeypatch.setenv("LIGHTNING_AUTH_TOKEN_FILE", str(tmp_path / "missing-token")) - for key in ("LIGHTNING_USER_ID", "LIGHTNING_API_KEY", "LIGHTNING_AUTH_TOKEN"): - monkeypatch.delenv(key, raising=False) + monkeypatch.delenv("LIGHTNING_USER_ID", raising=False) + monkeypatch.delenv("LIGHTNING_API_KEY", raising=False) + monkeypatch.delenv("LIGHTNING_AUTH_TOKEN", raising=False) + monkeypatch.delenv(OIDC_LOGIN_ENV, raising=False) @pytest.mark.parametrize( @@ -46,104 +47,110 @@ def test_resource_commands_never_start_browser_authentication(monkeypatch, tmp_p assert "Run `lightning login`" in str(result.exception) -def test_explicit_login_starts_the_login_flow(monkeypatch, tmp_path) -> None: +def test_explicit_login_can_start_browser_authentication(monkeypatch, tmp_path) -> None: _clear_auth(monkeypatch, tmp_path) - login = MagicMock() - monkeypatch.setattr("lightning_sdk.cli.entrypoint.cli_login.login", login) - monkeypatch.setattr("lightning_sdk.cli.entrypoint._get_authed_user", lambda: MagicMock(name="Test User")) + browser_auth_started = False + + def _complete_browser_auth(auth: Auth) -> None: + nonlocal browser_auth_started + browser_auth_started = True + auth.save(auth_token="sess_user-1.secret") + + monkeypatch.setattr(Auth, "_run_server", _complete_browser_auth) result = CliRunner().invoke(main_cli, ["login"]) assert result.exit_code == 0, result.output - login.assert_called_once_with(no_browser=False, insecure_storage=False, orgs=None, project=None) + assert browser_auth_started -def test_login_passes_every_option_to_the_login_flow(monkeypatch, tmp_path) -> None: +def test_login_reports_success_and_failure(monkeypatch, tmp_path) -> None: _clear_auth(monkeypatch, tmp_path) - login = MagicMock() - monkeypatch.setattr("lightning_sdk.cli.entrypoint.cli_login.login", login) - monkeypatch.setattr("lightning_sdk.cli.entrypoint._get_authed_user", lambda: MagicMock(name="Test User")) + user = MagicMock() + user.name = "Test User" + monkeypatch.setattr("lightning_sdk.cli.entrypoint._get_authed_user", lambda: user) - result = CliRunner().invoke( - main_cli, - ["login", "--no-browser", "--org", "acme, beta", "--project", "team", "--insecure-storage"], - ) + def _save_token(auth: Auth) -> None: + auth.save(token="sess_user-1.secret", user_id="user-1") - assert result.exit_code == 0, result.output - login.assert_called_once_with( - no_browser=True, - insecure_storage=True, - orgs=["acme", "beta"], - project="team", - ) + monkeypatch.setattr(Auth, "_run_server", _save_token) + ok = CliRunner().invoke(main_cli, ["login"]) + assert ok.exit_code == 0, ok.output + assert 'Successfully logged in as "Test User".' in ok.output + + # the browser callback finished without saving a token + _clear_auth(monkeypatch, tmp_path / "second") + monkeypatch.setattr(Auth, "_run_server", lambda _auth: None) + failed = CliRunner().invoke(main_cli, ["login"]) + assert failed.exit_code == 1 + assert "Login failed" in str(failed.exception) -def test_login_reports_success(monkeypatch, tmp_path) -> None: +def test_login_no_browser_starts_a_fresh_paste_login_instead_of_the_browser(monkeypatch, tmp_path) -> None: _clear_auth(monkeypatch, tmp_path) + monkeypatch.setenv(OIDC_LOGIN_ENV, "1") + Auth().save(auth_token="sess_expired") user = MagicMock() user.name = "Test User" monkeypatch.setattr("lightning_sdk.cli.entrypoint._get_authed_user", lambda: user) - monkeypatch.setattr("lightning_sdk.cli.entrypoint.cli_login.login", MagicMock()) + monkeypatch.setattr( + Auth, + "_run_server", + lambda _self: (_ for _ in ()).throw(AssertionError("browser auth reached")), + ) - result = CliRunner().invoke(main_cli, ["login"]) + def _approve(auth: Auth) -> None: + assert not Auth.secrets_file.exists() + auth.save(auth_token="sess_new") + + monkeypatch.setattr("lightning_sdk.cli.entrypoint.login_with_pasted_code", _approve) + + result = CliRunner().invoke(main_cli, ["login", "--no-browser"]) assert result.exit_code == 0, result.output assert 'Successfully logged in as "Test User".' in result.output + auth = Auth() + assert auth.load() + assert auth.auth_token == "sess_new" -def test_login_reports_a_refused_login(monkeypatch, tmp_path) -> None: +def test_login_no_browser_is_refused_without_the_opt_in(monkeypatch, tmp_path) -> None: _clear_auth(monkeypatch, tmp_path) + Auth().save(auth_token="sess_kept") monkeypatch.setattr( - "lightning_sdk.cli.entrypoint.cli_login.login", - MagicMock(side_effect=cli_login.LoginError("No OS keychain is available")), + "lightning_sdk.cli.entrypoint.login_with_pasted_code", + lambda _auth: (_ for _ in ()).throw(AssertionError("paste login reached")), ) - result = CliRunner().invoke(main_cli, ["login"]) + result = CliRunner().invoke(main_cli, ["login", "--no-browser"]) - assert result.exit_code == 1 - assert "No OS keychain is available" in str(result.exception) + assert result.exit_code == 2 + assert "not available yet" in result.output + auth = Auth() + assert auth.load() + assert auth.auth_token == "sess_kept" def test_logout_reports_success_and_failure(monkeypatch, tmp_path) -> None: _clear_auth(monkeypatch, tmp_path) - logout = MagicMock() - monkeypatch.setattr("lightning_sdk.cli.entrypoint.cli_login.logout", logout) - ok = CliRunner().invoke(main_cli, ["logout"]) assert ok.exit_code == 0, ok.output assert "Successfully logged out." in ok.output - logout.assert_called_once_with() - monkeypatch.setattr( - "lightning_sdk.cli.entrypoint.cli_login.logout", - MagicMock(side_effect=cli_login.LoginError("nobody home")), - ) + monkeypatch.setattr("lightning_sdk.cli.entrypoint.logout_session", MagicMock(side_effect=OSError("denied"))) failed = CliRunner().invoke(main_cli, ["logout"]) assert failed.exit_code == 1 - assert "Logout failed: nobody home" in str(failed.exception) - - -def test_logout_revokes_a_legacy_session(monkeypatch, tmp_path) -> None: - _clear_auth(monkeypatch, tmp_path) - Auth().save(auth_token="sess_old") - monkeypatch.setattr("lightning_sdk.cli.entrypoint.cli_login.logout", MagicMock()) - - with patch("lightning_sdk.cli.entrypoint.requests.post") as post: - result = CliRunner().invoke(main_cli, ["logout"]) - - assert result.exit_code == 0, result.output - post.assert_called_once() - assert post.call_args.kwargs["headers"] == {"Authorization": "Bearer sess_old"} + assert "Logout failed: denied" in str(failed.exception) def test_cli_browser_auth_policy_is_reset_after_command(monkeypatch, tmp_path) -> None: _clear_auth(monkeypatch, tmp_path) browser_auth_starts = 0 - def _complete_browser_auth(_auth: Auth) -> None: + def _complete_browser_auth(auth: Auth) -> None: nonlocal browser_auth_starts browser_auth_starts += 1 - monkeypatch.setenv("LIGHTNING_AUTH_TOKEN", "token") + auth.save(auth_token="sess_user-1.secret") monkeypatch.setattr(Auth, "_run_server", _complete_browser_auth) @@ -152,7 +159,7 @@ def _complete_browser_auth(_auth: Auth) -> None: assert result.exit_code == 1 assert "Run `lightning login`" in str(result.exception) assert browser_auth_starts == 0 - assert Auth().authenticate() == "Bearer token" + assert Auth().authenticate().startswith("Bearer ") assert browser_auth_starts == 1 @@ -310,10 +317,7 @@ def test_login_help(): result_text = assert_help_contains("lightning login --help", "Usage: lightning login [OPTIONS]") assert "Sign in to Lightning AI." in result_text - assert "--no-browser" in result_text - assert "--org" in result_text - assert "--project" in result_text - assert "--insecure-storage" in result_text + assert "--no-browser" not in result_text @mock_command_logging @@ -323,19 +327,40 @@ def test_logout_help(): assert "Sign out of the current session." in result_text +def test_login_replaces_existing_session(monkeypatch, tmp_path): + _clear_auth(monkeypatch, tmp_path) + Auth().save(auth_token="sess_expired") + + def _complete_browser_auth(auth: Auth) -> None: + assert not Auth.secrets_file.exists() + auth.save(auth_token="sess_new") + + monkeypatch.setattr(Auth, "_run_server", _complete_browser_auth) + + result = CliRunner().invoke(main_cli, ["login"]) + + assert result.exit_code == 0, result.output + auth = Auth() + assert auth.load() + assert auth.auth_token == "sess_new" + + @pytest.mark.parametrize("env_var", ["LIGHTNING_AUTH_TOKEN", "LIGHTNING_API_KEY"]) def test_login_not_required_with_credentials_in_env(monkeypatch, tmp_path, env_var): _clear_auth(monkeypatch, tmp_path) monkeypatch.setenv(env_var, "from-env") - monkeypatch.setattr( - "lightning_sdk.cli.entrypoint.cli_login.login", - MagicMock(side_effect=AssertionError("login should not start")), - ) + Auth().save(auth_token="sess_saved") + + def _fail_browser_auth(auth: Auth) -> None: + raise AssertionError("browser authentication should not start") + + monkeypatch.setattr(Auth, "_run_server", _fail_browser_auth) result = CliRunner().invoke(main_cli, ["login"]) assert result.exit_code == 0, result.output assert "credentials are set in the environment" in result.output + assert Auth.secrets_file.exists() @patch.dict(os.environ, {"LIGHTNING_INTERACTIVE": "true", "LIGHTNING_CLOUD_SPACE_ID": "cloud-space-id"}, clear=True) diff --git a/python/tests/cli/utils/test_auth.py b/python/tests/cli/utils/test_auth.py index 360b847e2..1292fec5e 100644 --- a/python/tests/cli/utils/test_auth.py +++ b/python/tests/cli/utils/test_auth.py @@ -1,43 +1,99 @@ +import base64 +import hashlib +import json from unittest.mock import MagicMock, PropertyMock, patch +from urllib.parse import parse_qs, urlparse import pytest import rich_click as click from lightning_sdk.cli.legacy.deploy._auth import _AuthLitServe, _AuthMode -from lightning_sdk.cli.utils.auth import browser_authentication, require_auth_header -from lightning_sdk.lightning_cloud.credentials import CredentialSource +from lightning_sdk.cli.utils.auth import ( + OIDC_LOGIN_ENV, + browser_authentication, + login_with_pasted_code, + logout_session, + require_auth_header, +) from lightning_sdk.lightning_cloud.login import Auth -def _no_credentials(monkeypatch, tmp_path) -> None: - for key in ("LIGHTNING_USER_ID", "LIGHTNING_API_KEY", "LIGHTNING_AUTH_TOKEN"): +def _response(status_code: int, body=None) -> MagicMock: + response = MagicMock(status_code=status_code, ok=status_code < 400) + response.json.return_value = body + return response + + +def _login_env(monkeypatch, tmp_path): + """Point credentials at tmp_path and return the secrets file.""" + secrets_file = tmp_path / "credentials.json" + monkeypatch.setattr(Auth, "secrets_file", secrets_file) + for key in ( + "LIGHTNING_USER_ID", + "LIGHTNING_API_KEY", + "LIGHTNING_AUTH_TOKEN", + "LIGHTNING_DEPLOYMENT_ID", + OIDC_LOGIN_ENV, + ): monkeypatch.delenv(key, raising=False) - monkeypatch.setenv("LIGHTNING_AUTH_TOKEN_FILE", str(tmp_path / "missing-token")) - monkeypatch.setattr(Auth, "secrets_file", tmp_path / "credentials.json") - monkeypatch.setattr( - Auth, - "_run_server", - lambda self: (_ for _ in ()).throw(AssertionError("browser authentication reached")), - ) + return secrets_file + +def _challenge(verifier: str) -> str: + return base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest()).rstrip(b"=").decode() -def test_require_auth_header_uses_the_credential_source() -> None: - """An available header must be returned without loading or authenticating.""" - auth = MagicMock(auth_header="Bearer token") + +def _capture_browser_login(monkeypatch, tmp_path, token_response) -> dict: + """Run the loopback login with a fake server and browser; return what was captured.""" + captured: dict = {"secrets_file": _login_env(monkeypatch, tmp_path)} + monkeypatch.setenv(OIDC_LOGIN_ENV, "1") + + class FakeServer: + def __init__(self, config) -> None: + captured["config"] = config + captured["route"] = next(r for r in config.app.routes if getattr(r, "path", None) == "/callback") + + def run(self) -> None: + captured["run"](captured["route"].endpoint) + + post = MagicMock(side_effect=token_response) + monkeypatch.setattr("lightning_sdk.cli.utils.auth._generate_state", lambda: "fixed-state") + monkeypatch.setattr("lightning_sdk.cli.utils.auth.requests.head", MagicMock()) + monkeypatch.setattr("lightning_sdk.cli.utils.auth.requests.post", post) + monkeypatch.setattr("webbrowser.open", lambda url: captured.setdefault("auth_url", url) and True) + monkeypatch.setattr("uvicorn.Server", FakeServer) + captured["post"] = post + return captured + + +def test_require_auth_header_uses_environment_credentials() -> None: + """An environment header must be used without loading or authenticating.""" + auth = MagicMock(auth_header="Basic encoded") with patch("lightning_sdk.cli.utils.auth.Auth", return_value=auth): - assert require_auth_header() == "Bearer token" + assert require_auth_header() == "Basic encoded" auth.load.assert_not_called() auth.authenticate.assert_not_called() +def test_require_auth_header_loads_saved_credentials() -> None: + """Saved credentials must work without starting browser authentication.""" + auth = MagicMock() + type(auth).auth_header = PropertyMock(side_effect=[ValueError("missing"), "Bearer token"]) + auth.load.return_value = True + + with patch("lightning_sdk.cli.utils.auth.Auth", return_value=auth): + assert require_auth_header() == "Bearer token" + + auth.authenticate.assert_not_called() + + def test_require_auth_header_instructs_login_without_opening_browser() -> None: """Missing credentials must stop ordinary commands and direct users to login.""" auth = MagicMock() - type(auth).auth_header = PropertyMock( - side_effect=ValueError("No Lightning credentials are available. Run `lightning login` first.") - ) + type(auth).auth_header = PropertyMock(side_effect=ValueError("missing")) + auth.load.return_value = False with patch("lightning_sdk.cli.utils.auth.Auth", return_value=auth), pytest.raises( click.UsageError, match="lightning login" @@ -47,36 +103,340 @@ def test_require_auth_header_instructs_login_without_opening_browser() -> None: auth.authenticate.assert_not_called() -def test_browser_authentication_blocks_the_source_from_opening_a_browser(monkeypatch, tmp_path) -> None: - _no_credentials(monkeypatch, tmp_path) +def test_browser_authentication_blocks_subclass_callback_server(monkeypatch, tmp_path) -> None: + """The guard must also block SDK auth subclasses with custom callback servers.""" + monkeypatch.setattr(Auth, "secrets_file", tmp_path / "credentials.json") + monkeypatch.delenv("LIGHTNING_AUTH_TOKEN", raising=False) + + class CustomAuth(Auth): + browser_started = False - with browser_authentication(False), pytest.raises(ValueError, match="Run `lightning login` first"): - CredentialSource().header() + def load(self) -> bool: + return False + def _run_server(self) -> None: + self.browser_started = True -def test_browser_authentication_allows_the_source_to_open_a_browser(monkeypatch, tmp_path) -> None: - _no_credentials(monkeypatch, tmp_path) - started = [] - monkeypatch.setattr(Auth, "_run_server", lambda self: started.append(True)) + auth = CustomAuth(api_key="", auth_token="") - with browser_authentication(True), pytest.raises(ValueError, match="couldn't find any credentials"): - CredentialSource().header() + with browser_authentication(False), pytest.raises(ValueError, match="lightning login"): + auth.authenticate() - assert started == [True] + assert not auth.browser_started -def test_litserve_auth_uses_the_credential_source(monkeypatch) -> None: - """Deployment authentication must not start callback auth when a credential exists.""" - monkeypatch.setenv("LIGHTNING_AUTH_TOKEN", "token") +def test_litserve_auth_uses_saved_credentials_without_starting_callback_server() -> None: + """Deployment authentication must not start callback auth when credentials exist.""" auth = _AuthLitServe(_AuthMode.DEPLOY) + auth.api_key = "key" - with patch.object(auth, "_run_server", side_effect=AssertionError("callback server must not start")): - assert auth.authenticate() == "Bearer token" + with patch.object( + auth, + "_run_server", + side_effect=AssertionError("callback server must not start"), + ): + assert auth.authenticate().startswith("Bearer ") -def test_litserve_auth_instructs_login_when_callback_auth_would_be_required(monkeypatch, tmp_path) -> None: +def test_litserve_auth_instructs_login_when_callback_auth_would_be_required() -> None: """Deployment callback auth must tell the user to run the explicit login command.""" - _no_credentials(monkeypatch, tmp_path) - - with browser_authentication(False), pytest.raises(click.UsageError, match="lightning login"): + with pytest.raises(click.UsageError, match="lightning login"): _AuthLitServe(_AuthMode.DEPLOY)._run_server() + + +def test_auth_save_stores_session_token_as_auth_token(monkeypatch, tmp_path) -> None: + """Browser login callback passes the session token as `token`; it must be saved as `auth_token`.""" + secrets_file = tmp_path / "credentials.json" + monkeypatch.setattr(Auth, "secrets_file", secrets_file) + monkeypatch.delenv("LIGHTNING_USER_ID", raising=False) + monkeypatch.delenv("LIGHTNING_API_KEY", raising=False) + monkeypatch.delenv("LIGHTNING_AUTH_TOKEN", raising=False) + + auth = Auth() + auth.save(token="sess_user-1.secret", user_id="user-1", api_key="") + + assert auth.auth_token == "sess_user-1.secret" + assert auth.api_key == "" + assert auth.auth_header == "Bearer sess_user-1.secret" + assert json.loads(secrets_file.read_text()) == { + "user_id": "user-1", + "api_key": "", + "auth_token": "sess_user-1.secret", + } + assert (secrets_file.stat().st_mode & 0o777) == 0o600 + + +def test_logout_session_revokes_server_session_and_clears_credentials(monkeypatch, tmp_path) -> None: + """`lightning logout` revokes the saved session token on the server and removes local credentials.""" + secrets_file = tmp_path / "credentials.json" + monkeypatch.setattr(Auth, "secrets_file", secrets_file) + monkeypatch.delenv("LIGHTNING_USER_ID", raising=False) + monkeypatch.delenv("LIGHTNING_API_KEY", raising=False) + monkeypatch.delenv("LIGHTNING_AUTH_TOKEN", raising=False) + + auth = Auth() + auth.save(token="sess_user-1.secret", user_id="user-1") + + with patch("lightning_sdk.cli.utils.auth.requests.post") as post_mock: + logout_session(Auth()) + + post_mock.assert_called_once() + assert post_mock.call_args.kwargs["headers"] == {"Authorization": "Bearer sess_user-1.secret"} + assert not secrets_file.exists() + + +def test_login_with_browser_exchanges_the_code_with_pkce_on_loopback(monkeypatch, tmp_path) -> None: + """The callback trades the code and the verifier matching the sent challenge for a session.""" + from lightning_sdk.lightning_cloud.login import AuthServer + + captured = _capture_browser_login( + monkeypatch, tmp_path, [_response(200, {"access_token": "sess_user-1.secret", "token_type": "Bearer"})] + ) + captured["run"] = lambda endpoint: captured.setdefault("location", endpoint(code="the-code", state="fixed-state")) + + auth = Auth() + AuthServer().login_with_browser(auth) + + url = urlparse(captured["auth_url"]) + query = {k: v[0] for k, v in parse_qs(url.query).items()} + assert url.path == "/oidc/authorize" + assert query["client_id"] == "lightning-cli" + assert query["response_type"] == "code" + assert query["scope"] == "openid" + assert query["state"] == "fixed-state" + assert query["code_challenge_method"] == "S256" + assert query["redirect_uri"] == f"http://127.0.0.1:{captured['config'].port}/callback" + assert captured["config"].host == "127.0.0.1" + + data = captured["post"].call_args.kwargs["data"] + assert captured["post"].call_args.args[0].endswith("/oidc/token") + assert data["grant_type"] == "authorization_code" + assert data["code"] == "the-code" + assert data["redirect_uri"] == query["redirect_uri"] + assert _challenge(data["code_verifier"]) == query["code_challenge"] + + assert captured["location"].headers["location"].endswith("/cli-login-successful") + assert auth.auth_token == "sess_user-1.secret" + assert json.loads(captured["secrets_file"].read_text())["auth_token"] == "sess_user-1.secret" + assert (captured["secrets_file"].stat().st_mode & 0o777) == 0o600 + + +@pytest.mark.parametrize( + "params", + [ + {"code": "the-code", "state": ""}, + {"code": "the-code", "state": "forged-state"}, + {"code": "", "state": "fixed-state", "error": "access_denied"}, + {"code": "", "state": "fixed-state"}, + ], + ids=["no state", "forged state", "denied", "no code"], +) +def test_login_with_browser_refuses_a_bad_callback_without_exchanging(monkeypatch, tmp_path, params) -> None: + from lightning_sdk.lightning_cloud.login import AuthServer + + captured = _capture_browser_login(monkeypatch, tmp_path, []) + captured["run"] = lambda endpoint: captured.setdefault("location", endpoint(**params)) + + auth = Auth() + AuthServer().login_with_browser(auth) + + assert captured["location"].headers["location"].endswith("/cli-login-failed") + captured["post"].assert_not_called() + assert not captured["secrets_file"].exists() + assert not auth.auth_token + + +def test_login_with_browser_reports_a_refused_exchange(monkeypatch, tmp_path) -> None: + from lightning_sdk.lightning_cloud.login import AuthServer + + captured = _capture_browser_login(monkeypatch, tmp_path, [_response(400, {"error": "invalid_grant"})]) + captured["run"] = lambda endpoint: captured.setdefault("location", endpoint(code="used", state="fixed-state")) + + AuthServer().login_with_browser(Auth()) + + assert captured["location"].headers["location"].endswith("/cli-login-failed") + assert not captured["secrets_file"].exists() + + +def test_login_with_browser_without_a_browser_falls_back_to_paste(monkeypatch, tmp_path) -> None: + from lightning_sdk.lightning_cloud.login import AuthServer + + secrets_file = _login_env(monkeypatch, tmp_path) + monkeypatch.setenv(OIDC_LOGIN_ENV, "1") + monkeypatch.setattr("lightning_sdk.cli.utils.auth.requests.head", MagicMock()) + monkeypatch.setattr("webbrowser.open", lambda url: False) + monkeypatch.setattr("uvicorn.Server", lambda config: (_ for _ in ()).throw(AssertionError("server started"))) + monkeypatch.setattr( + "lightning_sdk.cli.utils.auth.login_with_pasted_code", lambda auth: auth.save(auth_token="sess_x") + ) + + AuthServer().login_with_browser(Auth()) + + assert json.loads(secrets_file.read_text())["auth_token"] == "sess_x" + + +def test_login_with_browser_in_a_deployment_without_a_browser_raises(monkeypatch, tmp_path) -> None: + from lightning_sdk.lightning_cloud.login import AuthServer + + _login_env(monkeypatch, tmp_path) + monkeypatch.setenv("LIGHTNING_DEPLOYMENT_ID", "dep") + monkeypatch.setattr("lightning_sdk.cli.utils.auth.requests.head", MagicMock()) + monkeypatch.setattr("webbrowser.open", lambda url: False) + + with pytest.raises(RuntimeError, match="Include SDK credentials"): + AuthServer().login_with_browser(Auth()) + + +def test_login_with_browser_defaults_to_the_session_token_login(monkeypatch, tmp_path) -> None: + """Without the opt-in, the sign-in page hands a session token to `/login-complete`, and the callback + redirects to `/cli-login-successful` without forwarding it in the URL.""" + import asyncio + + from lightning_sdk.lightning_cloud.login import AuthServer + + _login_env(monkeypatch, tmp_path) + + captured: dict = {} + + class FakeServer: + def __init__(self, config) -> None: + captured["app"] = config.app + + def run(self) -> None: + route = next(r for r in captured["app"].routes if getattr(r, "path", None) == "/login-complete") + resp = asyncio.run( + route.endpoint(token="sess_user-1.secret", key="", state="fixed-state", user_id="user-1") + ) + captured["location"] = resp.headers["location"] + + monkeypatch.setattr("lightning_sdk.cli.utils.auth._generate_state", lambda: "fixed-state") + monkeypatch.setattr("lightning_sdk.cli.utils.auth.requests.head", MagicMock()) + monkeypatch.setattr("webbrowser.open", lambda url: captured.setdefault("auth_url", url) or True) + monkeypatch.setattr("uvicorn.Server", FakeServer) + + auth = Auth() + AuthServer().login_with_browser(auth) + + assert "/sign-in?" in captured["auth_url"] + assert "state=fixed-state" in captured["auth_url"] or "state%3Dfixed-state" in captured["auth_url"] + assert captured["location"].endswith("/cli-login-successful") + assert "token=" not in captured["location"] + assert auth.auth_token == "sess_user-1.secret" + + +def test_session_token_login_rejects_callback_without_matching_state(monkeypatch, tmp_path) -> None: + """Omitting or forging `state` must not store the callback token or key.""" + import asyncio + + from lightning_sdk.lightning_cloud.login import AuthServer + + secrets_file = _login_env(monkeypatch, tmp_path) + + captured: dict = {} + + class FakeServer: + def __init__(self, config) -> None: + captured["app"] = config.app + + def run(self) -> None: + captured["route"] = next(r for r in captured["app"].routes if getattr(r, "path", None) == "/login-complete") + + monkeypatch.setattr("lightning_sdk.cli.utils.auth._generate_state", lambda: "fixed-state") + monkeypatch.setattr("lightning_sdk.cli.utils.auth.requests.head", MagicMock()) + monkeypatch.setattr("webbrowser.open", lambda url: True) + monkeypatch.setattr("uvicorn.Server", FakeServer) + + auth = Auth() + AuthServer().login_with_browser(auth) + + for state in ("", "forged-state"): + resp = asyncio.run( + captured["route"].endpoint( + token="sess_attacker.secret", key="attacker-key", state=state, user_id="attacker" + ) + ) + assert resp.headers["location"].endswith("/cli-login-failed") + + assert not secrets_file.exists() + assert not auth.auth_token + assert not auth.api_key + + +def _paste_login(monkeypatch, tmp_path, codes: list, responses: list) -> tuple: + secrets_file = _login_env(monkeypatch, tmp_path) + post = MagicMock(side_effect=responses) + monkeypatch.setattr("lightning_sdk.cli.utils.auth.requests.post", post) + monkeypatch.setattr("lightning_sdk.cli.utils.auth.click.prompt", MagicMock(side_effect=codes)) + return secrets_file, post + + +def test_login_with_pasted_code_strips_a_trailing_slash_from_the_origin(monkeypatch, tmp_path, capsys) -> None: + monkeypatch.setattr("lightning_sdk.lightning_cloud.env.LIGHTNING_CLOUD_URL", "http://localhost:9800/") + _paste_login(monkeypatch, tmp_path, ["the-code"], [_response(200, {"access_token": "sess_user-1.secret"})]) + + login_with_pasted_code(Auth()) + + url = next(word for word in capsys.readouterr().out.split() if "/oidc/authorize" in word) + query = {k: v[0] for k, v in parse_qs(urlparse(url).query).items()} + assert url.startswith("http://localhost:9800/oidc/authorize?") + assert query["redirect_uri"] == "http://localhost:9800/cli-login" + assert query["scope"] == "openid" + + +def test_login_with_pasted_code_prints_the_authorize_url_and_saves_the_session(monkeypatch, tmp_path, capsys) -> None: + secrets_file, post = _paste_login( + monkeypatch, tmp_path, [" the-code \n"], [_response(200, {"access_token": "sess_user-1.secret"})] + ) + + login_with_pasted_code(Auth()) + + output = capsys.readouterr().out + url = next(word for word in output.split() if "/oidc/authorize" in word) + query = {k: v[0] for k, v in parse_qs(urlparse(url).query).items()} + assert query["client_id"] == "lightning-cli" + assert query["scope"] == "openid" + assert query["redirect_uri"].endswith("/cli-login") + assert "//cli-login" not in query["redirect_uri"] + assert query["code_challenge_method"] == "S256" + assert "never give it to anyone" in output + + data = post.call_args.kwargs["data"] + assert data["code"] == "the-code" + assert data["redirect_uri"] == query["redirect_uri"] + assert _challenge(data["code_verifier"]) == query["code_challenge"] + assert json.loads(secrets_file.read_text())["auth_token"] == "sess_user-1.secret" + + +def test_login_with_pasted_code_prompts_again_on_invalid_grant(monkeypatch, tmp_path) -> None: + secrets_file, post = _paste_login( + monkeypatch, + tmp_path, + ["typo", "the-code"], + [_response(400, {"error": "invalid_grant"}), _response(200, {"access_token": "sess_user-1.secret"})], + ) + + login_with_pasted_code(Auth()) + + assert [c.kwargs["data"]["code"] for c in post.call_args_list] == ["typo", "the-code"] + assert json.loads(secrets_file.read_text())["auth_token"] == "sess_user-1.secret" + + +def test_login_with_pasted_code_gives_up_after_three_bad_codes(monkeypatch, tmp_path) -> None: + refused = [_response(400, {"error": "invalid_grant"})] * 3 + secrets_file, post = _paste_login(monkeypatch, tmp_path, ["a", "b", "c"], refused) + + with pytest.raises(RuntimeError, match="not accepted"): + login_with_pasted_code(Auth()) + + assert post.call_count == 3 + assert not secrets_file.exists() + + +def test_login_with_pasted_code_stops_on_any_other_error(monkeypatch, tmp_path) -> None: + secrets_file, post = _paste_login(monkeypatch, tmp_path, ["a"], [_response(404, {})]) + + with pytest.raises(RuntimeError, match="HTTP 404"): + login_with_pasted_code(Auth()) + + assert post.call_count == 1 + assert not secrets_file.exists() diff --git a/python/tests/lightning_cloud/test_cli_login.py b/python/tests/lightning_cloud/test_cli_login.py index 6b9b9e33a..adab8aa59 100644 --- a/python/tests/lightning_cloud/test_cli_login.py +++ b/python/tests/lightning_cloud/test_cli_login.py @@ -5,7 +5,7 @@ import pytest -from lightning_sdk.lightning_cloud import cli_login, env +from lightning_sdk.lightning_cloud import cli_login, credentials, env from lightning_sdk.lightning_cloud.login import Auth @@ -92,7 +92,7 @@ def test_token_returns_the_access_token_and_logout_revokes_it(platform): assert cli_login.token() == "lit_ua_access" assert len(calls) == 1 - cli_login.logout() + assert credentials.logout() == () assert calls[-1][0] == "http://control-plane.test/oidc/revoke" assert calls[-1][1] == {"client_id": "lightning-cli", "token": "lit_ur_refresh"} diff --git a/python/tests/lightning_cloud/test_credentials.py b/python/tests/lightning_cloud/test_credentials.py index 0f698453f..6ddea0cab 100644 --- a/python/tests/lightning_cloud/test_credentials.py +++ b/python/tests/lightning_cloud/test_credentials.py @@ -49,15 +49,14 @@ def request(method, url, headers=None, **kwargs): assert seen == ["Bearer sess_one", "Bearer sess_two"] -def test_saved_api_key_is_bearer_and_never_basic(isolated_credentials, monkeypatch) -> None: +def test_a_saved_api_key_is_never_sent(isolated_credentials, monkeypatch) -> None: isolated_credentials.write_text(json.dumps({"user_id": "user-1", "api_key": "key-1"})) monkeypatch.setenv("LIGHTNING_USER_ID", "user-1") monkeypatch.setenv("LIGHTNING_API_KEY", "key-1") - assert CredentialSource().header() == "Bearer key-1" - assert Auth().authenticate() == "Bearer key-1" - assert Auth().auth_header == "Bearer key-1" - assert Auth().get_auth_header("api_key") == "Bearer key-1" + assert CredentialSource().resolve() is None + with pytest.raises(ValueError, match="lightning login"): + Auth().get_auth_header() def test_a_legacy_session_login_still_sends_bearer(isolated_credentials) -> None: From c77c85c58f93a61bef36fd0d05ec8c1cd21793ab Mon Sep 17 00:00:00 2001 From: fnwoke-lightning Date: Wed, 4 Nov 2026 18:31:47 -0500 Subject: [PATCH 04/27] test(lightning_cloud): port grid's login harness and pin M4-9 and M4-10 Port grid's fake platform, fake browser, keyrings and login worker under tests.lightning_cloud, and extend the fake platform with the interim /sign-in to /login-complete flow and POST /v1/auth/logout. The conftest overrides the SDK's autouse threading.Thread mock, clears the SDK's dummy LIGHTNING_API_KEY and LIGHTNING_AUTH_TOKEN, and puts back the vendored interim login that main's CLI replaces when imported. New cases cover the interim login, the switch at every entry point, the options only the new login has, replacing either kind of login with either kind, logout of both kinds, the five minute stop, and the exact reuse text. Co-authored-by: Cursor --- python/tests/lightning_cloud/__init__.py | 0 python/tests/lightning_cloud/conftest.py | 118 +++ .../tests/lightning_cloud/test_cli_login.py | 536 ++++++++++--- .../tests/lightning_cloud/test_credentials.py | 703 +++++++++++++++++- .../lightning_cloud/test_interim_login.py | 220 ++++++ .../lightning_cloud/test_replaced_login.py | 79 ++ .../lightning_cloud/test_shared_login.py | 94 +++ .../tests/lightning_cloud/utils/__init__.py | 0 .../lightning_cloud/utils/fake_platform.py | 439 +++++++++++ .../tests/lightning_cloud/utils/keyrings.py | 71 ++ .../lightning_cloud/utils/login_worker.py | 41 + 11 files changed, 2178 insertions(+), 123 deletions(-) create mode 100644 python/tests/lightning_cloud/__init__.py create mode 100644 python/tests/lightning_cloud/conftest.py create mode 100644 python/tests/lightning_cloud/test_interim_login.py create mode 100644 python/tests/lightning_cloud/test_replaced_login.py create mode 100644 python/tests/lightning_cloud/test_shared_login.py create mode 100644 python/tests/lightning_cloud/utils/__init__.py create mode 100644 python/tests/lightning_cloud/utils/fake_platform.py create mode 100644 python/tests/lightning_cloud/utils/keyrings.py create mode 100644 python/tests/lightning_cloud/utils/login_worker.py diff --git a/python/tests/lightning_cloud/__init__.py b/python/tests/lightning_cloud/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/tests/lightning_cloud/conftest.py b/python/tests/lightning_cloud/conftest.py new file mode 100644 index 000000000..3ee3db9fe --- /dev/null +++ b/python/tests/lightning_cloud/conftest.py @@ -0,0 +1,118 @@ +import importlib.util +import time +import webbrowser + +import keyring +import pytest + +from lightning_sdk.lightning_cloud import env, login +from lightning_sdk.lightning_cloud.login import Auth, AuthServer +from tests.lightning_cloud.utils.fake_platform import FakeBrowser, FakePlatform +from tests.lightning_cloud.utils.keyrings import MemoryKeyring + + +@pytest.fixture(autouse=True) +def keep_alive_mocker(): # noqa: PT004 + """Replace the SDK fixture that mocks threading.Thread for every test. + + The fake platform and browser each serve from a real thread. + """ + + +def _load_vendored_login(): + spec = importlib.util.spec_from_file_location("_vendored_login", login.__file__) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +_VENDORED = _load_vendored_login() + + +@pytest.fixture(autouse=True) +def _vendored_login(monkeypatch): + """Put back the vendored interim login, which the SDK CLI replaces when imported. + + Collecting any test that imports lightning_sdk.cli applies that replacement, + which keeps no reference to the vendored methods, so they come from a fresh load. + """ + monkeypatch.setattr(Auth, "save", _VENDORED.Auth.save) + monkeypatch.setattr(AuthServer, "login_with_browser", _VENDORED.AuthServer.login_with_browser) + + +@pytest.fixture(autouse=True) +def isolated_login(_clear_lightning_env_vars, tmp_path, monkeypatch): + """Keep every test away from the real ~/.lightning, keychain and browser. + + The SDK's conftest sets a dummy LIGHTNING_API_KEY and LIGHTNING_AUTH_TOKEN, + which these tests must choose for themselves. + """ + for name in ( + "LIGHTNING_AUTH_TOKEN_FILE", + "LIGHTNING_AUTH_TOKEN", + "LIGHTNING_API_KEY", + "LIGHTNING_USER_ID", + "LIGHTNING_CLI_OIDC_LOGIN", + "LIGHTNING_DEPLOYMENT_ID", + ): + monkeypatch.delenv(name, raising=False) + credentials = tmp_path / "credentials.json" + monkeypatch.setattr(Auth, "secrets_file", credentials) + monkeypatch.setenv("LIGHTNING_CREDENTIAL_PATH", str(credentials)) + login_path = tmp_path / "login.json" + monkeypatch.setattr(env, "LIGHTNING_LOGIN_PATH", str(login_path)) + monkeypatch.setenv("LIGHTNING_LOGIN_PATH", str(login_path)) + + def refuse_browser(url, *args, **kwargs): + raise AssertionError("a test opened the real browser") + + monkeypatch.setattr(webbrowser, "open", refuse_browser) + + previous = keyring.get_keyring() + memory = MemoryKeyring() + keyring.set_keyring(memory) + try: + yield memory + finally: + keyring.set_keyring(previous) + + +@pytest.fixture() +def _oidc_login(monkeypatch): + monkeypatch.setenv("LIGHTNING_CLI_OIDC_LOGIN", "1") + + +@pytest.fixture() +def platform(monkeypatch): + fake = FakePlatform() + monkeypatch.setattr(env, "LIGHTNING_CLOUD_URL", fake.url) + monkeypatch.setenv("LIGHTNING_CLOUD_URL", fake.url) + yield fake + fake.stop() + + +@pytest.fixture() +def browser(platform, monkeypatch): + fake = FakeBrowser(platform) + monkeypatch.setattr(webbrowser, "open", fake.open) + yield fake + fake.finish() + + +class Clock: + def __init__(self): + self.now = time.time() + + def __call__(self): + return self.now + + def advance(self, seconds): + self.now += seconds + + +@pytest.fixture() +def clock(monkeypatch): + """One clock for the SDK and the fake platform, moved only by the test.""" + fake = Clock() + monkeypatch.setattr(time, "time", fake) + return fake diff --git a/python/tests/lightning_cloud/test_cli_login.py b/python/tests/lightning_cloud/test_cli_login.py index adab8aa59..0e5925ebc 100644 --- a/python/tests/lightning_cloud/test_cli_login.py +++ b/python/tests/lightning_cloud/test_cli_login.py @@ -1,114 +1,470 @@ +import getpass import json +import logging import os -from datetime import datetime, timedelta, timezone -from urllib.parse import parse_qs, urlparse +import pathlib +import threading +import webbrowser +import keyring import pytest +from keyring.backends import fail -from lightning_sdk.lightning_cloud import cli_login, credentials, env +from lightning_sdk.lightning_cloud import cli_login, env +from lightning_sdk.lightning_cloud.cli_login import LoginError +from lightning_sdk.lightning_cloud.credentials import CredentialKind, logout from lightning_sdk.lightning_cloud.login import Auth +from tests.lightning_cloud.utils.fake_platform import FakeBrowser, OtherDevice, SilentBrowser, raised_within +from tests.lightning_cloud.utils.keyrings import MemoryKeyring +_HOUR = 3600 -class _Response: - def __init__(self, status_code, body): - self.status_code = status_code - self._body = body - - def json(self): - return self._body - - -@pytest.fixture() -def platform(monkeypatch, tmp_path): - login_path = tmp_path / "login.json" - monkeypatch.setattr(env, "LIGHTNING_LOGIN_PATH", str(login_path)) - monkeypatch.setattr(env, "LIGHTNING_CLOUD_URL", "http://control-plane.test") - monkeypatch.setattr(Auth, "secrets_file", tmp_path / "credentials.json") - monkeypatch.setattr(cli_login.getpass, "getpass", lambda prompt: "the-code") - calls = [] - - def post(url, data=None, headers=None, timeout=None): - calls.append((url, dict(data or {}))) - if url.endswith("/oidc/revoke"): - return _Response(200, {}) - return _Response( - 200, - { - "access_token": "lit_ua_access", - "refresh_token": "lit_ur_refresh", - "expires_in": 3600, - "refresh_expires_at": (datetime.now(timezone.utc) + timedelta(days=90)) - .isoformat() - .replace("+00:00", "Z"), - "lightning_grant": {"orgs": "all", "permissions": ["read", "write"]}, - }, - ) - - monkeypatch.setattr(cli_login.requests, "post", post) - return login_path, calls - - -def test_login_stores_the_pair_and_current_login_reads_it(platform, capsys): - login_path, calls = platform - - login = cli_login.login( - no_browser=True, - insecure_storage=True, - orgs=["acme", "beta"], - project="team", - ) - assert login.access_token == "lit_ua_access" - assert login.refresh_token == "lit_ur_refresh" - assert login.grant == {"orgs": "all", "permissions": ["read", "write"]} - current = cli_login.current_login() - assert current.storage == "file" - assert current.grant == {"orgs": "all", "permissions": ["read", "write"]} +def _keychain_entry(url): + return keyring.get_password("lightning-cli", url) or "" + + +def _login_file(): + path = pathlib.Path(env.LIGHTNING_LOGIN_PATH) + return path.read_text() if path.exists() else "" + + +def _nothing_saved(url): + assert cli_login.current_login() is None + assert _login_file() == "" + assert _keychain_entry(url) == "" + + +def _use_browser(monkeypatch, browser): + monkeypatch.setattr(webbrowser, "open", browser.open) + return browser + + +def _paste_codes(monkeypatch, platform, capsys, typed): + terminal = [] + + def read_terminal(): + captured = capsys.readouterr() + terminal.append(captured.out + captured.err) + return "".join(terminal) + + device = OtherDevice(platform, read_terminal, typed) + monkeypatch.setattr(getpass, "getpass", device.getpass) + return device, read_terminal + + +def test_browser_login_keeps_the_pair_in_the_keychain_and_out_of_urls_logs_and_terminal( + platform, browser, capsys, caplog +): + caplog.set_level(logging.DEBUG) + caplog.set_level(logging.DEBUG, logger="urllib3") + + made = cli_login.login() + browser.finish() + + entry = _keychain_entry(platform.url) + assert made.access_token in entry + assert made.refresh_token in entry + assert platform.accepts(made.access_token) + assert browser.landed == f"{platform.url}/cli-login-successful" + assert len(platform.authorizations) == 1 + exchange = platform.calls_to("/oidc/token")[0] + assert exchange["storage"] == "keychain" + + captured = capsys.readouterr() + sdk_logs = [r.getMessage() for r in caplog.records if r.thread == threading.get_ident()] + assert sdk_logs + shown = captured.out + captured.err + "\n".join(sdk_logs) + repr(made) + _login_file() + for secret in (made.access_token, made.refresh_token, exchange["code"], exchange["code_verifier"]): + assert secret not in shown + for url in [*browser.visited, browser.landed]: + assert made.access_token not in url + assert made.refresh_token not in url + assert exchange["code_verifier"] not in url + + +def test_status_shows_grant_storage_and_expiries_without_tokens(platform, browser, clock): + cli_login.login() + + status = cli_login.current_login() + + assert status.url == platform.url + assert status.storage == "keychain" + assert status.grant == {"orgs": "all", "permissions": ["read", "write"]} + assert status.access_expires_at == pytest.approx(clock.now + _HOUR) + assert status.refresh_expires_at == pytest.approx(clock.now + 14 * 24 * _HOUR, abs=1) + assert not hasattr(status, "access_token") + assert not hasattr(status, "refresh_token") + + +def test_org_and_project_hints_reach_the_consent_page(platform, browser): + cli_login.login(orgs=["acme", "beta"], project="research") + + consent = platform.authorizations[0] + assert consent["orgs_hint"] == "acme,beta" + assert consent["project_hint"] == "research" + + +def test_paste_back_sends_the_same_hints_as_the_browser(platform, monkeypatch, capsys): + _paste_codes(monkeypatch, platform, capsys, ["CODE"]) + + cli_login.login(no_browser=True, orgs=["acme", "beta"], project="research") + + consent = platform.authorizations[0] + assert consent["orgs_hint"] == "acme,beta" + assert consent["project_hint"] == "research" + + +def test_the_loopback_answers_other_paths_with_404_and_keeps_waiting(platform, monkeypatch): + browser = _use_browser(monkeypatch, FakeBrowser(platform, stray_paths=["/favicon.ico"])) + + made = cli_login.login() + browser.finish() + + assert browser.stray_statuses == [404] + assert platform.accepts(made.access_token) + - saved = json.loads(login_path.read_text()) - assert saved["storage"] == "file" - assert saved["access_token"] == "lit_ua_access" - assert (login_path.stat().st_mode & 0o777) == 0o600 +def test_a_callback_whose_state_does_not_match_is_refused_and_saves_nothing(platform, monkeypatch): + browser = _use_browser(monkeypatch, FakeBrowser(platform, state="forged")) - url = next(word for word in capsys.readouterr().err.split() if "/oidc/authorize" in word) - query = {key: values[0] for key, values in parse_qs(urlparse(url).query).items()} - assert query["orgs_hint"] == "acme,beta" - assert query["project_hint"] == "team" - assert query["redirect_uri"] == "http://control-plane.test/cli-login" + with pytest.raises(LoginError, match="didn't match"): + cli_login.login() + browser.finish() + + assert browser.landed == f"{platform.url}/cli-login-failed" + assert platform.calls_to("/oidc/token") == [] + _nothing_saved(platform.url) + + +def test_a_denied_login_saves_nothing(platform, monkeypatch): + browser = _use_browser(monkeypatch, FakeBrowser(platform, deny=True)) + + with pytest.raises(LoginError, match="denied"): + cli_login.login() + browser.finish() + + assert browser.landed == f"{platform.url}/cli-login-failed" + assert platform.calls_to("/oidc/token") == [] + _nothing_saved(platform.url) + + +def test_a_browser_login_with_no_callback_in_five_minutes_stops_and_suggests_no_browser(platform, clock, monkeypatch): + browser = _use_browser(monkeypatch, SilentBrowser(clock, 5 * 60 + 1)) + + stopped = raised_within(10, cli_login.login) + + assert isinstance(stopped, LoginError) + assert len(browser.opened) == 1 + assert "`lightning login`" in str(stopped) + assert "--no-browser" in str(stopped) + assert platform.calls_to("/oidc/token") == [] + _nothing_saved(platform.url) + + +def test_a_browser_login_answered_within_five_minutes_completes(platform, clock, monkeypatch): + real = FakeBrowser(platform) + + def open_late(url, *args, **kwargs): + clock.advance(5 * 60 - 1) + return real.open(url) + + monkeypatch.setattr(webbrowser, "open", open_late) + made = cli_login.login() + real.finish() + + assert platform.accepts(made.access_token) - form = calls[0][1] - assert calls[0][0] == "http://control-plane.test/oidc/token" - assert form["client_id"] == "lightning-cli" - assert form["grant_type"] == "authorization_code" - assert form["code"] == "the-code" - assert form["code_verifier"] - assert form["redirect_uri"] == "http://control-plane.test/cli-login" +def test_paste_back_asks_again_after_a_wrong_code(platform, monkeypatch, capsys): + device, read_terminal = _paste_codes(monkeypatch, platform, capsys, ["wrong", "CODE"]) -def test_token_returns_the_access_token_and_logout_revokes_it(platform): - login_path, calls = platform - cli_login.login(no_browser=True, insecure_storage=True) + made = cli_login.login(no_browser=True) - assert cli_login.token() == "lit_ua_access" - assert len(calls) == 1 + assert device.prompts == 2 + assert platform.accepts(made.access_token) + exchanges = platform.calls_to("/oidc/token") + assert [form["code"] for form in exchanges] == ["wrong", device.code] + assert {form["redirect_uri"] for form in exchanges} == {f"{platform.url}/cli-login"} + terminal = read_terminal() + assert "never into a website, form, chat or email" in terminal + assert device.code not in terminal + assert made.access_token not in terminal - assert credentials.logout() == () - assert calls[-1][0] == "http://control-plane.test/oidc/revoke" - assert calls[-1][1] == {"client_id": "lightning-cli", "token": "lit_ur_refresh"} - assert not login_path.exists() +def test_three_wrong_codes_end_the_login(platform, monkeypatch, capsys): + device, _ = _paste_codes(monkeypatch, platform, capsys, ["one", "two", "three", "CODE"]) + with pytest.raises(LoginError, match="3 tries"): + cli_login.login(no_browser=True) -def test_login_refuses_without_a_keychain(monkeypatch, tmp_path): - monkeypatch.setattr(env, "LIGHTNING_LOGIN_PATH", str(tmp_path / "login.json")) - monkeypatch.setattr(cli_login, "_keychain_available", lambda: False) + assert device.prompts == 3 + assert len(platform.calls_to("/oidc/token")) == 3 + _nothing_saved(platform.url) - with pytest.raises(cli_login.LoginError, match="insecure-storage"): + +def test_a_browser_that_cannot_open_falls_back_to_paste_back(platform, monkeypatch, capsys): + monkeypatch.setattr(webbrowser, "open", lambda url, *args, **kwargs: False) + device, _ = _paste_codes(monkeypatch, platform, capsys, ["CODE"]) + + made = cli_login.login() + + assert device.prompts == 1 + assert platform.accepts(made.access_token) + + +def test_a_deployment_without_a_browser_says_to_include_sdk_credentials(platform, monkeypatch): + monkeypatch.setattr(webbrowser, "open", lambda url, *args, **kwargs: False) + monkeypatch.setenv("LIGHTNING_DEPLOYMENT_ID", "the deployment id") + + with pytest.raises(LoginError, match="Include SDK credentials"): + cli_login.login() + + +def test_without_a_keychain_login_refuses_and_names_insecure_storage(platform, browser): + keyring.set_keyring(fail.Keyring()) + + with pytest.raises(LoginError, match="--insecure-storage"): cli_login.login() + assert browser.visited == [] + assert platform.calls == [] + + +def test_a_keychain_that_refuses_the_save_revokes_the_new_login_and_keeps_the_old_one(platform, browser): + old = platform.mint_session("user-1") + Auth().save(auth_token=old) + + class Refusing(MemoryKeyring): + def set_password(self, service, username, password): + raise RuntimeError("locked") + + keyring.set_keyring(Refusing()) + + with pytest.raises(LoginError, match="--insecure-storage"): + cli_login.login() + + [revoked] = platform.calls_to("/oidc/revoke") + assert not platform.family_active(revoked["token"]) + assert cli_login.current_login() is None + assert json.loads(Auth.secrets_file.read_text())["auth_token"] == old + assert platform.session_active(old) + assert platform.logouts() == [] + + +def test_insecure_storage_keeps_the_login_in_a_file_only_the_user_can_read(platform, browser, clock): + keyring.set_keyring(fail.Keyring()) + + made = cli_login.login(insecure_storage=True) + + assert platform.calls_to("/oidc/token")[0]["storage"] == "file" + assert platform.family_storage(made.refresh_token) == "file" + assert made.refresh_expires_at == pytest.approx(clock.now + 7 * 24 * _HOUR, abs=1) + assert cli_login.token() == made.access_token + if os.name != "nt": + assert os.stat(env.LIGHTNING_LOGIN_PATH).st_mode & 0o777 == 0o600 + + +def test_a_new_login_removes_the_interim_credentials(platform, browser): + Auth.secrets_file.write_text(json.dumps({"auth_token": "sess_legacy"})) + + cli_login.login() + + assert not Auth.secrets_file.exists() + + +def test_token_with_twenty_minutes_left_refreshes_first_and_survives_a_later_refresh(platform, browser, clock): + made = cli_login.login() + clock.advance(_HOUR - 20 * 60) + + printed = cli_login.token() + + assert printed != made.access_token + assert len(platform.refreshes()) == 1 + assert cli_login.current_login().access_expires_at - clock.now > 50 * 60 + + clock.advance(55 * 60) + later = cli_login.token() + assert later != printed + assert len(platform.refreshes()) == 2 + assert platform.accepts(printed) + + assert logout() == () + assert not platform.accepts(printed) + assert not platform.accepts(later) + + +def test_token_with_enough_time_left_does_not_refresh(platform, browser, clock): + made = cli_login.login() + clock.advance(20 * 60) + + assert cli_login.token() == made.access_token + assert platform.refreshes() == [] + + +def test_token_without_a_login_says_to_log_in(platform): + with pytest.raises(LoginError, match="lightning login"): + cli_login.token() + + +def test_logout_revokes_both_kinds_on_the_server_and_forgets_both_locally(platform, browser): + made = cli_login.login() + session = platform.mint_session("user-1") + Auth.secrets_file.write_text(json.dumps({"auth_token": session})) + + assert logout() == () + + assert [form["token"] for form in platform.calls_to("/oidc/revoke")] == [made.refresh_token] + assert platform.logouts() == [session] + assert not platform.family_active(made.refresh_token) + assert not platform.accepts(made.access_token) + assert not platform.accepts(session) + assert not Auth.secrets_file.exists() + _nothing_saved(platform.url) + + +@pytest.mark.parametrize( + ("refusing", "failed"), + [("revoke", (CredentialKind.LOGIN,)), ("logout", (CredentialKind.SESSION,))], +) +def test_logout_still_revokes_the_other_kind_and_forgets_both_when_one_revoke_fails( + platform, browser, monkeypatch, refusing, failed +): + made = cli_login.login() + session = platform.mint_session("user-1") + Auth.secrets_file.write_text(json.dumps({"auth_token": session})) + monkeypatch.setattr(platform, refusing, lambda *args: (503, {})) + + assert logout() == failed + + assert platform.accepts(session) == (refusing == "logout") + assert platform.family_active(made.refresh_token) == (refusing == "revoke") + assert not Auth.secrets_file.exists() + _nothing_saved(platform.url) + + +def test_logout_that_cannot_reach_the_server_still_forgets_locally(platform, browser): + cli_login.login() + Auth.secrets_file.write_text(json.dumps({"auth_token": "sess_legacy"})) + platform.stop() + + assert logout() == (CredentialKind.LOGIN, CredentialKind.SESSION) + + assert not Auth.secrets_file.exists() + _nothing_saved(platform.url) + + +def test_logout_with_nothing_saved_revokes_nothing(platform): + assert logout() == () + assert platform.requests == [] + + +def test_every_token_and_revoke_call_names_the_client(platform, browser, clock): + cli_login.login() + clock.advance(_HOUR) + cli_login.token() + cli_login.login() + logout() + + paths = [path for path, _, _ in platform.calls] + assert paths == ["/oidc/token", "/oidc/token", "/oidc/token", "/oidc/revoke", "/oidc/revoke"] + for _, form, user_agent in platform.calls: + assert form["client_id"] == "lightning-cli" + assert user_agent == env.USER_AGENT + + +def test_a_login_for_another_url_is_ignored(platform, browser, monkeypatch): + cli_login.login() + monkeypatch.setattr(env, "LIGHTNING_CLOUD_URL", "https://other.example") + + assert cli_login.current_login() is None + with pytest.raises(LoginError, match="not logged in"): + cli_login.token() + assert platform.calls_to("/oidc/token")[1:] == [] + + +def test_a_refused_refresh_forgets_the_login(platform, browser, clock): + cli_login.login() + platform.blocked = True + clock.advance(_HOUR) + + with pytest.raises(LoginError) as refused: + cli_login.token() + + assert str(refused.value) == "Your login has expired or was revoked. Run `lightning login`." + _nothing_saved(platform.url) + with pytest.raises(LoginError, match="not logged in"): + cli_login.token() + assert len(platform.refreshes()) == 1 + + +def test_a_refresh_token_used_elsewhere_says_the_login_was_copied(platform, browser, clock): + made = cli_login.login() + status, _ = platform.token( + { + "client_id": "lightning-cli", + "grant_type": "refresh_token", + "refresh_token": made.refresh_token, + } + ) + assert status == 200 + clock.advance(_HOUR) + + with pytest.raises(LoginError) as refused: + cli_login.token() + + assert str(refused.value) == ( + "This login was revoked because a copy of it was used elsewhere. " + "Run `lightning login` and consider rotating any keys this machine held." + ) + _nothing_saved(platform.url) + + +@pytest.mark.parametrize( + "description", + ["refresh token was reused", "the login is revoked", "Refresh token was reused; the login is revoked."], +) +def test_only_the_exact_reuse_text_gives_the_reuse_message(platform, browser, clock, monkeypatch, description): + cli_login.login() + monkeypatch.setattr( + platform, "token", lambda form: (400, {"error": "invalid_grant", "error_description": description}) + ) + clock.advance(_HOUR) + + with pytest.raises(LoginError) as refused: + cli_login.token() + + assert str(refused.value) == "Your login has expired or was revoked. Run `lightning login`." + + +def test_token_says_so_when_it_cannot_refresh(platform, browser, clock, monkeypatch): + cli_login.login() + clock.advance(_HOUR - 20 * 60) + platform.stop() + + with pytest.raises(LoginError, match="Couldn't refresh"): + cli_login.token() + + assert cli_login.current_login() is not None + + +def test_a_save_interrupted_between_keychain_and_login_file_converges(platform, browser, clock, monkeypatch): + cli_login.login() + clock.advance(_HOUR) + real_replace = os.replace + + def crash(source, target): + raise OSError("disk went away") + + monkeypatch.setattr(os, "replace", crash) + with pytest.raises(OSError, match="disk went away"): + cli_login.token() + monkeypatch.setattr(os, "replace", real_replace) -def test_file_login_is_readable_only_by_the_user(platform): - login_path, _ = platform - cli_login.login(no_browser=True, insecure_storage=True) + printed = cli_login.token() - assert os.stat(login_path).st_mode & 0o777 == 0o600 + assert platform.accepts(printed) + assert platform.reuse_events == [] + assert platform.family_active(printed) + assert len(platform.refreshes()) == 2 + assert cli_login.token() == printed diff --git a/python/tests/lightning_cloud/test_credentials.py b/python/tests/lightning_cloud/test_credentials.py index 6ddea0cab..44ff5faea 100644 --- a/python/tests/lightning_cloud/test_credentials.py +++ b/python/tests/lightning_cloud/test_credentials.py @@ -1,38 +1,117 @@ +import importlib +import importlib.metadata import json +import os +import re +import threading import pytest +import requests -from lightning_sdk.lightning_cloud import env -from lightning_sdk.lightning_cloud.credentials import CredentialSource +from lightning_sdk.lightning_cloud import cli_login, env +from lightning_sdk.lightning_cloud import rest_client as rest_client_module +from lightning_sdk.lightning_cloud.__version__ import __version__ +from lightning_sdk.lightning_cloud.cli_login import LoginError +from lightning_sdk.lightning_cloud.credentials import ( + Credential, + CredentialKind, + CredentialSource, + browser_login, + logout, + start_login, +) from lightning_sdk.lightning_cloud.login import Auth -from lightning_sdk.lightning_cloud.rest_client import GridRestClient, create_swagger_client +from lightning_sdk.lightning_cloud.openapi.rest import ApiException +from lightning_sdk.lightning_cloud.rest_client import ( + GridRestClient, + LightningClient, + create_swagger_client, +) +from tests.lightning_cloud.utils.fake_platform import FakeBrowser +_MISSING = ( + "We couldn't find any credentials linked to your account. " + "Please try logging in using the CLI command `lightning login`" +) +_LOGIN_FAILED = "Authentication failed. Please run `lightning login`." +_LOGIN_FIRST = "No Lightning credentials are available. Run `lightning login` first." +_HOUR = 3600 +_SWITCH = pytest.mark.parametrize("switch", ["interim", "oidc"]) + + +def _exactly(message): + return f"^{re.escape(message)}$" + + +@pytest.fixture(autouse=True) +def isolated_credentials(isolated_login, monkeypatch): + # Without the fake platform, a login nobody expected fails fast instead of + # reaching the real Lightning. + monkeypatch.setattr(env, "LIGHTNING_CLOUD_URL", "http://127.0.0.1:9") + return Auth.secrets_file -class _Api: - def __init__(self, request): - self.request = request +def _set_switch(monkeypatch, switch): + if switch == "oidc": + monkeypatch.setenv("LIGHTNING_CLI_OIDC_LOGIN", "1") + else: + monkeypatch.delenv("LIGHTNING_CLI_OIDC_LOGIN", raising=False) -@pytest.fixture() -def isolated_credentials(monkeypatch, tmp_path): - for name in ( - "LIGHTNING_AUTH_TOKEN_FILE", - "LIGHTNING_AUTH_TOKEN", - "LIGHTNING_API_KEY", - "LIGHTNING_USER_ID", - ): - monkeypatch.delenv(name, raising=False) - credentials = tmp_path / "credentials.json" - monkeypatch.setattr(Auth, "secrets_file", credentials) - monkeypatch.setattr( - Auth, - "_run_server", - lambda self: (_ for _ in ()).throw(AssertionError("browser login is not expected")), + +def _sent_to_the_api(platform): + return platform.authorizations_sent("/v1/auth/user") + + +def test_precedence_is_token_file_then_env_then_saved_login(isolated_credentials, monkeypatch, tmp_path): + assert env.LIGHTNING_AUTH_TOKEN_FILE == "LIGHTNING_AUTH_TOKEN_FILE" + isolated_credentials.write_text( + json.dumps( + { + "user_id": "user-1", + "api_key": "key-saved", + "auth_token": "sess_saved", + } + ) ) - return credentials + monkeypatch.setenv("LIGHTNING_AUTH_TOKEN", "sess_env") + token_file = tmp_path / "token" + token_file.write_text(" sess_file \n") + monkeypatch.setenv(env.LIGHTNING_AUTH_TOKEN_FILE, str(token_file)) + + source = CredentialSource() + assert source.header() == "Bearer sess_file" + token_file.write_text(" \n") + assert source.header() == "Bearer sess_env" -def test_token_file_is_reread_between_requests(isolated_credentials, monkeypatch, tmp_path) -> None: + token_file.unlink() + assert source.header() == "Bearer sess_env" + + monkeypatch.delenv("LIGHTNING_AUTH_TOKEN") + assert source.header() == "Bearer sess_saved" + + +def test_resolve_names_each_kind_in_precedence_order(platform, browser, monkeypatch, tmp_path): + made = cli_login.login() + browser.finish() + Auth.secrets_file.write_text(json.dumps({"auth_token": "sess_saved"})) + monkeypatch.setenv("LIGHTNING_AUTH_TOKEN", "sess_env") + token_file = tmp_path / "token" + token_file.write_text("lit_w_file") + monkeypatch.setenv(env.LIGHTNING_AUTH_TOKEN_FILE, str(token_file)) + source = CredentialSource() + + assert source.resolve() == Credential(CredentialKind.TOKEN_FILE, "lit_w_file") + token_file.unlink() + assert source.resolve() == Credential(CredentialKind.ENV_TOKEN, "sess_env") + monkeypatch.delenv("LIGHTNING_AUTH_TOKEN") + assert source.resolve() == Credential(CredentialKind.LOGIN, made.access_token) + cli_login._forget() + assert source.resolve() == Credential(CredentialKind.SESSION, "sess_saved") + assert "sess_saved" not in repr(source.resolve()) + + +def test_token_file_is_reread_between_requests(isolated_credentials, monkeypatch, tmp_path): token_file = tmp_path / "token" token_file.write_text("sess_one\n") monkeypatch.setenv(env.LIGHTNING_AUTH_TOKEN_FILE, str(token_file)) @@ -40,37 +119,595 @@ def test_token_file_is_reread_between_requests(isolated_credentials, monkeypatch def request(method, url, headers=None, **kwargs): seen.append(headers["Authorization"]) + return "ok" client = GridRestClient(api_client=_Api(request), with_auth=True) client.api_client.request("GET", "http://example") token_file.write_text("sess_two") client.api_client.request("GET", "http://example") - assert seen == ["Bearer sess_one", "Bearer sess_two"] -def test_a_saved_api_key_is_never_sent(isolated_credentials, monkeypatch) -> None: - isolated_credentials.write_text(json.dumps({"user_id": "user-1", "api_key": "key-1"})) +def test_no_request_sends_basic_whatever_user_id_and_the_saved_login_hold(platform, isolated_credentials, monkeypatch): + session = platform.mint_session("user-1") + isolated_credentials.write_text(json.dumps({"user_id": "user-1", "api_key": "key-1", "auth_token": session})) monkeypatch.setenv("LIGHTNING_USER_ID", "user-1") monkeypatch.setenv("LIGHTNING_API_KEY", "key-1") + assert LightningClient().auth_service_get_user().id == "user-1" + assert Auth().create_api_client().default_headers["Authorization"] == f"Bearer {session}" + headers = [ + *_sent_to_the_api(platform), + Auth().authenticate(), + Auth().auth_header, + Auth().get_auth_header(), + Auth().get_auth_header("api_key"), + ] + assert all(header.startswith("Bearer ") for header in headers), headers + assert _sent_to_the_api(platform) == [f"Bearer {session}"] + + +def test_no_request_sends_lightning_api_key_or_a_saved_api_key(platform, isolated_credentials, monkeypatch): + isolated_credentials.write_text(json.dumps({"user_id": "user-1", "api_key": "key-saved"})) + monkeypatch.setenv("LIGHTNING_USER_ID", "user-1") + monkeypatch.setenv("LIGHTNING_API_KEY", "key-env") + assert CredentialSource().resolve() is None - with pytest.raises(ValueError, match="lightning login"): - Auth().get_auth_header() + with browser_login(False): + with pytest.raises(ValueError, match=_exactly(_LOGIN_FIRST)): + LightningClient().auth_service_get_user() + with pytest.raises(ValueError, match=_exactly(_LOGIN_FIRST)): + Auth().authenticate() + with pytest.raises(ValueError, match=_exactly(_LOGIN_FIRST)): + _ = Auth().auth_header + assert platform.requests == [] + + +@_SWITCH +def test_with_no_credential_the_source_logs_in_once_with_the_login_the_switch_selects( + platform, browser, monkeypatch, switch +): + _set_switch(monkeypatch, switch) + source = CredentialSource() + + first = source.header() + browser.finish() + + assert source.header() == first + if switch == "oidc": + assert first == f"Bearer {cli_login.token(min_remaining=0)}" + assert (len(platform.authorizations), len(platform.sign_ins)) == (1, 0) + else: + assert first == f"Bearer {json.loads(Auth.secrets_file.read_text())['auth_token']}" + assert (len(platform.authorizations), len(platform.sign_ins)) == (0, 1) + assert not [path for _, path, _ in platform.requests if path.startswith("/oidc/")] + +@_SWITCH +def test_a_token_login_then_a_browser_login_as_another_user_sends_the_second_users_token(platform, monkeypatch, switch): + platform.token_keys["key-a"] = "user-a" + Auth().token_login("key-a") + assert LightningClient().auth_service_get_user().id == "user-a" + first = json.loads(Auth.secrets_file.read_text())["auth_token"] -def test_a_legacy_session_login_still_sends_bearer(isolated_credentials) -> None: - isolated_credentials.write_text(json.dumps({"user_id": "user-1", "auth_token": "sess_interim"})) + _set_switch(monkeypatch, switch) + browser = FakeBrowser(platform, user="user-b") + monkeypatch.setattr("webbrowser.open", browser.open) + start_login() + browser.finish() + + assert LightningClient().auth_service_get_user().id == "user-b" + assert not platform.accepts(first) + + +def test_a_later_login_replaces_the_earlier_one_and_the_next_request_sends_it(isolated_credentials, platform, browser): + Auth().save(auth_token=platform.mint_session("user-a")) seen = [] def request(method, url, headers=None, **kwargs): seen.append(headers["Authorization"]) + Auth()._run_server() GridRestClient(api_client=_Api(request), with_auth=True).api_client.request("GET", "http://example") + saved = json.loads(isolated_credentials.read_text())["auth_token"] + assert seen == [f"Bearer {saved}"] + assert platform.accepts(saved) + + +@_SWITCH +def test_a_session_saved_by_the_stable_sdk_still_sends_bearer(isolated_credentials, monkeypatch, switch): + _set_switch(monkeypatch, switch) + isolated_credentials.write_text(json.dumps({"user_id": "user-1", "api_key": "", "auth_token": "sess_interim"})) + seen = [] + def request(method, url, headers=None, **kwargs): + seen.append(headers["Authorization"]) + + GridRestClient(api_client=_Api(request), with_auth=True).api_client.request("GET", "http://example") assert seen == ["Bearer sess_interim"] -def test_the_user_agent_names_the_installed_sdk() -> None: - assert env.USER_AGENT.startswith("lightning-sdk/") - assert create_swagger_client().user_agent == env.USER_AGENT +@_SWITCH +def test_a_saved_new_login_is_sent_whatever_the_switch_says(platform, browser, monkeypatch, switch): + made = cli_login.login() + _set_switch(monkeypatch, switch) + + assert LightningClient().auth_service_get_user().id == "user-1" + assert _sent_to_the_api(platform) == [f"Bearer {made.access_token}"] + + +def test_an_expired_session_gives_one_run_lightning_login_message_without_a_loop(platform, isolated_credentials): + session = platform.mint_session("user-1") + isolated_credentials.write_text(json.dumps({"auth_token": session})) + platform.logout(f"Bearer {session}") + + with pytest.raises(ConnectionError) as refused: + LightningClient().auth_service_get_user() + + assert str(refused.value) == _LOGIN_FAILED + assert len(_sent_to_the_api(platform)) == 1 + + +def test_requests_keep_working_across_access_token_expiry_with_one_refresh_each(platform, browser, clock): + cli_login.login() + client = GridRestClient() + + for _ in range(3): + assert client.auth_service_get_user().id == "user-1" + assert platform.refreshes() == [] + + clock.advance(_HOUR - 100) + for _ in range(3): + assert client.auth_service_get_user().id == "user-1" + assert len(platform.refreshes()) == 1 + + clock.advance(_HOUR) + for _ in range(3): + assert client.auth_service_get_user().id == "user-1" + assert len(platform.refreshes()) == 2 + assert platform.reuse_events == [] + + +def test_an_access_token_the_server_refuses_is_refreshed_and_the_request_retried_once(platform, browser): + cli_login.login() + client = GridRestClient() + platform.expire_access_tokens() + + assert client.auth_service_get_user().id == "user-1" + assert len(platform.refreshes()) == 1 + + +def _one_client_two_threads(platform, monkeypatch, other_thread_token): + """A request whose 401 arrives after another thread used the same client.""" + other = [] + client = None + + def request(method, url, headers=None, **kwargs): + response = requests.get( + f"{platform.url}/v1/auth/user", + headers={"Authorization": headers["Authorization"]}, + timeout=10, + ) + if response.status_code == 401 and threading.current_thread() is threading.main_thread() and not other: + with monkeypatch.context() as patch: + if other_thread_token is not None: + patch.setenv("LIGHTNING_AUTH_TOKEN", other_thread_token) + thread = threading.Thread( + target=lambda: other.append(client.api_client.request("GET", "http://example")) + ) + thread.start() + thread.join(timeout=10) + if response.status_code == 401: + raise ApiException(status=401) + return response.json()["id"] + + client = GridRestClient(api_client=_Api(request), with_auth=True) + return client, other + + +def test_a_refresh_takes_the_token_its_request_sent_after_another_thread_read_another_kind( + platform, browser, monkeypatch +): + cli_login.login() + platform.expire_access_tokens() + client, other = _one_client_two_threads(platform, monkeypatch, platform.mint_session("user-2")) + + assert client.api_client.request("GET", "http://example") == "user-1" + assert other == ["user-2"] + assert len(platform.refreshes()) == 1 + + +def test_a_refresh_after_another_thread_refreshed_the_same_login_does_not_refresh_again(platform, browser, monkeypatch): + cli_login.login() + platform.expire_access_tokens() + client, other = _one_client_two_threads(platform, monkeypatch, None) + + assert client.api_client.request("GET", "http://example") == "user-1" + assert other == ["user-1"] + assert len(platform.refreshes()) == 1 + assert platform.reuse_events == [] + + +def test_a_refused_refresh_says_run_lightning_login_once_and_forgets_the_login(platform, browser, clock): + cli_login.login() + client = GridRestClient() + platform.blocked = True + clock.advance(_HOUR) + + with pytest.raises(LoginError) as refused: + client.auth_service_get_user() + assert str(refused.value) == "Your login has expired or was revoked. Run `lightning login`." + + with pytest.raises(ValueError, match=_exactly(_MISSING)): + client.auth_service_get_user() + assert len(platform.authorizations) == 1 + assert len(platform.refreshes()) == 1 + assert cli_login.current_login() is None + + +def test_a_network_error_while_refreshing_keeps_the_unexpired_token(platform, browser, clock, monkeypatch): + made = cli_login.login() + client = GridRestClient() + clock.advance(_HOUR - 100) + real_post = requests.post + + def unreachable(url, *args, **kwargs): + if url.endswith("/oidc/token"): + raise requests.ConnectionError("network down") + return real_post(url, *args, **kwargs) + + monkeypatch.setattr(requests, "post", unreachable) + assert client.auth_service_get_user().id == "user-1" + assert cli_login.token(min_remaining=0) == made.access_token + + +def test_a_5xx_while_refreshing_keeps_the_unexpired_token(platform, browser, clock, monkeypatch): + made = cli_login.login() + client = GridRestClient() + clock.advance(_HOUR - 100) + + def unavailable(form): + return 503, {"error": "temporarily_unavailable"} + + monkeypatch.setattr(platform, "token", unavailable) + assert client.auth_service_get_user().id == "user-1" + assert len(platform.refreshes()) == 1 + assert cli_login.token(min_remaining=0) == made.access_token + + +def test_a_401_on_an_explicit_authorization_does_not_refresh_the_login(platform, browser): + cli_login.login() + client = GridRestClient() + assert client.auth_service_get_user().id == "user-1" + + with pytest.raises(ConnectionError): + client.api_client.call_api( + "/v1/auth/user", + "GET", + header_params={"Authorization": "Bearer someone-else"}, + response_type="V1GetUserResponse", + ) + assert platform.refreshes() == [] + + +def test_a_login_for_another_url_is_never_sent(platform, browser, monkeypatch): + cli_login.login() + monkeypatch.setattr(env, "LIGHTNING_CLOUD_URL", "https://other.example") + + def request(method, url, headers=None, **kwargs): + raise AssertionError("no request is expected") + + client = GridRestClient(api_client=_Api(request), with_auth=True) + with browser_login(False), pytest.raises(ValueError, match=_exactly(_LOGIN_FIRST)): + client.api_client.request("GET", "http://example") + + +def test_with_browser_login_off_a_request_says_run_lightning_login_first(): + def request(method, url, headers=None, **kwargs): + raise AssertionError("no request is expected") + + client = GridRestClient(api_client=_Api(request), with_auth=True) + with browser_login(False): + with pytest.raises(ValueError, match=_exactly(_LOGIN_FIRST)): + client.api_client.request("GET", "http://example") + with pytest.raises(ValueError, match=_exactly(_LOGIN_FIRST)): + Auth().authenticate() + + +def test_browser_login_off_holds_for_an_auth_subclass_that_would_log_in(): + started = [] + + class LogsIn(Auth): + def _run_server(self): + started.append(1) + + with browser_login(False), pytest.raises(ValueError, match=_exactly(_LOGIN_FIRST)): + LogsIn().authenticate() + assert started == [] + + +def test_auth_header_never_opens_a_browser(): + with pytest.raises(ValueError, match=_exactly(_LOGIN_FIRST)): + _ = Auth().auth_header + + +def test_login_that_saves_nothing_raises_the_login_message(monkeypatch): + started = [] + + def run_server(self): + started.append(1) + + monkeypatch.setattr(Auth, "_run_server", run_server) + source = CredentialSource() + with pytest.raises(ValueError, match=_exactly(_MISSING)): + source.header() + with pytest.raises(ValueError, match=_exactly(_MISSING)): + source.header() + assert started == [1] + + +def test_a_login_turned_off_where_it_starts_says_so_on_every_request(monkeypatch): + def refuse(self): + raise ValueError(_LOGIN_FIRST) + + monkeypatch.setattr(Auth, "_run_server", refuse) + source = CredentialSource() + for _ in range(2): + with pytest.raises(ValueError, match=_exactly(_LOGIN_FIRST)): + source.header() + + +def test_token_file_env_and_session_cannot_refresh(isolated_credentials, monkeypatch, tmp_path): + isolated_credentials.write_text(json.dumps({"auth_token": "sess_saved"})) + source = CredentialSource() + assert source.refresh("sess_saved") is False + + monkeypatch.setenv("LIGHTNING_AUTH_TOKEN", "sess_env") + assert source.refresh("sess_env") is False + + token_file = tmp_path / "token" + token_file.write_text("lit_w_file") + monkeypatch.setenv(env.LIGHTNING_AUTH_TOKEN_FILE, str(token_file)) + assert source.refresh("lit_w_file") is False + + +def test_a_token_file_rotated_since_the_refused_request_counts_as_refreshed(monkeypatch, tmp_path): + token_file = tmp_path / "token" + token_file.write_text("lit_w_old") + monkeypatch.setenv(env.LIGHTNING_AUTH_TOKEN_FILE, str(token_file)) + seen = [] + + def request(method, url, headers=None, **kwargs): + seen.append(headers["Authorization"]) + if len(seen) == 1: + token_file.write_text("lit_w_new") + raise ApiException(status=401) + return "ok" + + client = GridRestClient(api_client=_Api(request), with_auth=True) + assert client.api_client.request("GET", "http://example") == "ok" + assert seen == ["Bearer lit_w_old", "Bearer lit_w_new"] + + +def test_retries_once_when_refresh_succeeds(monkeypatch): + source = _Refreshable(["old", "new"]) + seen = [] + + def request(method, url, headers=None, **kwargs): + seen.append(headers["Authorization"]) + if len(seen) == 1: + raise ApiException(status=401) + return "ok" + + monkeypatch.setattr(rest_client_module, "CredentialSource", lambda: source) + client = GridRestClient(api_client=_Api(request), with_auth=True) + assert client.api_client.request("GET", "http://example") == "ok" + assert seen == ["Bearer old", "Bearer new"] + assert source.refused == ["old"] + + +def test_second_401_after_refresh_raises_once(monkeypatch): + source = _Refreshable(["same", "same"]) + seen = [] + + def request(method, url, headers=None, **kwargs): + seen.append(headers["Authorization"]) + raise ApiException(status=401) + + monkeypatch.setattr(rest_client_module, "CredentialSource", lambda: source) + client = GridRestClient(api_client=_Api(request), with_auth=True) + with pytest.raises(ConnectionError) as exc: + client.api_client.request("GET", "http://example") + assert str(exc.value) == _LOGIN_FAILED + assert len(seen) == 2 + assert source.refused == ["same"] + + +def test_non_refreshable_source_raises_on_the_first_401(monkeypatch): + source = _Static() + seen = [] + + def request(method, url, headers=None, **kwargs): + seen.append(1) + raise ApiException(status=401) + + monkeypatch.setattr(rest_client_module, "CredentialSource", lambda: source) + client = GridRestClient(api_client=_Api(request), with_auth=True) + with pytest.raises(ConnectionError) as exc: + client.api_client.request("GET", "http://example") + assert str(exc.value) == _LOGIN_FAILED + assert seen == [1] + assert source.refused == ["static"] + + +def test_non_401_is_unchanged(monkeypatch): + source = _Static() + + def request(method, url, headers=None, **kwargs): + raise ApiException(status=503) + + monkeypatch.setattr(rest_client_module, "CredentialSource", lambda: source) + client = GridRestClient(api_client=_Api(request), with_auth=True) + with pytest.raises(ApiException) as exc: + client.api_client.request("GET", "http://example") + assert exc.value.status == 503 + assert source.refused == [] + + +def test_explicit_authorization_is_left_in_place(monkeypatch): + source = _Static() + caller_headers = {"Authorization": "Bearer explicit", "Accept": "text/plain"} + + def request(method, url, headers=None, **kwargs): + return dict(headers) + + monkeypatch.setattr(rest_client_module, "CredentialSource", lambda: source) + client = GridRestClient(api_client=_Api(request), with_auth=True) + sent = client.api_client.request("GET", "http://example", headers=caller_headers) + assert sent["Authorization"] == "Bearer explicit" + assert sent["Accept"] == "text/plain" + assert caller_headers == { + "Authorization": "Bearer explicit", + "Accept": "text/plain", + } + assert source.headers == 0 + + +def test_without_auth_keeps_the_401_wrapper(): + def request(method, url, headers=None, **kwargs): + assert headers is None or "Authorization" not in headers + raise ApiException(status=401) + + client = GridRestClient(api_client=_Api(request), with_auth=False) + with pytest.raises(ConnectionError) as exc: + client.api_client.request("GET", "http://example") + assert str(exc.value) == _LOGIN_FAILED + + +def test_lightning_client_without_auth_does_not_send_authorization(monkeypatch): + def request(method, url, headers=None, **kwargs): + return headers + + monkeypatch.setattr(rest_client_module, "create_swagger_client", lambda **kwargs: _Api(request)) + client = LightningClient(retry=False, with_auth=False) + assert client.api_client.request("GET", "http://example") is None + + +def test_save_replaces_the_saved_login(isolated_credentials): + path = isolated_credentials + auth = Auth() + auth.save(user_id="user-1", api_key="key-1") + auth.save(token="ignored", auth_token="sess_1") + assert json.loads(path.read_text()) == {"user_id": "", "api_key": "", "auth_token": "sess_1"} + assert not auth.user_id + assert not auth.api_key + assert auth.auth_token == "sess_1" + + +@pytest.mark.skipif(os.name == "nt", reason="Windows has no POSIX file modes") +def test_saved_login_is_readable_only_by_the_user(isolated_credentials): + path = isolated_credentials + Auth().save(auth_token="sess_1") + assert os.stat(path).st_mode & 0o777 == 0o600 + os.chmod(path, 0o644) + Auth().save(auth_token="sess_2") + assert os.stat(path).st_mode & 0o777 == 0o600 + + +def test_save_replaces_a_corrupt_credentials_file(isolated_credentials): + path = isolated_credentials + path.write_text('{"auth_token": "sess_1", "trunc') + + Auth().save(auth_token="sess_2") + + assert json.loads(path.read_text())["auth_token"] == "sess_2" + + +def test_user_agent_names_the_installed_sdk(monkeypatch): + def version(name): + if name == "lightning_sdk": + return "2026.11.3" + raise importlib.metadata.PackageNotFoundError(name) + + assert _user_agent_with(monkeypatch, version) == "lightning-sdk/2026.11.3" + + +def test_user_agent_without_the_sdk_names_lightning_cloud(monkeypatch): + def version(name): + raise importlib.metadata.PackageNotFoundError(name) + + assert _user_agent_with(monkeypatch, version) == f"lightning-cloud/{__version__}" + + +def test_every_request_names_the_installed_sdk(platform, browser, monkeypatch): + monkeypatch.setenv("LIGHTNING_CLI_OIDC_LOGIN", "1") + start_login() + LightningClient().auth_service_get_user() + monkeypatch.delenv("LIGHTNING_CLI_OIDC_LOGIN") + start_login() + browser.finish() + LightningClient().auth_service_get_user() + logout() + + browser_pages = {("GET", "/oidc/authorize"), ("GET", "/sign-in")} + sent = {(method, path): headers.get("User-Agent") for method, path, headers in platform.requests} + assert set(sent) - browser_pages == { + ("POST", "/oidc/token"), + ("POST", "/oidc/revoke"), + ("HEAD", "/sign-in"), + ("GET", "/v1/auth/user"), + ("POST", "/v1/auth/logout"), + } + assert {sent[request] for request in set(sent) - browser_pages} == { + f"lightning-sdk/{importlib.metadata.version('lightning_sdk')}" + } + + +def _user_agent_with(monkeypatch, version): + with monkeypatch.context() as patch: + patch.setattr(importlib.metadata, "version", version) + importlib.reload(env) + try: + api = create_swagger_client(with_auth=True) + finally: + patch.undo() + importlib.reload(env) + assert "Authorization" not in api.default_headers + return api.user_agent + + +class _Api: + def __init__(self, request): + self.request = request + + +class _Refreshable(CredentialSource): + def __init__(self, tokens): + super().__init__() + self.tokens = list(tokens) + self.refused = [] + self._index = 0 + + def header(self): + token = self.tokens[min(self._index, len(self.tokens) - 1)] + return f"Bearer {token}" + + def refresh(self, token): + self.refused.append(token) + self._index += 1 + return True + + +class _Static(CredentialSource): + def __init__(self): + super().__init__() + self.refused = [] + self.headers = 0 + + def header(self): + self.headers += 1 + return "Bearer static" + + def refresh(self, token): + self.refused.append(token) + return False diff --git a/python/tests/lightning_cloud/test_interim_login.py b/python/tests/lightning_cloud/test_interim_login.py new file mode 100644 index 000000000..0ac4ef773 --- /dev/null +++ b/python/tests/lightning_cloud/test_interim_login.py @@ -0,0 +1,220 @@ +import json +import os +import re +import webbrowser +from urllib.parse import parse_qs, urlsplit + +import pytest + +from lightning_sdk.lightning_cloud import cli_login +from lightning_sdk.lightning_cloud.cli_login import LoginError +from lightning_sdk.lightning_cloud.credentials import ( + CredentialKind, + CredentialSource, + LoginUnavailable, + start_login, +) +from lightning_sdk.lightning_cloud.login import Auth, AuthServer +from lightning_sdk.lightning_cloud.rest_client import LightningClient +from tests.lightning_cloud.utils.fake_platform import FakeBrowser, SilentBrowser, raised_within + +_ENTRY_POINTS = { + "start_login": start_login, + "Auth._run_server": lambda: Auth()._run_server(), + "Auth.authenticate": lambda: Auth().authenticate(), + "CredentialSource.header": lambda: CredentialSource().header(), +} + + +def _saved(): + return json.loads(Auth.secrets_file.read_text()) if Auth.secrets_file.exists() else None + + +def _oidc_requests(platform): + return [path for _, path, _ in platform.requests if path.startswith("/oidc/")] + + +def test_the_interim_login_saves_the_sess_as_auth_token_and_the_next_request_sends_it(platform, browser): + assert start_login() is None + browser.finish() + + [sign_in] = platform.sign_ins + assert re.fullmatch(r"http://localhost:\d+/login-complete\?state=[\w-]+", sign_in["redirectTo"]) + saved = _saved() + assert saved == {"user_id": "user-1", "api_key": "", "auth_token": saved["auth_token"]} + assert saved["auth_token"].startswith("sess_") + assert browser.landed == f"{platform.url}/cli-login-successful" + assert CredentialSource().resolve().kind is CredentialKind.SESSION + + assert LightningClient().auth_service_get_user().id == "user-1" + assert platform.authorizations_sent("/v1/auth/user") == [f"Bearer {saved['auth_token']}"] + assert _oidc_requests(platform) == [] + assert cli_login.current_login() is None + + +@pytest.mark.skipif(os.name == "nt", reason="Windows has no POSIX file modes") +def test_the_interim_login_is_saved_in_a_file_only_the_user_can_read(platform, browser): + Auth.secrets_file.write_text("{}") + os.chmod(Auth.secrets_file, 0o644) + + start_login() + + assert os.stat(Auth.secrets_file).st_mode & 0o777 == 0o600 + + +def test_an_interim_callback_whose_state_does_not_match_is_refused_and_saves_nothing(platform, monkeypatch): + browser = FakeBrowser(platform, state="forged") + monkeypatch.setattr(webbrowser, "open", browser.open) + + with pytest.raises(LoginError, match="didn't match"): + start_login() + browser.finish() + + assert browser.landed == f"{platform.url}/cli-login-failed" + assert _saved() is None + + +def test_an_interim_callback_without_a_token_saves_nothing(platform, monkeypatch): + browser = FakeBrowser(platform) + monkeypatch.setattr(webbrowser, "open", browser.open) + monkeypatch.setattr(platform, "mint_session", lambda user="user-1": "") + + with pytest.raises(LoginError, match="no token"): + start_login() + browser.finish() + + assert browser.landed == f"{platform.url}/cli-login-failed" + assert _saved() is None + + +def test_sign_in_params_reach_the_interim_sign_in_url(platform, browser): + start_login(sign_in_params={"okbhrt": "x334uv8t7v", "litaimodel": "lightning-ai/model"}) + browser.finish() + + [sign_in] = platform.sign_ins + assert sign_in["okbhrt"] == "x334uv8t7v" + assert sign_in["litaimodel"] == "lightning-ai/model" + assert _saved()["auth_token"].startswith("sess_") + + +def test_a_subclass_that_builds_its_own_sign_in_url_still_gets_its_state_checked(platform, monkeypatch): + class OwnUrl(AuthServer): + def get_auth_url(self, port): + return f"{platform.url}/sign-in?redirectTo=http://localhost:{port}/login-complete&okbhrt=code" + + browser = FakeBrowser(platform) + monkeypatch.setattr(webbrowser, "open", browser.open) + + OwnUrl().login_with_browser(Auth()) + browser.finish() + + [sign_in] = platform.sign_ins + assert sign_in["okbhrt"] == "code" + assert "state=" in sign_in["redirectTo"] + assert _saved()["auth_token"].startswith("sess_") + + +def test_the_interim_login_never_saves_the_legacy_key_the_callback_carries(platform, browser): + start_login() + browser.finish() + + assert _saved()["api_key"] == "" + assert "key_legacy" not in Auth.secrets_file.read_text() + + +def test_an_interim_login_with_no_callback_in_five_minutes_stops(platform, clock, monkeypatch): + browser = SilentBrowser(clock, 5 * 60 + 1) + monkeypatch.setattr(webbrowser, "open", browser.open) + + stopped = raised_within(10, start_login) + + assert isinstance(stopped, LoginError) + assert len(browser.opened) == 1 + assert "5 minutes" in str(stopped) + assert "`lightning login` again" in str(stopped) + assert "--no-browser" not in str(stopped) + assert _saved() is None + + +def test_the_interim_login_says_so_when_no_browser_opens(platform, monkeypatch): + monkeypatch.setattr(webbrowser, "open", lambda url, *args, **kwargs: False) + + with pytest.raises(LoginError, match="No browser could be opened"): + start_login() + + assert _saved() is None + + +@pytest.mark.parametrize("entry", list(_ENTRY_POINTS)) +@pytest.mark.parametrize("switch", [None, "1"]) +def test_the_switch_chooses_the_login_every_entry_point_starts(platform, browser, monkeypatch, entry, switch): + if switch: + monkeypatch.setenv("LIGHTNING_CLI_OIDC_LOGIN", switch) + + _ENTRY_POINTS[entry]() + browser.finish() + + if switch: + assert (len(platform.authorizations), platform.sign_ins) == (1, []) + assert cli_login.current_login() is not None + assert _saved() is None + else: + assert (platform.authorizations, len(platform.sign_ins)) == ([], 1) + assert _oidc_requests(platform) == [] + assert _saved()["auth_token"].startswith("sess_") + + +def test_the_switch_is_read_when_each_login_starts(platform, monkeypatch): + browser = FakeBrowser(platform) + monkeypatch.setattr(webbrowser, "open", browser.open) + + start_login() + browser.finish() + monkeypatch.setenv("LIGHTNING_CLI_OIDC_LOGIN", "1") + made = start_login() + browser.finish() + + assert len(platform.sign_ins) == 1 + assert len(platform.authorizations) == 1 + assert cli_login.token(min_remaining=0) == made.access_token + + +@pytest.mark.parametrize( + ("options", "named"), + [ + ({"no_browser": True}, "--no-browser"), + ({"insecure_storage": True}, "--insecure-storage"), + ({"orgs": ["acme"]}, "--org"), + ({"project": "research"}, "--project"), + ], +) +def test_without_the_switch_an_option_only_the_new_login_has_is_refused(platform, monkeypatch, options, named): + monkeypatch.setattr(webbrowser, "open", lambda url, *args, **kwargs: pytest.fail("a browser opened")) + + with pytest.raises(LoginUnavailable) as refused: + start_login(**options) + + assert isinstance(refused.value, ValueError) + assert named in str(refused.value) + assert "LIGHTNING_CLI_OIDC_LOGIN" in str(refused.value) + assert platform.requests == [] + + +def test_with_the_switch_the_new_login_takes_the_hints_and_ignores_sign_in_params(platform, browser, _oidc_login): + made = start_login(orgs=["acme"], project="research", sign_in_params={"okbhrt": "code"}) + + [consent] = platform.authorizations + assert (consent["orgs_hint"], consent["project_hint"]) == ("acme", "research") + assert "okbhrt" not in consent + assert made.grant == {"orgs": "all", "permissions": ["read", "write"]} + + +def test_the_interim_login_url_carries_no_credential(platform, browser): + start_login() + browser.finish() + + saved = _saved()["auth_token"] + for url in browser.visited[:1] + [browser.landed]: + assert saved not in url + assert "token=" not in urlsplit(url).query + assert parse_qs(urlsplit(browser.landed).query) == {} diff --git a/python/tests/lightning_cloud/test_replaced_login.py b/python/tests/lightning_cloud/test_replaced_login.py new file mode 100644 index 000000000..30043d940 --- /dev/null +++ b/python/tests/lightning_cloud/test_replaced_login.py @@ -0,0 +1,79 @@ +import json +import webbrowser + +import keyring +import pytest + +from lightning_sdk.lightning_cloud import cli_login +from lightning_sdk.lightning_cloud.cli_login import LoginError +from lightning_sdk.lightning_cloud.credentials import start_login +from lightning_sdk.lightning_cloud.login import Auth +from lightning_sdk.lightning_cloud.rest_client import LightningClient +from tests.lightning_cloud.utils.fake_platform import FakeBrowser + +_KINDS = ["interim", "oidc"] + + +class _Saved: + """What one login left: the token requests send, and what revokes it.""" + + def __init__(self, kind, token, refresh_token=None): + self.kind = kind + self.token = token + self.refresh_token = refresh_token + + +def _log_in(platform, monkeypatch, kind, user, **browser): + if kind == "oidc": + monkeypatch.setenv("LIGHTNING_CLI_OIDC_LOGIN", "1") + else: + monkeypatch.delenv("LIGHTNING_CLI_OIDC_LOGIN", raising=False) + fake = FakeBrowser(platform, user=user, **browser) + monkeypatch.setattr(webbrowser, "open", fake.open) + try: + made = start_login() + finally: + fake.finish() + if kind == "oidc": + return _Saved(kind, made.access_token, made.refresh_token) + return _Saved(kind, json.loads(Auth.secrets_file.read_text())["auth_token"]) + + +def _revoked(platform, saved): + if saved.kind == "oidc": + return [form["token"] for form in platform.calls_to("/oidc/revoke")] == [saved.refresh_token] + return platform.logouts() == [saved.token] + + +@pytest.mark.parametrize("second", _KINDS) +@pytest.mark.parametrize("first", _KINDS) +def test_logging_in_again_revokes_the_login_it_replaced(platform, monkeypatch, first, second): + old = _log_in(platform, monkeypatch, first, "user-a") + assert LightningClient().auth_service_get_user().id == "user-a" + + _log_in(platform, monkeypatch, second, "user-b") + + assert LightningClient().auth_service_get_user().id == "user-b" + assert not platform.accepts(old.token) + assert _revoked(platform, old) + if first == "oidc": + assert not platform.family_active(old.refresh_token) + if second == "interim": + assert cli_login.current_login() is None + assert keyring.get_password("lightning-cli", platform.url) is None + else: + assert not Auth.secrets_file.exists() + + +@pytest.mark.parametrize("second", _KINDS) +@pytest.mark.parametrize("first", _KINDS) +def test_a_login_that_fails_leaves_the_old_login_saved_and_live(platform, monkeypatch, first, second): + old = _log_in(platform, monkeypatch, first, "user-a") + + with pytest.raises(LoginError, match="didn't match"): + _log_in(platform, monkeypatch, second, "user-b", state="forged") + + assert LightningClient().auth_service_get_user().id == "user-a" + assert platform.accepts(old.token) + assert platform.calls_to("/oidc/revoke") == [] + assert platform.logouts() == [] diff --git a/python/tests/lightning_cloud/test_shared_login.py b/python/tests/lightning_cloud/test_shared_login.py new file mode 100644 index 000000000..f15379ed8 --- /dev/null +++ b/python/tests/lightning_cloud/test_shared_login.py @@ -0,0 +1,94 @@ +import multiprocessing +import os +import queue +import time + +import keyring +import pytest + +from lightning_sdk.lightning_cloud import cli_login, env +from tests.lightning_cloud.utils import login_worker +from tests.lightning_cloud.utils.keyrings import KEYRING_PATH, FileKeyring + +_PROCESSES = 8 + + +@pytest.fixture() +def _shared_keyring(tmp_path, monkeypatch): + monkeypatch.setenv(KEYRING_PATH, str(tmp_path / "keyring.json")) + monkeypatch.setenv("PYTHON_KEYRING_BACKEND", "tests.lightning_cloud.utils.keyrings.FileKeyring") + keyring.set_keyring(FileKeyring()) + + +def _run(target, args_for): + """Run ``target`` in fresh processes and return what each reported.""" + context = multiprocessing.get_context("spawn") + results = context.Queue() + # Process.start() drops its args, and a spawned child can only open a + # barrier the parent still holds. + shared = args_for(context) + processes = [context.Process(target=target, args=(*args, results)) for args in shared] + for process in processes: + process.start() + try: + reports = [] + deadline = time.monotonic() + 90 + for _ in processes: + try: + reports.append(results.get(timeout=max(deadline - time.monotonic(), 0.1))) + except queue.Empty: + break + for process in processes: + process.join(timeout=max(deadline - time.monotonic(), 0.1)) + finally: + for process in processes: + if process.is_alive(): + process.kill() + errors = [detail for status, detail in reports if status != "ok"] + assert errors == [] + assert len(reports) == len(processes) + assert [process.exitcode for process in processes] == [0] * len(processes) + return [detail for _, detail in reports] + + +def test_eight_processes_share_one_login_through_several_refreshes_without_reuse(platform, browser, _shared_keyring): + # Anything under five minutes left is refreshed, so each token lasts one second. + platform.access_ttl = 301 + cli_login.login() + + def args_for(context): + barrier = context.Barrier(_PROCESSES) + return [(barrier, 4.0)] * _PROCESSES + + calls = _run(login_worker.call_the_api, args_for) + + assert all(count > 0 for count in calls) + assert platform.reuse_events == [] + assert 3 <= len(platform.refreshes()) <= 10 + assert platform.family_active(cli_login.token(min_remaining=0)) + + +def test_a_fresh_process_repairs_a_save_interrupted_between_keychain_and_login_file( + platform, browser, _shared_keyring, monkeypatch +): + cli_login.login() + real_replace = os.replace + + def crash_on_the_login_file(source, target): + if os.fspath(target) == env.LIGHTNING_LOGIN_PATH: + raise OSError("the process died here") + real_replace(source, target) + + monkeypatch.setattr(os, "replace", crash_on_the_login_file) + with pytest.raises(OSError, match="the process died here"): + cli_login.token(min_remaining=10 * 3600) + monkeypatch.setattr(os, "replace", real_replace) + + [printed] = _run(login_worker.take_a_token, lambda context: [()]) + + assert platform.accepts(printed) + assert platform.reuse_events == [] + assert len(platform.refreshes()) == 2 + status = cli_login.current_login() + assert status.access_expires_at == pytest.approx(time.time() + 3600, abs=60) + assert cli_login.token() == printed diff --git a/python/tests/lightning_cloud/utils/__init__.py b/python/tests/lightning_cloud/utils/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/tests/lightning_cloud/utils/fake_platform.py b/python/tests/lightning_cloud/utils/fake_platform.py new file mode 100644 index 000000000..1f9a3eb8e --- /dev/null +++ b/python/tests/lightning_cloud/utils/fake_platform.py @@ -0,0 +1,439 @@ +import base64 +import hashlib +import json +import re +import secrets +import threading +import time +from datetime import datetime, timezone +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from urllib.parse import parse_qs, urlencode, urlsplit, urlunsplit + +import requests + +CLIENT_ID = "lightning-cli" +_LOOPBACK = re.compile(r"http://127\.0\.0\.1:\d+/callback") +_INTERIM_CALLBACK = re.compile(r"http://localhost:\d+/login-complete\?state=[^&]+") +_DAY = 24 * 3600 + + +class _Server(ThreadingHTTPServer): + # Eight processes connect at once, and the default backlog of 5 resets some. + request_queue_size = 64 + daemon_threads = True + + +class FakePlatform: + """The control plane's login endpoints, both logins, and one API call. + + The new login is the lightning-cli OAuth client. The interim login is the + sign-in page handing a ``sess_`` token to a localhost callback. Refresh has + no reuse grace, so any refresh race shows up as a reuse event. + """ + + def __init__(self, access_ttl=3600): + self.access_ttl = access_ttl + self.calls = [] + self.requests = [] + self.authorizations = [] + self.sign_ins = [] + self.reuse_events = [] + self.blocked = False + self.token_keys = {} + self._codes = {} + self._families = {} + self._refresh = {} + self._access = {} + self._sessions = {} + self._lock = threading.Lock() + self._server = _Server(("127.0.0.1", 0), _handler(self)) + self.url = f"http://127.0.0.1:{self._server.server_address[1]}" + self._thread = threading.Thread(target=self._server.serve_forever, kwargs={"poll_interval": 0.02}, daemon=True) + self._thread.start() + + def stop(self): + self._server.shutdown() + self._server.server_close() + + def approve(self, authorize_query, user="user-1"): + """Click Approve on the consent page: issue a code for that request.""" + code = "code_" + secrets.token_urlsafe(16) + with self._lock: + self._codes[code] = { + "redirect_uri": authorize_query["redirect_uri"], + "challenge": authorize_query["code_challenge"], + "user": user, + } + return code + + def mint_session(self, user="user-1"): + """What the sign-in page hands the interim login: a ``sess_`` for ``user``.""" + token = "sess_" + secrets.token_urlsafe(16) + with self._lock: + self._sessions[token] = {"user": user, "active": True} + return token + + def accepts(self, token): + return self.user_of(token) is not None + + def user_of(self, token): + """The user ``GET /v1/auth/user`` answers for ``token``, or None for a 401.""" + with self._lock: + session = self._sessions.get(token) + if session is not None: + return session["user"] if session["active"] else None + if self._access_ok(token): + return self._access[token]["family"]["user"] + return None + + def session_active(self, token): + with self._lock: + return self._sessions[token]["active"] + + def family_active(self, token): + with self._lock: + family = self._family_of(token) + return family is not None and family["active"] + + def family_storage(self, token): + with self._lock: + return self._family_of(token)["storage"] + + def expire_access_tokens(self): + with self._lock: + for access in self._access.values(): + access["expires_at"] = 0 + + def calls_to(self, path): + return [form for call_path, form, _ in self.calls if call_path == path] + + def refreshes(self): + return [form for form in self.calls_to("/oidc/token") if form.get("grant_type") == "refresh_token"] + + def logouts(self): + """The bearer tokens ``POST /v1/auth/logout`` received.""" + return [ + headers.get("Authorization", "")[len("Bearer ") :] + for method, path, headers in self.requests + if (method, path) == ("POST", "/v1/auth/logout") + ] + + def authorizations_sent(self, path): + return [headers.get("Authorization") for _, request_path, headers in self.requests if request_path == path] + + def authorize(self, query): + problems = [] + if query.get("response_type") != "code": + problems.append("response_type") + if query.get("client_id") != CLIENT_ID: + problems.append("client_id") + redirect = query.get("redirect_uri", "") + if not (_LOOPBACK.fullmatch(redirect) or redirect == f"{self.url}/cli-login"): + problems.append("redirect_uri") + if "openid" not in query.get("scope", "").split(): + problems.append("scope") + if not query.get("state"): + problems.append("state") + if not 43 <= len(query.get("code_challenge", "")) <= 128: + problems.append("code_challenge") + if query.get("code_challenge_method") != "S256": + problems.append("code_challenge_method") + if not problems: + self.authorizations.append(query) + return problems + + def sign_in(self, query, user): + """The sign-in page: hand a new ``sess_`` to the ``redirectTo`` callback.""" + redirect = query.get("redirectTo", "") + if not _INTERIM_CALLBACK.fullmatch(redirect): + return None + self.sign_ins.append(query) + callback = urlsplit(redirect) + params = parse_qs(callback.query) + params.update({"token": [self.mint_session(user)], "userID": [user], "key": ["key_legacy"]}) + return urlunsplit(callback._replace(query=urlencode(params, doseq=True))) + + def token(self, form): + if form.get("client_id") != CLIENT_ID: + return 401, {"error": "invalid_client"} + with self._lock: + if form.get("grant_type") == "authorization_code": + return self._exchange(form) + if form.get("grant_type") == "refresh_token": + return self._rotate(form) + return 400, {"error": "unsupported_grant_type"} + + def revoke(self, form): + if form.get("client_id") != CLIENT_ID: + return 401, {"error": "invalid_client"} + with self._lock: + family = self._family_of(form.get("token", "")) + if family is not None: + family["active"] = False + return 200, {} + + def logout(self, bearer): + token = bearer[len("Bearer ") :] if bearer.startswith("Bearer ") else "" + with self._lock: + session = self._sessions.get(token) + if session is None or not session["active"]: + return 401, {"code": 16, "message": "unauthenticated"} + session["active"] = False + return 200, {} + + def token_login(self, body): + user = self.token_keys.get(body.get("tokenKey")) + if user is None: + return 401, {"code": 16, "message": "unauthenticated"} + return 200, {"token": self.mint_session(user)} + + def _exchange(self, form): + storage = form.get("storage") + if storage not in ("keychain", "file"): + return 400, {"error": "invalid_request", "error_description": "storage"} + code = self._codes.get(form.get("code", "")) + if code is None: + return 400, {"error": "invalid_grant", "error_description": "unknown authorization code"} + if code["redirect_uri"] != form.get("redirect_uri"): + return 400, {"error": "invalid_grant", "error_description": "redirect_uri mismatch"} + verifier = form.get("code_verifier", "").encode() + challenge = base64.urlsafe_b64encode(hashlib.sha256(verifier).digest()).rstrip(b"=").decode() + if challenge != code["challenge"]: + return 400, {"error": "invalid_grant", "error_description": "PKCE verification failed"} + del self._codes[form["code"]] + family = {"active": True, "storage": storage, "current": None, "user": code["user"]} + self._families[id(family)] = family + return 200, self._issue(family) + + def _rotate(self, form): + presented = form.get("refresh_token", "") + family = self._family_of(presented) + if family is None or not family["active"]: + return 400, {"error": "invalid_grant", "error_description": "invalid refresh token"} + if presented != family["current"]: + family["active"] = False + self.reuse_events.append(presented) + return 400, { + "error": "invalid_grant", + "error_description": "refresh token was reused; the login is revoked", + } + if self.blocked: + return 400, {"error": "access_denied"} + return 200, self._issue(family) + + def _issue(self, family): + access = "lit_ua_" + secrets.token_urlsafe(24) + refresh = "lit_ur_" + secrets.token_urlsafe(24) + now = time.time() + family["current"] = refresh + self._refresh[refresh] = family + self._access[access] = {"family": family, "expires_at": now + self.access_ttl} + idle = 7 * _DAY if family["storage"] == "file" else 14 * _DAY + refresh_expires_at = datetime.fromtimestamp(now + idle, timezone.utc) + return { + "access_token": access, + "token_type": "Bearer", + "expires_in": self.access_ttl, + "refresh_token": refresh, + "refresh_expires_at": refresh_expires_at.strftime("%Y-%m-%dT%H:%M:%S.123456789Z"), + "lightning_grant": {"orgs": "all", "permissions": ["read", "write"]}, + } + + def _family_of(self, token): + if token in self._refresh: + return self._refresh[token] + if token in self._access: + return self._access[token]["family"] + return None + + def _access_ok(self, token): + access = self._access.get(token) + return access is not None and access["family"]["active"] and time.time() < access["expires_at"] + + +def _handler(platform): + class Handler(BaseHTTPRequestHandler): + def do_HEAD(self): # noqa: N802 + platform.requests.append(("HEAD", urlsplit(self.path).path, dict(self.headers))) + self.send_response(200) + self.end_headers() + + def do_GET(self): # noqa: N802 + parts = urlsplit(self.path) + platform.requests.append(("GET", parts.path, dict(self.headers))) + query = {key: values[0] for key, values in parse_qs(parts.query).items()} + if parts.path == "/oidc/authorize": + problems = platform.authorize(query) + if problems: + self._json(400, {"error": "invalid_request", "error_description": ",".join(problems)}) + return + self.send_response(302) + self.send_header("Location", f"/cli-login?{parts.query}") + self.end_headers() + elif parts.path == "/sign-in": + landing = platform.sign_in(query, self.headers.get("X-Signed-In-As", "user-1")) + if landing is None: + self._json(400, {"error": "redirectTo"}) + return + self.send_response(302) + self.send_header("Location", landing) + self.end_headers() + elif parts.path == "/v1/auth/user": + bearer = self.headers.get("Authorization", "") + token = bearer[len("Bearer ") :] if bearer.startswith("Bearer ") else "" + user = platform.user_of(token) + if user is not None: + self._json(200, {"id": user}) + else: + self._json(401, {"code": 16, "message": "unauthenticated"}) + else: + self._json(404, {}) + + def do_POST(self): # noqa: N802 + length = int(self.headers.get("Content-Length", 0)) + raw = self.rfile.read(length).decode() + platform.requests.append(("POST", self.path, dict(self.headers))) + if self.path == "/v1/auth/logout": + self._json(*platform.logout(self.headers.get("Authorization", ""))) + return + if self.path == "/v1/auth/token-login": + self._json(*platform.token_login(json.loads(raw or "{}"))) + return + form = {key: values[0] for key, values in parse_qs(raw).items()} + platform.calls.append((self.path, form, self.headers.get("User-Agent"))) + if self.path == "/oidc/token": + self._json(*platform.token(form)) + elif self.path == "/oidc/revoke": + self._json(*platform.revoke(form)) + else: + self._json(404, {}) + + def _json(self, status, body): + payload = json.dumps(body).encode() + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + def log_message(self, format, *args): + pass + + return Handler + + +class FakeBrowser: + """Stands in for webbrowser.open, for either login. + + It signs in or approves as ``user`` and follows the redirect to the + localhost callback. Like a real browser it runs on its own, so the CLI can + serve the callback. + """ + + def __init__(self, platform, state=None, deny=False, stray_paths=(), user="user-1"): + self.platform = platform + self.state = state + self.deny = deny + self.stray_paths = stray_paths + self.user = user + self.visited = [] + self.stray_statuses = [] + self.landed = None + self._thread = None + + def open(self, url, *args, **kwargs): + self._thread = threading.Thread(target=self._visit, args=(url,), daemon=True) + self._thread.start() + return True + + def finish(self): + if self._thread is not None: + self._thread.join(timeout=10) + + def _visit(self, url): + self.visited.append(url) + landing = self._sign_in(url) if urlsplit(url).path == "/sign-in" else self._approve(url) + self.visited.append(landing) + response = requests.get(landing, allow_redirects=False, timeout=10) + self.landed = response.headers.get("Location") + + def _sign_in(self, url): + page = requests.get(url, headers={"X-Signed-In-As": self.user}, allow_redirects=False, timeout=10) + assert page.status_code == 302, page.text + landing = urlsplit(page.headers["Location"]) + self._visit_strays(f"{landing.scheme}://{landing.netloc}") + if self.state is None: + return landing.geturl() + params = parse_qs(landing.query) + params["state"] = [self.state] + return urlunsplit(landing._replace(query=urlencode(params, doseq=True))) + + def _approve(self, url): + consent = requests.get(url, allow_redirects=False, timeout=10) + assert consent.status_code == 302, consent.text + query = {key: values[0] for key, values in parse_qs(urlsplit(url).query).items()} + callback = query["redirect_uri"] + self._visit_strays(callback[: -len("/callback")]) + if self.deny: + params = {"error": "access_denied", "state": query["state"]} + else: + params = {"code": self.platform.approve(query, self.user), "state": self.state or query["state"]} + return f"{callback}?{urlencode(params)}" + + def _visit_strays(self, origin): + for path in self.stray_paths: + self.stray_statuses.append(requests.get(origin + path, timeout=10).status_code) + + +class SilentBrowser: + """A browser that opens but never comes back, while the clock runs past the login's wait.""" + + def __init__(self, clock, seconds): + self.clock = clock + self.seconds = seconds + self.opened = [] + + def open(self, url, *args, **kwargs): + self.opened.append(url) + self.clock.advance(self.seconds) + return True + + +def raised_within(seconds, call): + """What ``call`` raised, failing rather than hanging when it is still running after ``seconds``.""" + outcome = [] + + def run(): + try: + call() + except BaseException as err: + outcome.append(err) + else: + outcome.append(None) + + worker = threading.Thread(target=run, daemon=True) + worker.start() + worker.join(seconds) + assert outcome, f"still waiting after {seconds} seconds" + return outcome[0] + + +class OtherDevice: + """Someone approving a paste-back login on another device, then typing codes.""" + + def __init__(self, platform, read_terminal, typed): + self.platform = platform + self.read_terminal = read_terminal + self.typed = list(typed) + self.prompts = 0 + self.code = None + + def getpass(self, prompt=""): + self.prompts += 1 + if self.code is None: + url = re.search(r"https?://\S+/oidc/authorize\?\S+", self.read_terminal()).group(0) + requests.get(url, allow_redirects=False, timeout=10).raise_for_status() + query = {key: values[0] for key, values in parse_qs(urlsplit(url).query).items()} + self.code = self.platform.approve(query) + typed = self.typed.pop(0) + return self.code if typed == "CODE" else typed diff --git a/python/tests/lightning_cloud/utils/keyrings.py b/python/tests/lightning_cloud/utils/keyrings.py new file mode 100644 index 000000000..eb8a6d1ea --- /dev/null +++ b/python/tests/lightning_cloud/utils/keyrings.py @@ -0,0 +1,71 @@ +import json +import os +import pathlib +import tempfile +import time + +from keyring.backend import KeyringBackend +from keyring.errors import PasswordDeleteError + +KEYRING_PATH = "LIGHTNING_TEST_KEYRING" + + +class MemoryKeyring(KeyringBackend): + priority = 1 + + def __init__(self): + super().__init__() + self.entries = {} + + def get_password(self, service, username): + return self.entries.get((service, username)) + + def set_password(self, service, username, password): + self.entries[(service, username)] = password + + def delete_password(self, service, username): + if self.entries.pop((service, username), None) is None: + raise PasswordDeleteError(username) + + +class FileKeyring(KeyringBackend): + """A keyring several processes share, in the file named by LIGHTNING_TEST_KEYRING.""" + + priority = 1 + + def get_password(self, service, username): + return self._read().get(f"{service} {username}") + + def set_password(self, service, username, password): + entries = self._read() + entries[f"{service} {username}"] = password + self._write(entries) + + def delete_password(self, service, username): + entries = self._read() + if entries.pop(f"{service} {username}", None) is None: + raise PasswordDeleteError(username) + self._write(entries) + + def _path(self): + return pathlib.Path(os.environ[KEYRING_PATH]) + + def _read(self): + try: + return json.loads(self._path().read_text()) + except FileNotFoundError: + return {} + + def _write(self, entries): + path = self._path() + fd, tmp = tempfile.mkstemp(dir=path.parent) + with os.fdopen(fd, "w") as file: + json.dump(entries, file) + for _ in range(100): + try: + os.replace(tmp, path) + return + except PermissionError: + # Windows refuses to replace a file another process has open. + time.sleep(0.01) + os.replace(tmp, path) diff --git a/python/tests/lightning_cloud/utils/login_worker.py b/python/tests/lightning_cloud/utils/login_worker.py new file mode 100644 index 000000000..c4a6a9db1 --- /dev/null +++ b/python/tests/lightning_cloud/utils/login_worker.py @@ -0,0 +1,41 @@ +"""Code a spawned process runs on the login its parent test made. + +The parent passes LIGHTNING_CLOUD_URL, LIGHTNING_LOGIN_PATH and a shared +keyring through the environment. +""" +import time +import webbrowser + + +def _refuse_browser(url, *args, **kwargs): + raise AssertionError("a worker opened the real browser") + + +def call_the_api(barrier, seconds, results): + webbrowser.open = _refuse_browser + try: + from lightning_sdk.lightning_cloud.rest_client import GridRestClient + + client = GridRestClient() + barrier.wait(timeout=60) + deadline = time.monotonic() + seconds + calls = 0 + while time.monotonic() < deadline: + assert client.auth_service_get_user().id == "user-1" + calls += 1 + time.sleep(0.02) + results.put(("ok", calls)) + except BaseException as err: + results.put(("error", repr(err))) + raise + + +def take_a_token(results): + webbrowser.open = _refuse_browser + try: + from lightning_sdk.lightning_cloud import cli_login + + results.put(("ok", cli_login.token())) + except BaseException as err: + results.put(("error", repr(err))) + raise From b1fbac8c95398f76148f6852da6a2020e2533324 Mon Sep 17 00:00:00 2001 From: fnwoke-lightning Date: Wed, 4 Nov 2026 19:46:49 -0500 Subject: [PATCH 05/27] feat(api): attach credentials when each raw request is sent A 401 retries once after the source refreshes, and a folder download hands its workers that handler instead of a header captured once. Co-authored-by: Cursor --- python/lightning_sdk/api/billing_api.py | 4 +- python/lightning_sdk/api/filesystem_api.py | 15 +- python/lightning_sdk/api/llm_api.py | 62 ++---- python/lightning_sdk/api/logs_api.py | 2 +- python/lightning_sdk/api/sandbox_api.py | 16 +- python/lightning_sdk/api/studio_api.py | 21 +- python/lightning_sdk/api/teamspace_api.py | 17 +- python/lightning_sdk/api/utils.py | 49 ++++- .../lightning_sdk/cli/legacy/deploy/_auth.py | 21 +- .../cli/utils/lockdown_status.py | 12 +- python/tests/api/test_billing_api.py | 29 +-- python/tests/api/test_credential_requests.py | 183 ++++++++++++++++++ python/tests/api/test_filesystem_api.py | 82 ++------ python/tests/api/test_library_login.py | 102 ++++++++++ python/tests/api/test_studio_api.py | 80 ++------ python/tests/api/test_teamspace_api.py | 71 +------ python/tests/api/test_utils.py | 32 +-- .../core/filesystem/test_filesystem_copy.py | 16 +- .../filesystem/test_filesystem_listdir.py | 13 -- .../test_filesystem_teamspace_lookup.py | 5 - .../core/filesystem/test_filesystem_walk.py | 21 -- 21 files changed, 460 insertions(+), 393 deletions(-) create mode 100644 python/tests/api/test_credential_requests.py create mode 100644 python/tests/api/test_library_login.py diff --git a/python/lightning_sdk/api/billing_api.py b/python/lightning_sdk/api/billing_api.py index c5ca06e1d..d590b16d9 100644 --- a/python/lightning_sdk/api/billing_api.py +++ b/python/lightning_sdk/api/billing_api.py @@ -27,9 +27,9 @@ import requests from lightning_sdk.api.utils import ( - _authenticate_and_get_auth_headers, _raise_for_download_status, cached_lightning_client, + credential_auth, ) @@ -309,7 +309,7 @@ def _download_activity_csv( r = requests.get( f"{self._client.api_client.configuration.host}{endpoint}", params=query_params, - headers=_authenticate_and_get_auth_headers(), + auth=credential_auth(), stream=True, allow_redirects=True, ) diff --git a/python/lightning_sdk/api/filesystem_api.py b/python/lightning_sdk/api/filesystem_api.py index c42431a58..efa85e5cb 100644 --- a/python/lightning_sdk/api/filesystem_api.py +++ b/python/lightning_sdk/api/filesystem_api.py @@ -12,7 +12,6 @@ _DOWNLOAD_MAX_TRIES, _DOWNLOAD_READ_TIMEOUT_SECONDS, _RETRYABLE_DOWNLOAD_STATUSES, - _authenticate_and_get_auth_headers, _BlobUploader, _collect_download_results, _paged_tree_entries, @@ -22,6 +21,8 @@ _stream_download_to_file, _TransientDownloadError, cached_lightning_client, + credential_auth, + request_auth, ) from lightning_sdk.lightning_cloud.rest_client import LightningClient @@ -74,7 +75,7 @@ def fetch_page(query_params: Dict[str, str]) -> Dict: r = requests.get( f"{self._client.api_client.configuration.host}/v1/projects/{teamspace_id}/artifacts/trees/{path}", params=query_params, - headers=_authenticate_and_get_auth_headers(), + auth=credential_auth(), ) if r.status_code != 200: raise _RemoteApiError(f"Failed to list files: {r.status_code}", status_code=r.status_code) @@ -134,6 +135,7 @@ def _download_single_file( pbar: Optional[tqdm] = None, progress_bar: bool = False, expected_size: Optional[int] = None, + auth: Optional[object] = None, ) -> None: """Download a single artifact file. @@ -156,6 +158,7 @@ def _download_single_file( """ url = f"{self._client.api_client.configuration.host}/v1/projects/{teamspace_id}/artifacts/blobs/{remote_path}" owned_pbar = None + request = request_auth(auth) @backoff.on_exception(backoff.expo, _TransientDownloadError, max_tries=_DOWNLOAD_MAX_TRIES) def download() -> None: @@ -163,7 +166,7 @@ def download() -> None: try: r = requests.get( url, - headers=_authenticate_and_get_auth_headers(), + auth=request, stream=True, allow_redirects=True, timeout=(_DOWNLOAD_CONNECT_TIMEOUT_SECONDS, _DOWNLOAD_READ_TIMEOUT_SECONDS), @@ -252,6 +255,7 @@ def download_folder( mininterval=1, ) + auth = credential_auth() try: with ThreadPoolExecutor(max_workers=num_workers) as executor: futures = [ @@ -262,6 +266,7 @@ def download_folder( teamspace_id, pbar, expected_size=entry.get("size"), + auth=auth, ) for entry in files ] @@ -287,7 +292,7 @@ def delete_file(self, teamspace_id: str, remote_path: str) -> None: r = requests.delete( f"{self._client.api_client.configuration.host}" f"/v1/projects/{teamspace_id}/artifacts/blobs/{remote_path.strip('/')}", - headers=_authenticate_and_get_auth_headers(), + auth=credential_auth(), ) if r.status_code not in (200, 204): raise _RemoteApiError( @@ -310,7 +315,7 @@ def delete_folder(self, teamspace_id: str, remote_path: str) -> None: r = requests.delete( f"{self._client.api_client.configuration.host}" f"/v1/projects/{teamspace_id}/artifacts/trees/{remote_path.strip('/')}", - headers=_authenticate_and_get_auth_headers(), + auth=credential_auth(), ) if r.status_code not in (200, 204): raise _RemoteApiError( diff --git a/python/lightning_sdk/api/llm_api.py b/python/lightning_sdk/api/llm_api.py index 712bf1933..ea120ab8b 100644 --- a/python/lightning_sdk/api/llm_api.py +++ b/python/lightning_sdk/api/llm_api.py @@ -7,13 +7,11 @@ import time import warnings from typing import Any, AsyncGenerator, Dict, Generator, List, Optional, Union -from urllib.parse import urlencode from rich.prompt import Confirm from lightning_sdk.api.utils import cached_lightning_client -from lightning_sdk.lightning_cloud import env -from lightning_sdk.lightning_cloud.login import Auth, AuthServer +from lightning_sdk.lightning_cloud.credentials import CredentialSource, browser_login, start_login from lightning_sdk.lightning_cloud.openapi.models import ( StreamResultOfV1ConversationResponseChunk, V1ConversationResponseChunk, @@ -25,45 +23,27 @@ LITAI_CODE = os.environ.get("LITAI_CODE", "x334uv8t7v") -class _AuthServer(AuthServer): - def __init__(self, model: str, *args: Any, **kwargs: Any) -> None: - self._model = model - super().__init__(*args, **kwargs) - - def get_auth_url(self, port: int) -> str: - redirect_uri = f"http://localhost:{port}/login-complete" - params = urlencode({"redirectTo": redirect_uri, "okbhrt": LITAI_CODE, "litaimodel": self._model}) - return f"{env.LIGHTNING_CLOUD_URL}/sign-in?{params}" - - -class _AuthLitAI(Auth): - def __init__(self, model: str, shall_confirm: bool = False) -> None: - super().__init__() - self._model = model - self._shall_confirm = shall_confirm - - def _run_server(self) -> None: - if self._shall_confirm: - proceed = Confirm.ask( - "[bold yellow]LitAI needs to authenticate with Lightning AI.[/bold yellow]\n" - "This will open a browser window for login.\n" - "Do you want to continue?", - default=True, - ) - if not proceed: - raise RuntimeError( - "Login cancelled. Please login at https://platform.lightning.ai/sign-up. " - "Or run `lightning login` to login." - ) - print("Opening browser for authentication...") - print("Please come back to the terminal after logging in.") - time.sleep(1) - _AuthServer(self._model).login_with_browser(self) - - def authenticate(model: str, shall_confirm: bool = True) -> None: - auth = _AuthLitAI(model, shall_confirm) - auth.authenticate() + with browser_login(False): + existing = CredentialSource().resolve() + if existing is not None: + return + if shall_confirm: + proceed = Confirm.ask( + "[bold yellow]LitAI needs to authenticate with Lightning AI.[/bold yellow]\n" + "This will open a browser window for login.\n" + "Do you want to continue?", + default=True, + ) + if not proceed: + raise RuntimeError( + "Login cancelled. Please login at https://platform.lightning.ai/sign-up. " + "Or run `lightning login` to login." + ) + print("Opening browser for authentication...") + print("Please come back to the terminal after logging in.") + time.sleep(1) + start_login(sign_in_params={"okbhrt": LITAI_CODE, "litaimodel": model}) class LLMApi: diff --git a/python/lightning_sdk/api/logs_api.py b/python/lightning_sdk/api/logs_api.py index 7b8de8055..322b21076 100644 --- a/python/lightning_sdk/api/logs_api.py +++ b/python/lightning_sdk/api/logs_api.py @@ -520,7 +520,7 @@ def follow( while True: ws = None try: - auth_header = Auth().authenticate() + auth_header = Auth().auth_header # The handshake gets a longer budget than the deliberately short read poll. ws = websocket.create_connection( url, header=[f"Authorization: {auth_header}"], timeout=_CONNECT_TIMEOUT diff --git a/python/lightning_sdk/api/sandbox_api.py b/python/lightning_sdk/api/sandbox_api.py index 03c611e99..7c8029488 100644 --- a/python/lightning_sdk/api/sandbox_api.py +++ b/python/lightning_sdk/api/sandbox_api.py @@ -4,7 +4,7 @@ from typing import Any from lightning_sdk.lightning_cloud import env as lightning_env -from lightning_sdk.lightning_cloud.credentials import CredentialSource +from lightning_sdk.lightning_cloud.credentials import CredentialSource, browser_login from lightning_sdk.lightning_cloud.openapi import JobsServiceApi, SandboxesServiceApi from lightning_sdk.lightning_cloud.openapi.models import ( SandboxesServiceCreateSandboxDirectoryBody, @@ -228,9 +228,14 @@ def _apply_auth_and_host(self) -> None: def _ensure_auth(self) -> None: if self._auth_configured: return - self._client.api_client.request = credential_request_wrapper( - self._client.api_client.request, CredentialSource() - ) + source = CredentialSource() + wrapped = credential_request_wrapper(self._client.api_client.request, source) + + def request(*args: Any, **kwargs: Any) -> Any: + with browser_login(False): + return wrapped(*args, **kwargs) + + self._client.api_client.request = request self._auth_configured = True @property @@ -248,7 +253,8 @@ def auth_header(self) -> str: headers = getattr(self._client.api_client, "default_headers", {}) auth = headers.get("Authorization", "") if isinstance(headers, dict) else "" return str(auth or "") - return CredentialSource().header() + with browser_login(False): + return CredentialSource().header() def sandboxes(self) -> SandboxesServiceApi: self._ensure_auth() diff --git a/python/lightning_sdk/api/studio_api.py b/python/lightning_sdk/api/studio_api.py index 5552d46a6..7602dcd3d 100644 --- a/python/lightning_sdk/api/studio_api.py +++ b/python/lightning_sdk/api/studio_api.py @@ -9,7 +9,6 @@ from tqdm import tqdm from lightning_sdk.api.utils import ( - _authenticate_and_get_auth_headers, _BlobUploader, _collect_download_results, _DummyBody, @@ -20,6 +19,8 @@ _stream_download_to_file, _tree_path_info, cached_lightning_client, + credential_auth, + request_auth, ) from lightning_sdk.api.utils import ( _get_cloud_url as _cloud_url, @@ -1001,7 +1002,7 @@ def get_tree( r = requests.get( f"{self._client.api_client.configuration.host}/v1/projects/{teamspace_id}/artifacts/cloudspaces/{studio_id}/trees/{path}", params=query_params, - headers=_authenticate_and_get_auth_headers(), + auth=credential_auth(), ) return r.json() @@ -1111,7 +1112,7 @@ def download_file( r = requests.get( f"{self._client.api_client.configuration.host}/v1/projects/{teamspace_id}/artifacts/cloudspaces/{studio_id}/blobs/{path}", params=query_params, - headers=_authenticate_and_get_auth_headers(), + auth=credential_auth(), stream=True, allow_redirects=True, ) @@ -1139,7 +1140,7 @@ def _download_single_studio_file( download_dir: Path, studio_id: str, teamspace_id: str, - auth_headers: Dict[str, str], + auth: object, pbar: Optional[tqdm], ) -> None: """Download a single file from Studio with progress tracking. @@ -1150,7 +1151,7 @@ def _download_single_studio_file( download_dir: Local directory where the downloaded file is written. studio_id: Studio (cloud space) ID containing the file. teamspace_id: ID of the owning teamspace. - auth_headers: HTTP headers carrying the caller's credentials. + auth: Credential handler for this file's request. pbar: Optional tqdm progress bar to update as bytes are written. Raises: @@ -1164,7 +1165,7 @@ def _download_single_studio_file( r = requests.get( f"{self._client.api_client.configuration.host}/v1/projects/{teamspace_id}/artifacts/cloudspaces/{studio_id}/blobs/{file_path}", - headers=auth_headers, + auth=request_auth(auth), stream=True, ) @@ -1212,7 +1213,7 @@ def download_folder( print(f"No files found in {path}") return - auth_headers = _authenticate_and_get_auth_headers() + auth = credential_auth() total_size = sum(f.get("size", 0) for f in files) @@ -1237,7 +1238,7 @@ def download_folder( download_dir, studio_id, teamspace_id, - auth_headers, + auth, pbar, ) for file_info in files @@ -1275,7 +1276,7 @@ def remove_file(self, studio_id: str, teamspace_id: str, path: str) -> None: client_host = self._client.api_client.configuration.host url = f"{client_host}/v1/projects/{teamspace_id}/artifacts/cloudspaces/{studio_id}/blobs/{path}" - r = requests.delete(url, headers=_authenticate_and_get_auth_headers(), timeout=30) + r = requests.delete(url, auth=credential_auth(), timeout=30) if r.status_code == 204: return @@ -1306,7 +1307,7 @@ def remove_folder(self, studio_id: str, teamspace_id: str, path: str) -> None: client_host = self._client.api_client.configuration.host url = f"{client_host}/v1/projects/{teamspace_id}/artifacts/cloudspaces/{studio_id}/trees/{path}" - r = requests.delete(url, headers=_authenticate_and_get_auth_headers(), timeout=30) + r = requests.delete(url, auth=credential_auth(), timeout=30) if r.status_code == 204: return diff --git a/python/lightning_sdk/api/teamspace_api.py b/python/lightning_sdk/api/teamspace_api.py index fc1a15a08..711507cdc 100644 --- a/python/lightning_sdk/api/teamspace_api.py +++ b/python/lightning_sdk/api/teamspace_api.py @@ -11,7 +11,6 @@ from lightning_sdk.api.auth_api import AuthApi from lightning_sdk.api.utils import ( Experiment, - _authenticate_and_get_auth_headers, _BlobUploader, _collect_download_results, _download_model_files, @@ -23,6 +22,8 @@ _stream_download_to_file, _tree_path_info, cached_lightning_client, + credential_auth, + request_auth, ) from lightning_sdk.data_connection import BucketCredentials from lightning_sdk.lightning_cloud.openapi import ( @@ -764,7 +765,7 @@ def get_tree(self, teamspace_id: str, path: str, query_params: Optional[Dict[str r = requests.get( f"{self._client.api_client.configuration.host}/v1/projects/{teamspace_id}/artifacts/trees/{path}", params=query_params, - headers=_authenticate_and_get_auth_headers(), + auth=credential_auth(), ) return r.json() @@ -889,7 +890,7 @@ def download_file( r = requests.get( f"{self._client.api_client.configuration.host}/v1/projects/{teamspace_id}/artifacts/blobs/{path}", params=query_params, - headers=_authenticate_and_get_auth_headers(), + auth=credential_auth(), stream=True, ) @@ -915,7 +916,7 @@ def _download_single_file( base_path: str, download_dir: Path, teamspace_id: str, - auth_headers: Dict[str, str], + auth: object, cloud_account: Optional[str] = None, pbar: Optional[tqdm] = None, ) -> None: @@ -926,7 +927,7 @@ def _download_single_file( base_path: Base directory path prepended to the relative file path when building the request URL. download_dir: Local directory where the downloaded file is written. teamspace_id: ID of the owning teamspace. - auth_headers: HTTP headers carrying the caller's credentials. + auth: Credential handler for this file's request. cloud_account: Optional cloud account ID used to locate the artifact. pbar: Optional tqdm progress bar to update as bytes are written. @@ -946,7 +947,7 @@ def _download_single_file( r = requests.get( f"{self._client.api_client.configuration.host}/v1/projects/{teamspace_id}/artifacts/blobs/{file_path}", params=query_params, - headers=auth_headers, + auth=request_auth(auth), stream=True, ) @@ -990,7 +991,7 @@ def download_folder( print(f"No files found in {path}") return - auth_headers = _authenticate_and_get_auth_headers() + auth = credential_auth() total_size = sum(f.get("size", 0) for f in files) @@ -1014,7 +1015,7 @@ def download_folder( path, download_dir, teamspace_id, - auth_headers, + auth, cloud_account, pbar, ) diff --git a/python/lightning_sdk/api/utils.py b/python/lightning_sdk/api/utils.py index 61e97efad..825b076e7 100644 --- a/python/lightning_sdk/api/utils.py +++ b/python/lightning_sdk/api/utils.py @@ -29,10 +29,11 @@ import backoff import requests +from requests.auth import AuthBase from requests.exceptions import HTTPError from tqdm.auto import tqdm -from lightning_sdk.lightning_cloud.login import Auth +from lightning_sdk.lightning_cloud.credentials import CredentialSource from lightning_sdk.lightning_cloud.openapi import ( ModelsStoreCompleteMultiPartUploadBody, ModelsStoreCreateMultiPartUploadBody, @@ -298,7 +299,7 @@ def _post_blob_request(self, suffix: str, blob: Dict[str, Any], action: str) -> body: Dict[str, Any] = {"blobs": [blob]} if self.cluster_id: body["cluster_id"] = self.cluster_id - r = requests.post(url, json=body, headers=_authenticate_and_get_auth_headers(), timeout=30) + r = requests.post(url, json=body, auth=credential_auth(), timeout=30) if r.status_code >= 500 or r.status_code == 429: raise HTTPError( f"Transient error trying to {action} '{self.remote_path}'. Status code: {r.status_code}", response=r @@ -1271,6 +1272,44 @@ def to_iso_z(dt: datetime) -> str: return dt.isoformat(timespec="milliseconds") -def _authenticate_and_get_auth_headers() -> Dict[str, str]: - """Resolve the caller's credentials and return the HTTP headers carrying them.""" - return {"Authorization": Auth().authenticate()} +class CredentialAuth(AuthBase): + """Sets ``Authorization`` when a request is prepared, and retries one 401 after a refresh.""" + + def __init__(self, source: Optional[CredentialSource] = None) -> None: + self._source = source or CredentialSource() + + def __call__(self, request: requests.PreparedRequest) -> requests.PreparedRequest: + header = self._source.header() + request.headers["Authorization"] = header + sent = header[len("Bearer ") :] if header.startswith("Bearer ") else header + + def handle(response: requests.Response, **kwargs: Any) -> requests.Response: + if getattr(handle, "retried", False) or response.status_code != 401: + return response + if not self._source.refresh(sent): + return response + handle.retried = True + # The connection stays checked out until the body is read, so the retry can reuse it. + _ = response.content + response.close() + prepared = response.request.copy() + prepared.headers["Authorization"] = self._source.header() + retry = response.connection.send(prepared, **kwargs) + retry.history.append(response) + retry.request = prepared + return retry + + request.register_hook("response", handle) + return request + + +def credential_auth(source: Optional[CredentialSource] = None) -> CredentialAuth: + """A requests auth handler that reads :class:`CredentialSource` for each request.""" + return CredentialAuth(source) + + +def request_auth(auth: object = None) -> CredentialAuth: + """The auth object a download worker should use, or a fresh one when it wasn't given one.""" + if isinstance(auth, CredentialAuth): + return auth + return credential_auth() diff --git a/python/lightning_sdk/cli/legacy/deploy/_auth.py b/python/lightning_sdk/cli/legacy/deploy/_auth.py index fcbe61e11..2adb35ec8 100644 --- a/python/lightning_sdk/cli/legacy/deploy/_auth.py +++ b/python/lightning_sdk/cli/legacy/deploy/_auth.py @@ -11,7 +11,7 @@ from lightning_sdk.api.utils import cached_lightning_client from lightning_sdk.cli.utils.resource_resolution import resolve_teamspace from lightning_sdk.cli.utils.teamspace_option import resolve_teamspace as _resolve_teamspace_option -from lightning_sdk.lightning_cloud.login import Auth +from lightning_sdk.lightning_cloud.credentials import CredentialSource, browser_login from lightning_sdk.lightning_cloud.openapi import V1CloudSpace from lightning_sdk.utils.resolve import _get_authed_user @@ -23,24 +23,17 @@ class _AuthMode(Enum): DEPLOY = "deploy" -class _AuthLitServe(Auth): - def __init__(self, mode: _AuthMode, shall_confirm: bool = False) -> None: - super().__init__() - self._mode = mode - self._shall_confirm = shall_confirm - - def _run_server(self) -> None: - raise click.UsageError("No Lightning credentials are available. Run `lightning login` first.") - - def authenticate(mode: _AuthMode, shall_confirm: bool = True) -> None: """Validate existing Lightning AI credentials. Credentials must already be configured. Run ``lightning login`` first when - credentials are unavailable. + credentials are unavailable. This never opens a browser. """ - auth = _AuthLitServe(mode, shall_confirm) - auth.authenticate() + with browser_login(False): + try: + CredentialSource().header() + except ValueError as err: + raise click.UsageError("No Lightning credentials are available. Run `lightning login` first.") from err def select_teamspace(teamspace: Optional[str], org: Optional[str], user: Optional[str]) -> Teamspace: diff --git a/python/lightning_sdk/cli/utils/lockdown_status.py b/python/lightning_sdk/cli/utils/lockdown_status.py index 178096db4..a69753fda 100644 --- a/python/lightning_sdk/cli/utils/lockdown_status.py +++ b/python/lightning_sdk/cli/utils/lockdown_status.py @@ -80,16 +80,10 @@ def capability_closed(capability: str) -> bool: def _auth_header() -> Optional[str]: - """The caller's authorization header, or ``None`` when signed out. - - Reads the saved credentials unless the environment sets some, as the rest of the CLI does. - """ - auth = Auth() - if not (auth.api_key or auth.auth_token): - auth.load() + """The credential requests would send, or ``None`` when there isn't one.""" try: - return auth.get_auth_header() - except ValueError: + return Auth().auth_header + except (ValueError, ConnectionError): return None diff --git a/python/tests/api/test_billing_api.py b/python/tests/api/test_billing_api.py index 47e150c6b..ca09834a6 100644 --- a/python/tests/api/test_billing_api.py +++ b/python/tests/api/test_billing_api.py @@ -13,6 +13,7 @@ BillingResourceUsage, _build_activity_query_params, ) +from lightning_sdk.api.utils import CredentialAuth # ---- dataclasses: _from_api ------------------------------------------------ @@ -288,12 +289,8 @@ def test_get_activity_filter_values_with_project_id(mock_client): @mock.patch("requests.get", autospec=True) -@mock.patch( - "lightning_sdk.api.billing_api._authenticate_and_get_auth_headers", - return_value={"Authorization": "Bearer test-token"}, -) @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) -def test_get_session_activity_csv(mock_authenticate, mock_requests_get, tmp_path): +def test_get_session_activity_csv(mock_requests_get, tmp_path): mock_response = mock.Mock() mock_response.status_code = 200 mock_response.headers = {"content-length": "4"} @@ -307,21 +304,16 @@ def test_get_session_activity_csv(mock_authenticate, mock_requests_get, tmp_path assert result is None assert writer.getvalue() == "data" - mock_authenticate.assert_called_once_with() mock_requests_get.assert_called_once() call_args = mock_requests_get.call_args assert call_args[0][0].endswith("/v1/billing/usage-report/download/detailed") assert call_args[1]["params"] == {"orgId": "org-1"} - assert call_args[1]["headers"] == {"Authorization": "Bearer test-token"} + assert isinstance(call_args[1]["auth"], CredentialAuth) @mock.patch("requests.get", autospec=True) -@mock.patch( - "lightning_sdk.api.billing_api._authenticate_and_get_auth_headers", - return_value={"Authorization": "Bearer test-token"}, -) @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) -def test_get_session_activity_json_default_format(mock_authenticate, mock_requests_get, tmp_path): +def test_get_session_activity_json_default_format(mock_requests_get, tmp_path): mock_response = mock.Mock() mock_response.status_code = 200 mock_response.headers = {"content-length": "21"} @@ -340,12 +332,8 @@ def test_get_session_activity_json_default_format(mock_authenticate, mock_reques @mock.patch("requests.get", autospec=True) -@mock.patch( - "lightning_sdk.api.billing_api._authenticate_and_get_auth_headers", - return_value={"Authorization": "Bearer test-token"}, -) @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) -def test_get_resource_activity_csv(mock_authenticate, mock_requests_get, tmp_path): +def test_get_resource_activity_csv(mock_requests_get, tmp_path): mock_response = mock.Mock() mock_response.status_code = 200 mock_response.headers = {"content-length": "4"} @@ -365,19 +353,14 @@ def test_get_resource_activity_csv(mock_authenticate, mock_requests_get, tmp_pat assert result is None assert writer.getvalue() == "data" - mock_authenticate.assert_called_once_with() call_args = mock_requests_get.call_args assert call_args[0][0].endswith("/v1/billing/usage-report/download/summary") assert call_args[1]["params"] == {"orgId": "org-1", "projectIds": ["proj-1"], "limit": 5} @mock.patch("requests.get", autospec=True) -@mock.patch( - "lightning_sdk.api.billing_api._authenticate_and_get_auth_headers", - return_value={"Authorization": "Bearer test-token"}, -) @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) -def test_get_resource_activity_json_default_format(mock_authenticate, mock_requests_get, tmp_path): +def test_get_resource_activity_json_default_format(mock_requests_get, tmp_path): mock_response = mock.Mock() mock_response.status_code = 200 mock_response.headers = {"content-length": "21"} diff --git a/python/tests/api/test_credential_requests.py b/python/tests/api/test_credential_requests.py new file mode 100644 index 000000000..a44f5a4c6 --- /dev/null +++ b/python/tests/api/test_credential_requests.py @@ -0,0 +1,183 @@ +"""Raw requests take a credential handler, and a 401 is retried once with a fresh token.""" + +import threading +import webbrowser +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +import keyring +import pytest + +from lightning_sdk.api.billing_api import BillingApi +from lightning_sdk.api.filesystem_api import FilesystemApi +from lightning_sdk.api.studio_api import StudioApi +from lightning_sdk.api.teamspace_api import TeamspaceApi +from lightning_sdk.api.utils import _BlobUploader +from lightning_sdk.lightning_cloud import env +from lightning_sdk.lightning_cloud.credentials import CredentialSource, start_login +from lightning_sdk.lightning_cloud.login import Auth +from tests.lightning_cloud.utils.fake_platform import FakeBrowser, FakePlatform +from tests.lightning_cloud.utils.keyrings import MemoryKeyring + +_REAL_THREAD = threading.Thread + + +class _RetryServer: + """Answers 401 for ``old-token`` after rewriting the token file, then 200.""" + + def __init__(self, token_path): + seen = [] + + class Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_GET(self): # noqa: N802 + self._reply() + + def do_POST(self): # noqa: N802 + length = int(self.headers.get("Content-Length", "0")) + if length: + self.rfile.read(length) + self._reply() + + def _reply(self): + token = self.headers.get("Authorization", "").removeprefix("Bearer ").strip() + seen.append(token) + if token == "old-token": + token_path.write_text("new-token") + status, body = 401, b'{"error":"expired"}' + elif "usage-report" in self.path: + status, body = 200, b"col\n" + elif self.path.endswith("/blobs"): + status, body = 200, b'{"results":[]}' + else: + status, body = 200, b'{"tree":[]}' + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, fmt, *args): + return + + self.seen = seen + self._httpd = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + self.url = f"http://127.0.0.1:{self._httpd.server_address[1]}" + self._thread = _REAL_THREAD(target=self._httpd.serve_forever, daemon=True) + self._thread.start() + + def stop(self): + self._httpd.shutdown() + self._thread.join(timeout=5) + self._httpd.server_close() + + +@pytest.fixture() +def refreshed_token(monkeypatch, tmp_path): + threading.Thread = _REAL_THREAD + for name in ("LIGHTNING_AUTH_TOKEN", "LIGHTNING_API_KEY"): + monkeypatch.delenv(name, raising=False) + token_path = tmp_path / "token" + token_path.write_text("old-token") + monkeypatch.setenv("LIGHTNING_AUTH_TOKEN_FILE", str(token_path)) + server = _RetryServer(token_path) + monkeypatch.setenv("LIGHTNING_CLOUD_URL", server.url) + monkeypatch.setattr(env, "LIGHTNING_CLOUD_URL", server.url) + yield server + server.stop() + + +def test_filesystem_list_retries_a_401(refreshed_token): + FilesystemApi().list_files("ts", "data") + assert refreshed_token.seen == ["old-token", "new-token"] + + +def test_studio_tree_retries_a_401(refreshed_token): + StudioApi().get_tree("st", "ts", "data") + assert refreshed_token.seen == ["old-token", "new-token"] + + +def test_teamspace_tree_retries_a_401(refreshed_token): + TeamspaceApi().get_tree("ts", "data") + assert refreshed_token.seen == ["old-token", "new-token"] + + +def test_billing_csv_retries_a_401(refreshed_token): + BillingApi()._download_activity_csv("/v1/billing/usage-report/download/detailed", "org") + assert refreshed_token.seen == ["old-token", "new-token"] + + +def test_blob_upload_retries_a_401(refreshed_token, tmp_path): + local = tmp_path / "blob.bin" + local.write_bytes(b"hi") + uploader = _BlobUploader( + client=None, + endpoint_base=f"{refreshed_token.url}/v1/projects/ts/artifacts", + file_path=str(local), + remote_path="blob.bin", + progress_bar=False, + ) + uploader._post_blob_request("/blobs", {"path": "blob.bin"}, "request upload URLs for") + assert refreshed_token.seen == ["old-token", "new-token"] + + +def test_folder_download_refreshes_when_the_access_token_expires_midway(monkeypatch, tmp_path, inline_executor_cls): + threading.Thread = _REAL_THREAD + for name in ("LIGHTNING_AUTH_TOKEN", "LIGHTNING_API_KEY", "LIGHTNING_AUTH_TOKEN_FILE"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("LIGHTNING_CLI_OIDC_LOGIN", "1") + monkeypatch.setattr(Auth, "secrets_file", tmp_path / "credentials.json") + monkeypatch.setattr(env, "LIGHTNING_LOGIN_PATH", str(tmp_path / "login.json")) + previous = keyring.get_keyring() + keyring.set_keyring(MemoryKeyring()) + platform = FakePlatform() + monkeypatch.setattr(env, "LIGHTNING_CLOUD_URL", platform.url) + monkeypatch.setenv("LIGHTNING_CLOUD_URL", platform.url) + browser = FakeBrowser(platform) + monkeypatch.setattr(webbrowser, "open", browser.open) + seen = [] + + class Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_GET(self): # noqa: N802 + token = self.headers.get("Authorization", "").removeprefix("Bearer ").strip() + seen.append(token) + if len(seen) == 1: + platform.expire_access_tokens() + status, body = 401, b"expired" + elif platform.accepts(token): + status, body = 200, b"x" + else: + status, body = 401, b"rejected" + self.send_response(status) + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, fmt, *args): + return + + download = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + download_thread = _REAL_THREAD(target=download.serve_forever, daemon=True) + download_thread.start() + try: + start_login() + browser.finish() + first = CredentialSource().resolve().token + api = StudioApi() + api.list_files = lambda *_args, **_kwargs: [{"path": "a.bin", "size": 1}] + api._client.api_client.configuration.host = f"http://127.0.0.1:{download.server_address[1]}" + monkeypatch.setattr("lightning_sdk.api.studio_api.ThreadPoolExecutor", inline_executor_cls) + api.download_folder("folder", str(tmp_path / "out"), "st", "ts", "cluster", progress_bar=False) + assert seen[0] == first + assert seen[1] != first + assert platform.accepts(seen[1]) + assert (tmp_path / "out" / "a.bin").read_bytes() == b"x" + finally: + browser.finish() + platform.stop() + download.shutdown() + download_thread.join(timeout=5) + download.server_close() + keyring.set_keyring(previous) diff --git a/python/tests/api/test_filesystem_api.py b/python/tests/api/test_filesystem_api.py index 1678a2281..dcde53cbc 100644 --- a/python/tests/api/test_filesystem_api.py +++ b/python/tests/api/test_filesystem_api.py @@ -5,16 +5,12 @@ import requests from lightning_sdk.api.filesystem_api import FilesystemApi -from lightning_sdk.api.utils import _DOWNLOAD_MAX_TRIES +from lightning_sdk.api.utils import _DOWNLOAD_MAX_TRIES, CredentialAuth @mock.patch("requests.get", autospec=True) @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") -@mock.patch( - "lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers", - return_value={"Authorization": "Bearer token"}, -) -def test_download_file(_authenticate_mock, _mock_auth, mock_requests_get, tmpdir): +def test_download_file(_mock_auth, mock_requests_get, tmpdir): mock_response = mock.MagicMock() mock_response.status_code = 200 mock_response.headers = {"content-length": "0"} @@ -30,10 +26,8 @@ def test_download_file(_authenticate_mock, _mock_auth, mock_requests_get, tmpdir @mock.patch("lightning_sdk.api.filesystem_api._collect_download_results") @mock.patch("lightning_sdk.api.filesystem_api.tqdm") @mock.patch("lightning_sdk.api.filesystem_api.ThreadPoolExecutor") -@mock.patch("lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers") @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") -def test_download_folder(_mock_auth, authenticate_mock, mock_executor, mock_tqdm, mock_collect, tmpdir): - authenticate_mock.return_value = "test-token-123" +def test_download_folder(_mock_auth, mock_executor, mock_tqdm, mock_collect, tmpdir): filesystem_api = FilesystemApi() filesystem_api.list_files = mock.Mock( @@ -61,11 +55,7 @@ def test_download_folder(_mock_auth, authenticate_mock, mock_executor, mock_tqdm @mock.patch("requests.get", autospec=True) @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") -@mock.patch( - "lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers", - return_value={"Authorization": "Bearer token"}, -) -def test_download_file_rejects_an_error_document(_authenticate_mock, _mock_auth, mock_requests_get, tmp_path): +def test_download_file_rejects_an_error_document(_mock_auth, mock_requests_get, tmp_path): # An object store refusing a GET answers with an XML document, not with the bytes. body = b"InternalErrorWe encountered an internal error." mock_response = mock.Mock() @@ -87,11 +77,7 @@ def test_download_file_rejects_an_error_document(_authenticate_mock, _mock_auth, @mock.patch("requests.get", autospec=True) @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") -@mock.patch( - "lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers", - return_value={"Authorization": "Bearer token"}, -) -def test_download_file_retries_a_dropped_connection(_authenticate_mock, _mock_auth, mock_requests_get, tmp_path): +def test_download_file_retries_a_dropped_connection(_mock_auth, mock_requests_get, tmp_path): body = b"hello" ok_response = mock.Mock() ok_response.status_code = 200 @@ -111,11 +97,7 @@ def test_download_file_retries_a_dropped_connection(_authenticate_mock, _mock_au @mock.patch("requests.get", autospec=True) @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") -@mock.patch( - "lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers", - return_value={"Authorization": "Bearer token"}, -) -def test_download_file_retries_a_truncated_body(_authenticate_mock, _mock_auth, mock_requests_get, tmp_path): +def test_download_file_retries_a_truncated_body(_mock_auth, mock_requests_get, tmp_path): # A body shorter than Content-Length promised is a cut connection: retried, not fatal. body = b"hello" short_response = mock.Mock() @@ -140,11 +122,7 @@ def test_download_file_retries_a_truncated_body(_authenticate_mock, _mock_auth, @mock.patch("requests.get", autospec=True) @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") -@mock.patch( - "lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers", - return_value={"Authorization": "Bearer token"}, -) -def test_download_file_does_not_retry_a_denial(_authenticate_mock, _mock_auth, mock_requests_get, tmp_path): +def test_download_file_does_not_retry_a_denial(_mock_auth, mock_requests_get, tmp_path): mock_response = mock.Mock() mock_response.status_code = 404 mock_response.headers = {} @@ -159,14 +137,8 @@ def test_download_file_does_not_retry_a_denial(_authenticate_mock, _mock_auth, m assert mock_requests_get.call_count == 1 -@mock.patch( - "lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers", - return_value={"Authorization": "Bearer token"}, -) @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") -def test_download_folder_does_not_report_success_on_a_partial_download( - _mock_auth, _authenticate_mock, inline_executor_cls, tmp_path -): +def test_download_folder_does_not_report_success_on_a_partial_download(_mock_auth, inline_executor_cls, tmp_path): filesystem_api = FilesystemApi() filesystem_api.list_files = mock.Mock( return_value=[ @@ -190,11 +162,7 @@ def download(remote_path, *args, **kwargs): @mock.patch("requests.get", autospec=True) @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") -@mock.patch( - "lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers", - return_value={"Authorization": "Bearer token"}, -) -def test_list_files_returns_tree(_authenticate_mock, _mock_auth, mock_requests_get): +def test_list_files_returns_tree(_mock_auth, mock_requests_get): fake_files = [{"name": "test1.txt"}, {"name": "test2.txt"}] mock_response = mock.MagicMock() mock_response.status_code = 200 @@ -209,11 +177,7 @@ def test_list_files_returns_tree(_authenticate_mock, _mock_auth, mock_requests_g @mock.patch("requests.get", autospec=True) @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") -@mock.patch( - "lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers", - return_value={"Authorization": "Bearer token"}, -) -def test_list_files_returns_empty_when_tree_missing(_authenticate_mock, _mock_auth, mock_requests_get): +def test_list_files_returns_empty_when_tree_missing(_mock_auth, mock_requests_get): mock_response = mock.MagicMock() mock_response.status_code = 200 mock_response.json.return_value = {} @@ -227,11 +191,7 @@ def test_list_files_returns_empty_when_tree_missing(_authenticate_mock, _mock_au @mock.patch("requests.get", autospec=True) @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") -@mock.patch( - "lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers", - return_value={"Authorization": "Bearer token"}, -) -def test_list_files_passes_correct_url_and_params(_authenticate_mock, _mock_auth, mock_requests_get): +def test_list_files_passes_correct_url_and_params(_mock_auth, mock_requests_get): mock_response = mock.MagicMock() mock_response.status_code = 200 mock_response.json.return_value = {"tree": []} @@ -247,16 +207,12 @@ def test_list_files_passes_correct_url_and_params(_authenticate_mock, _mock_auth assert "ts-abc" in url assert "path/to/files" in url assert params["recursive"] == "true" - assert mock_requests_get.call_args[1]["headers"] == {"Authorization": "Bearer token"} + assert isinstance(mock_requests_get.call_args[1]["auth"], CredentialAuth) @mock.patch("requests.get", autospec=True) @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") -@mock.patch( - "lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers", - return_value={"Authorization": "Bearer token"}, -) -def test_list_files_non_recursive_by_default(_authenticate_mock, _mock_auth, mock_requests_get): +def test_list_files_non_recursive_by_default(_mock_auth, mock_requests_get): mock_response = mock.MagicMock() mock_response.status_code = 200 mock_response.json.return_value = {"tree": []} @@ -271,11 +227,7 @@ def test_list_files_non_recursive_by_default(_authenticate_mock, _mock_auth, moc @mock.patch("requests.get", autospec=True) @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") -@mock.patch( - "lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers", - return_value={"Authorization": "Bearer token"}, -) -def test_list_files_raises_on_non_200(_authenticate_mock, _mock_auth, mock_requests_get): +def test_list_files_raises_on_non_200(_mock_auth, mock_requests_get): mock_response = mock.MagicMock() mock_response.status_code = 404 mock_requests_get.return_value = mock_response @@ -287,11 +239,7 @@ def test_list_files_raises_on_non_200(_authenticate_mock, _mock_auth, mock_reque @mock.patch("requests.get", autospec=True) @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") -@mock.patch( - "lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers", - return_value={"Authorization": "Bearer token"}, -) -def test_list_files_follows_cursor_pages(_authenticate_mock, _mock_auth, mock_requests_get): +def test_list_files_follows_cursor_pages(_mock_auth, mock_requests_get): """A truncated listing is followed page by page until nextCursor comes back empty.""" pages = [ {"tree": [{"path": "a.txt"}], "truncated": True, "nextCursor": "cursor-1"}, diff --git a/python/tests/api/test_library_login.py b/python/tests/api/test_library_login.py new file mode 100644 index 000000000..e7e048821 --- /dev/null +++ b/python/tests/api/test_library_login.py @@ -0,0 +1,102 @@ +"""LitAI, LitServe and license validation stay on the credential source.""" + +import threading +import webbrowser + +import click +import pytest + +from lightning_sdk.api.license_api import LicenseApi +from lightning_sdk.api.llm_api import authenticate as litai_authenticate +from lightning_sdk.cli.legacy.deploy._auth import _AuthMode +from lightning_sdk.cli.legacy.deploy._auth import authenticate as litserve_authenticate +from lightning_sdk.lightning_cloud import env +from lightning_sdk.lightning_cloud.login import Auth +from tests.lightning_cloud.utils.fake_platform import FakeBrowser, FakePlatform + +_REAL_THREAD = threading.Thread + + +def _clear_credentials(monkeypatch, tmp_path): + threading.Thread = _REAL_THREAD + for name in ("LIGHTNING_AUTH_TOKEN", "LIGHTNING_API_KEY", "LIGHTNING_AUTH_TOKEN_FILE"): + monkeypatch.delenv(name, raising=False) + monkeypatch.delenv("LIGHTNING_CLI_OIDC_LOGIN", raising=False) + secrets = tmp_path / "credentials.json" + monkeypatch.setattr(Auth, "secrets_file", secrets) + return secrets + + +def test_litserve_does_not_open_a_browser_when_nothing_is_saved(monkeypatch, tmp_path): + _clear_credentials(monkeypatch, tmp_path) + + def opened(*_args, **_kwargs): + raise AssertionError("opened a browser") + + monkeypatch.setattr(webbrowser, "open", opened) + + with pytest.raises(click.UsageError, match="No Lightning credentials are available. Run `lightning login` first."): + litserve_authenticate(_AuthMode.DEPLOY) + + +def test_litserve_uses_an_existing_credential_without_a_browser(monkeypatch): + monkeypatch.setenv("LIGHTNING_AUTH_TOKEN", "env-token") + opened = [] + monkeypatch.setattr(webbrowser, "open", lambda *_args, **_kwargs: opened.append(True)) + + litserve_authenticate(_AuthMode.DEPLOY) + + assert opened == [] + + +def test_litai_puts_okbhrt_and_litaimodel_on_the_interim_sign_in(monkeypatch, tmp_path): + _clear_credentials(monkeypatch, tmp_path) + monkeypatch.setattr("lightning_sdk.api.llm_api.time.sleep", lambda _seconds: None) + monkeypatch.setattr("lightning_sdk.api.llm_api.Confirm.ask", lambda *_args, **_kwargs: True) + platform = FakePlatform() + monkeypatch.setattr(env, "LIGHTNING_CLOUD_URL", platform.url) + monkeypatch.setenv("LIGHTNING_CLOUD_URL", platform.url) + browser = FakeBrowser(platform) + monkeypatch.setattr(webbrowser, "open", browser.open) + try: + litai_authenticate("story-model") + query = platform.sign_ins[0] + assert query["okbhrt"] == "x334uv8t7v" + assert query["litaimodel"] == "story-model" + assert not platform.authorizations + finally: + browser.finish() + platform.stop() + + +def test_litai_does_nothing_when_a_credential_already_exists(monkeypatch): + monkeypatch.setenv("LIGHTNING_AUTH_TOKEN", "env-token") + asked = [] + monkeypatch.setattr( + "lightning_sdk.api.llm_api.Confirm.ask", + lambda *_args, **_kwargs: asked.append(True) or True, + ) + opened = [] + monkeypatch.setattr(webbrowser, "open", lambda *_args, **_kwargs: opened.append(True)) + + litai_authenticate("story-model") + + assert asked == [] + assert opened == [] + + +def test_license_validation_leaves_the_saved_login_unchanged(monkeypatch, tmp_path): + threading.Thread = _REAL_THREAD + secrets = tmp_path / "credentials.json" + original = b'{"auth_token": "sess_kept"}\n' + secrets.write_bytes(original) + monkeypatch.setattr(Auth, "secrets_file", secrets) + platform = FakePlatform() + platform.token_keys["license-token"] = "user-1" + monkeypatch.setattr(env, "LIGHTNING_CLOUD_URL", platform.url) + monkeypatch.setenv("LIGHTNING_CLOUD_URL", platform.url) + try: + LicenseApi("license-token") + assert secrets.read_bytes() == original + finally: + platform.stop() diff --git a/python/tests/api/test_studio_api.py b/python/tests/api/test_studio_api.py index aea360e7f..215db2554 100644 --- a/python/tests/api/test_studio_api.py +++ b/python/tests/api/test_studio_api.py @@ -7,6 +7,7 @@ import pytest from lightning_sdk.api.studio_api import StudioApi +from lightning_sdk.api.utils import CredentialAuth from lightning_sdk.lightning_cloud.openapi import ( CloudSpaceServiceCreateCloudSpaceBody, Externalv1CloudSpaceInstanceStatus, @@ -924,12 +925,8 @@ def test_duplicate_org_with_new_name( @mock.patch("requests.get", autospec=True) -@mock.patch( - "lightning_sdk.api.studio_api._authenticate_and_get_auth_headers", - return_value={"Authorization": "Bearer test-token"}, -) @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) -def test_get_tree(mock_authenticate, mock_requests_get): +def test_get_tree(mock_requests_get): """Test get_tree retrieves directory structure from studio.""" mock_response = mock.MagicMock() mock_response.json.return_value = { @@ -960,8 +957,7 @@ def test_get_tree(mock_authenticate, mock_requests_get): assert "/v1/projects/ts-abc/artifacts/cloudspaces/st-abc/trees/my-folder/" in call_args[0][0] assert call_args[1]["params"] is None - assert call_args[1]["headers"] == {"Authorization": "Bearer test-token"} - mock_authenticate.assert_called_once_with() + assert isinstance(call_args[1]["auth"], CredentialAuth) @pytest.mark.parametrize( @@ -1005,12 +1001,8 @@ def test_get_tree(mock_authenticate, mock_requests_get): ], ) @mock.patch("requests.get", autospec=True) -@mock.patch( - "lightning_sdk.api.studio_api._authenticate_and_get_auth_headers", - return_value={"Authorization": "Bearer test-token"}, -) @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) -def test_get_path_info(mock_authenticate, mock_requests_get, path, tree_response, expected_result): +def test_get_path_info(mock_requests_get, path, tree_response, expected_result): mock_response = mock.MagicMock() mock_response.json.return_value = tree_response mock_requests_get.return_value = mock_response @@ -1032,7 +1024,6 @@ def test_get_path_info(mock_authenticate, mock_requests_get, path, tree_response if path.strip("/") == "": mock_requests_get.assert_not_called() - mock_authenticate.assert_not_called() else: if "/" in path: expected_parent = path.rsplit("/", 1)[0] @@ -1041,7 +1032,6 @@ def test_get_path_info(mock_authenticate, mock_requests_get, path, tree_response # root level should include /trees/ call_url = mock_requests_get.call_args[0][0] assert "/trees/" in call_url - mock_authenticate.assert_called_once_with() @pytest.mark.parametrize( @@ -1104,13 +1094,8 @@ def test_get_path_info(mock_authenticate, mock_requests_get, path, tree_response ], ) @mock.patch("requests.get", autospec=True) -@mock.patch( - "lightning_sdk.api.studio_api._authenticate_and_get_auth_headers", - return_value={"Authorization": "Bearer test-token"}, -) @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_list_files( - mock_authenticate, mock_requests_get, path, mock_response, @@ -1140,7 +1125,6 @@ def test_list_files( assert call_args[1]["params"]["recursive"] == "true" assert result == expected_files - mock_authenticate.assert_called_once_with() def _blob_upload_create_response(path, upload_id, urls): @@ -1158,10 +1142,8 @@ def _blob_upload_create_response(path, upload_id, urls): @mock.patch("requests.post") @mock.patch("requests.put") @mock.patch("lightning_sdk.api.utils.tqdm") -@mock.patch("lightning_sdk.api.utils._authenticate_and_get_auth_headers") @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_upload_file( - authenticate_mock, tqdm_mock, requests_put_mock, requests_post_mock, @@ -1174,7 +1156,6 @@ def test_upload_file( # executor workers would never run; fall back to a plain map monkeypatch.setattr("lightning_sdk.api.utils.ThreadPoolExecutor.map", map) tqdm_mock.wrapattr.side_effect = lambda f, *args, **kwargs: f - authenticate_mock.return_value = {"Authorization": "Bearer test-token-123"} studio_api = StudioApi() @@ -1202,10 +1183,10 @@ def test_upload_file( # with the blob path in the body rather than the URL. create_call = requests_post_mock.call_args_list[0] assert create_call.args[0].endswith("/v1/projects/ts-abc/artifacts/cloudspaces/st-abc/blobs") - assert create_call.kwargs["headers"] == {"Authorization": "Bearer test-token-123"} + assert isinstance(create_call.kwargs["auth"], CredentialAuth) complete_call = requests_post_mock.call_args_list[1] assert complete_call.args[0].endswith("/v1/projects/ts-abc/artifacts/cloudspaces/st-abc/blobs/complete") - assert complete_call.kwargs["headers"] == {"Authorization": "Bearer test-token-123"} + assert isinstance(complete_call.kwargs["auth"], CredentialAuth) if single_part: assert create_call.kwargs["json"] == {"blobs": [{"path": "file1"}]} @@ -1236,10 +1217,8 @@ def test_upload_file( @mock.patch("requests.post") @mock.patch("requests.put") @mock.patch("lightning_sdk.api.utils.tqdm") -@mock.patch("lightning_sdk.api.utils._authenticate_and_get_auth_headers") @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_upload_file_leading_slash( - authenticate_mock, tqdm_mock, requests_put_mock, requests_post_mock, @@ -1249,7 +1228,6 @@ def test_upload_file_leading_slash( # A leading slash must not leak into the blob path sent to the upload # endpoints; both variants request the same clean path. tqdm_mock.wrapattr.side_effect = lambda f, *args, **kwargs: f - authenticate_mock.return_value = {"Authorization": "Bearer test-token-123"} studio_api = StudioApi() @@ -1271,11 +1249,8 @@ def test_upload_file_leading_slash( @mock.patch("requests.get", autospec=True) -@mock.patch( - "lightning_sdk.api.studio_api._authenticate_and_get_auth_headers", return_value={"Authorization": "Bearer token"} -) @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) -def test_download_file(mock_authenticate, mock_requests_get, tmpdir): +def test_download_file(mock_requests_get, tmpdir): mock_response = mock.Mock() mock_response.status_code = 200 mock_response.headers = {"content-length": "1024"} @@ -1286,15 +1261,11 @@ def test_download_file(mock_authenticate, mock_requests_get, tmpdir): filepath = os.path.join(tmpdir, "file1") studio_api.download_file("file1", filepath, "st-abc", "ts-abc", "cluster-abc") - mock_authenticate.assert_called_once_with() @mock.patch("requests.get", autospec=True) -@mock.patch( - "lightning_sdk.api.studio_api._authenticate_and_get_auth_headers", return_value={"Authorization": "Bearer token"} -) @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) -def test_download_file_non_200_status(mock_authenticate, mock_requests_get, tmpdir): +def test_download_file_non_200_status(mock_requests_get, tmpdir): mock_response = mock.Mock() mock_response.status_code = 404 mock_response.headers = {} @@ -1308,17 +1279,13 @@ def test_download_file_non_200_status(mock_authenticate, mock_requests_get, tmpd studio_api.download_file("file1", filepath, "st-abc", "ts-abc", "cluster-abc") assert not os.path.exists(filepath) - mock_authenticate.assert_called_once_with() @mock.patch("lightning_sdk.api.studio_api._collect_download_results") @mock.patch("lightning_sdk.api.studio_api.tqdm") @mock.patch("lightning_sdk.api.studio_api.ThreadPoolExecutor") -@mock.patch("lightning_sdk.api.studio_api._authenticate_and_get_auth_headers") @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) -def test_download_folder(authenticate_mock, mock_executor, mock_tqdm, mock_collect, tmpdir): - authenticate_mock.return_value = {"Authorization": "Bearer test-token-123"} - +def test_download_folder(mock_executor, mock_tqdm, mock_collect, tmpdir): studio_api = StudioApi() studio_api.list_files = mock.Mock( @@ -1393,13 +1360,8 @@ def test_download_single_studio_file_rejects_an_error_document_under_200(mock_re assert list(tmp_path.iterdir()) == [] -@mock.patch( - "lightning_sdk.api.studio_api._authenticate_and_get_auth_headers", return_value={"Authorization": "Bearer token"} -) @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) -def test_download_folder_does_not_report_success_on_a_partial_download( - _mock_authenticate, inline_executor_cls, tmp_path -): +def test_download_folder_does_not_report_success_on_a_partial_download(inline_executor_cls, tmp_path): studio_api = StudioApi() studio_api.list_files = mock.Mock( return_value=[ @@ -2065,15 +2027,10 @@ def test_get_port_url(mock_list_endpoints, port, name, endpoints, expected_url, ], ) @mock.patch("requests.delete", autospec=True) -@mock.patch( - "lightning_sdk.api.studio_api._authenticate_and_get_auth_headers", - return_value={"Authorization": "Bearer test-token"}, -) @mock.patch("lightning_sdk.api.studio_api.StudioApi.get_path_info", autospec=True) @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_remove_file( mock_get_path_info, - mock_authenticate, mock_requests_delete, path, path_info, @@ -2103,9 +2060,8 @@ def test_remove_file( call_args = mock_requests_delete.call_args assert f"/v1/projects/ts-abc/artifacts/cloudspaces/st-abc/blobs/{path}" in call_args[0][0] - assert call_args[1]["headers"] == {"Authorization": "Bearer test-token"} + assert isinstance(call_args[1]["auth"], CredentialAuth) assert call_args[1]["timeout"] == 30 - mock_authenticate.assert_called_once_with() mock_get_path_info.assert_called_once_with(mock.ANY, "st-abc", "ts-abc", path=path) @@ -2147,15 +2103,10 @@ def test_remove_file( ], ) @mock.patch("requests.delete", autospec=True) -@mock.patch( - "lightning_sdk.api.studio_api._authenticate_and_get_auth_headers", - return_value={"Authorization": "Bearer test-token"}, -) @mock.patch("lightning_sdk.api.studio_api.StudioApi.get_path_info", autospec=True) @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_remove_folder( mock_get_path_info, - mock_authenticate, mock_requests_delete, path, path_info, @@ -2185,20 +2136,15 @@ def test_remove_folder( call_args = mock_requests_delete.call_args assert f"/v1/projects/ts-abc/artifacts/cloudspaces/st-abc/trees/{path}" in call_args[0][0] - assert call_args[1]["headers"] == {"Authorization": "Bearer test-token"} + assert isinstance(call_args[1]["auth"], CredentialAuth) assert call_args[1]["timeout"] == 30 - mock_authenticate.assert_called_once_with() mock_get_path_info.assert_called_once_with(mock.ANY, "st-abc", "ts-abc", path=path) @mock.patch("requests.get", autospec=True) -@mock.patch( - "lightning_sdk.api.studio_api._authenticate_and_get_auth_headers", - return_value={"Authorization": "Bearer token"}, -) @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) -def test_list_files_follows_cursor_pages(_mock_authenticate, mock_requests_get): +def test_list_files_follows_cursor_pages(mock_requests_get): """A truncated listing is followed page by page until nextCursor comes back empty.""" pages = [ {"tree": [{"path": "a.txt", "type": "blob"}], "truncated": True, "nextCursor": "cursor-1"}, diff --git a/python/tests/api/test_teamspace_api.py b/python/tests/api/test_teamspace_api.py index 2d96ee92e..fc698eb75 100644 --- a/python/tests/api/test_teamspace_api.py +++ b/python/tests/api/test_teamspace_api.py @@ -8,6 +8,7 @@ from lightning_sdk import Machine from lightning_sdk.api.teamspace_api import SecretType, TeamspaceApi +from lightning_sdk.api.utils import CredentialAuth from lightning_sdk.lightning_cloud.openapi import ( ModelsStoreCreateModelBody, ModelsStoreCreateModelVersionBody, @@ -440,12 +441,8 @@ def test_get_model_errors(internal_teamspace_api_mocker): @mock.patch("requests.get", autospec=True) -@mock.patch( - "lightning_sdk.api.teamspace_api._authenticate_and_get_auth_headers", - return_value={"Authorization": "Bearer test-token"}, -) @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) -def test_get_tree(mock_authenticate, mock_requests_get): +def test_get_tree(mock_requests_get): """Test get_tree retrieves directory structure from teamspace drive.""" mock_response = mock.MagicMock() mock_response.json.return_value = { @@ -476,8 +473,7 @@ def test_get_tree(mock_authenticate, mock_requests_get): assert "/v1/projects/ts-abc/artifacts/trees/my-folder/" in call_args[0][0] assert call_args[1]["params"] is None - assert call_args[1]["headers"] == {"Authorization": "Bearer test-token"} - mock_authenticate.assert_called_once_with() + assert isinstance(call_args[1]["auth"], CredentialAuth) @pytest.mark.parametrize( @@ -521,12 +517,8 @@ def test_get_tree(mock_authenticate, mock_requests_get): ], ) @mock.patch("requests.get", autospec=True) -@mock.patch( - "lightning_sdk.api.teamspace_api._authenticate_and_get_auth_headers", - return_value={"Authorization": "Bearer test-token"}, -) @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) -def test_get_path_info(mock_authenticate, mock_requests_get, path, tree_response, expected_result): +def test_get_path_info(mock_requests_get, path, tree_response, expected_result): mock_response = mock.MagicMock() mock_response.json.return_value = tree_response mock_requests_get.return_value = mock_response @@ -548,7 +540,6 @@ def test_get_path_info(mock_authenticate, mock_requests_get, path, tree_response if path.strip("/") == "": mock_requests_get.assert_not_called() - mock_authenticate.assert_not_called() else: if "/" in path: expected_parent = path.rsplit("/", 1)[0] @@ -557,7 +548,6 @@ def test_get_path_info(mock_authenticate, mock_requests_get, path, tree_response # root level should include /trees/ call_url = mock_requests_get.call_args[0][0] assert "/trees/" in call_url - mock_authenticate.assert_called_once_with() @pytest.mark.parametrize( @@ -620,13 +610,8 @@ def test_get_path_info(mock_authenticate, mock_requests_get, path, tree_response ], ) @mock.patch("requests.get", autospec=True) -@mock.patch( - "lightning_sdk.api.teamspace_api._authenticate_and_get_auth_headers", - return_value={"Authorization": "Bearer test-token"}, -) @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_list_files( - mock_authenticate, mock_requests_get, path, mock_response, @@ -655,7 +640,6 @@ def test_list_files( assert call_args[1]["params"]["recursive"] == "true" assert result == expected_files - mock_authenticate.assert_called_once_with() def _blob_upload_create_response(path, upload_id, urls): @@ -680,10 +664,6 @@ def _make_upload_teamspace_api(): @mock.patch("requests.post") @mock.patch("requests.put") @mock.patch("lightning_sdk.api.utils.tqdm") -@mock.patch( - "lightning_sdk.api.utils._authenticate_and_get_auth_headers", - new=mock.MagicMock(return_value={"Authorization": "Bearer test-token"}), -) @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_upload_file( tqdm_mock, @@ -723,7 +703,7 @@ def test_upload_file( # URL request goes to the project-scoped batch endpoint with the cluster in the body. create_call = requests_post_mock.call_args_list[0] assert create_call.args[0] == "https://api.example.com/v1/projects/ts-abc/artifacts/blobs" - assert create_call.kwargs["headers"] == {"Authorization": "Bearer test-token"} + assert isinstance(create_call.kwargs["auth"], CredentialAuth) if single_part: assert create_call.kwargs["json"] == {"cluster_id": "cluster-abc", "blobs": [{"path": "file1"}]} @@ -757,10 +737,6 @@ def test_upload_file( @pytest.mark.parametrize("remote_path", ["uploads/file1", "Uploads/file1"]) @mock.patch("requests.post") @mock.patch("requests.put") -@mock.patch( - "lightning_sdk.api.utils._authenticate_and_get_auth_headers", - new=mock.MagicMock(return_value={"Authorization": "Bearer test-token"}), -) @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_upload_file_uploads_paths_pass_through( requests_put_mock, @@ -790,10 +766,6 @@ def test_upload_file_uploads_paths_pass_through( @mock.patch("requests.post") @mock.patch("requests.put") @mock.patch("lightning_sdk.api.utils.tqdm") -@mock.patch( - "lightning_sdk.api.utils._authenticate_and_get_auth_headers", - new=mock.MagicMock(return_value={"Authorization": "Bearer test-token"}), -) @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_upload_file_with_headers( tqdm_mock, @@ -837,10 +809,6 @@ def test_upload_file_with_headers( @mock.patch("requests.post") @mock.patch("requests.put") -@mock.patch( - "lightning_sdk.api.utils._authenticate_and_get_auth_headers", - new=mock.MagicMock(return_value={"Authorization": "Bearer test-token"}), -) @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) def test_upload_file_without_headers( requests_put_mock, @@ -882,11 +850,8 @@ def _stub_response(status_code: int, body: bytes) -> mock.Mock: @mock.patch("requests.get", autospec=True) -@mock.patch( - "lightning_sdk.api.teamspace_api._authenticate_and_get_auth_headers", return_value={"Authorization": "Bearer token"} -) @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) -def test_download_file(mock_authenticate, mock_requests_get, tmpdir): +def test_download_file(mock_requests_get, tmpdir): mock_requests_get.return_value = _stub_response(200, b"data" * 256) teamspace_api = TeamspaceApi() @@ -894,15 +859,11 @@ def test_download_file(mock_authenticate, mock_requests_get, tmpdir): filepath = os.path.join(tmpdir, "file1") teamspace_api.download_file("file1", filepath, "ts-abc", "cluster-abc") assert Path(filepath).read_bytes() == b"data" * 256 - mock_authenticate.assert_called_once_with() @mock.patch("requests.get", autospec=True) -@mock.patch( - "lightning_sdk.api.teamspace_api._authenticate_and_get_auth_headers", return_value={"Authorization": "Bearer token"} -) @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) -def test_download_file_rejects_an_error_document(_mock_authenticate, mock_requests_get, tmpdir): +def test_download_file_rejects_an_error_document(mock_requests_get, tmpdir): mock_requests_get.return_value = _stub_response(403, _S3_ERROR_BODY) teamspace_api = TeamspaceApi() @@ -930,13 +891,8 @@ def test_download_single_file_rejects_an_error_document_under_200(mock_requests_ assert list(tmp_path.iterdir()) == [] -@mock.patch( - "lightning_sdk.api.teamspace_api._authenticate_and_get_auth_headers", return_value={"Authorization": "Bearer token"} -) @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) -def test_download_folder_does_not_report_success_on_a_partial_download( - _mock_authenticate, inline_executor_cls, tmp_path -): +def test_download_folder_does_not_report_success_on_a_partial_download(inline_executor_cls, tmp_path): teamspace_api = TeamspaceApi() teamspace_api.list_files = mock.Mock( return_value=[{"path": "ok.bin", "size": 1000}, {"path": "single_2mib.bin", "size": 2 * 1024 * 1024}] @@ -958,11 +914,8 @@ def download(file_info, *args, **kwargs): @mock.patch("lightning_sdk.api.teamspace_api._collect_download_results") @mock.patch("lightning_sdk.api.teamspace_api.tqdm") @mock.patch("lightning_sdk.api.teamspace_api.ThreadPoolExecutor") -@mock.patch("lightning_sdk.api.teamspace_api._authenticate_and_get_auth_headers") @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) -def test_download_folder(authenticate_mock, mock_executor, mock_tqdm, mock_collect, tmpdir): - authenticate_mock.return_value = {"Authorization": "Bearer test-token-123"} - +def test_download_folder(mock_executor, mock_tqdm, mock_collect, tmpdir): teamspace_api = TeamspaceApi() teamspace_api.list_files = mock.Mock( @@ -1197,12 +1150,8 @@ def test_verify_secret_name(secret_name, expected): @mock.patch("requests.get", autospec=True) -@mock.patch( - "lightning_sdk.api.teamspace_api._authenticate_and_get_auth_headers", - return_value={"Authorization": "Bearer token"}, -) @mock.patch("lightning_sdk.lightning_cloud.credentials.Auth", new=mock.MagicMock()) -def test_list_files_follows_cursor_pages(_mock_authenticate, mock_requests_get): +def test_list_files_follows_cursor_pages(mock_requests_get): """A truncated listing is followed page by page until nextCursor comes back empty.""" pages = [ {"tree": [{"path": "a.txt", "type": "blob"}], "truncated": True, "nextCursor": "cursor-1"}, diff --git a/python/tests/api/test_utils.py b/python/tests/api/test_utils.py index b1befbea2..233df261e 100644 --- a/python/tests/api/test_utils.py +++ b/python/tests/api/test_utils.py @@ -9,6 +9,7 @@ import lightning_sdk.api.utils from lightning_sdk.api import utils from lightning_sdk.api.utils import ( + CredentialAuth, _BlobUploader, _collect_download_results, _download_model_files, @@ -37,9 +38,6 @@ def _make_mocked_blob_uploader(monkeypatch, file_path, remote_path, **kwargs): # Threadpools don't like mocks as input, so we just use a regular map here monkeypatch.setattr(lightning_sdk.api.utils.ThreadPoolExecutor, "map", map) - monkeypatch.setattr( - lightning_sdk.api.utils, "_authenticate_and_get_auth_headers", lambda: {"Authorization": "Bearer test-token"} - ) return _BlobUploader( client=Mock(), endpoint_base=_TEST_ENDPOINT_BASE, @@ -113,7 +111,7 @@ def test_blob_uploader_multipart(tmp_path, monkeypatch): create_call = post_mock.call_args_list[0] assert create_call.args[0] == f"{_TEST_ENDPOINT_BASE}/blobs" - assert create_call.kwargs["headers"] == {"Authorization": "Bearer test-token"} + assert isinstance(create_call.kwargs["auth"], CredentialAuth) assert create_call.kwargs["json"] == { "cluster_id": "test-cluster-id", "blobs": [{"path": "path/to/file/on/remote", "parts": 3, "part_size": 4}], @@ -790,16 +788,6 @@ def test_resolve_path_mappings(): assert path_mappings[1].connection_path == "" -@pytest.fixture() -def upload_auth(monkeypatch): - headers = {"Authorization": "Bearer test-token"} - monkeypatch.setattr( - "lightning_sdk.api.utils._authenticate_and_get_auth_headers", - lambda: headers, - ) - return headers - - def _make_single_part_blob_uploader(tmp_path, progress_bar=False, extra_headers=None, content=b"hello world"): file_path = tmp_path / "test_file.txt" file_path.write_bytes(content) @@ -818,7 +806,7 @@ def _configure_single_part_create(mock_requests, url="http://example.com/upload" @mock.patch("lightning_sdk.api.utils.requests") -def test_single_part_uploader_basic(mock_requests, tmp_path, upload_auth): +def test_single_part_uploader_basic(mock_requests, tmp_path): uploader = _make_single_part_blob_uploader(tmp_path) _configure_single_part_create(mock_requests) @@ -830,7 +818,7 @@ def test_single_part_uploader_basic(mock_requests, tmp_path, upload_auth): assert mock_requests.post.call_count == 1 post_call = mock_requests.post.call_args assert post_call.args[0] == f"{_TEST_ENDPOINT_BASE}/blobs" - assert post_call.kwargs["headers"] == {"Authorization": "Bearer test-token"} + assert isinstance(post_call.kwargs["auth"], CredentialAuth) # one PUT straight to the signed storage URL, with no token assert mock_requests.put.call_count == 1 @@ -842,7 +830,7 @@ def test_single_part_uploader_basic(mock_requests, tmp_path, upload_auth): @mock.patch("lightning_sdk.api.utils.requests") -def test_single_part_uploader_with_progress_bar(mock_requests, tmp_path, upload_auth): +def test_single_part_uploader_with_progress_bar(mock_requests, tmp_path): uploader = _make_single_part_blob_uploader(tmp_path, progress_bar=True) _configure_single_part_create(mock_requests) @@ -858,7 +846,7 @@ def test_single_part_uploader_with_progress_bar(mock_requests, tmp_path, upload_ @mock.patch("time.sleep", return_value=None) @mock.patch("lightning_sdk.api.utils.requests") -def test_single_part_uploader_retries_on_http_error(mock_requests, _, tmp_path, upload_auth): +def test_single_part_uploader_retries_on_http_error(mock_requests, _, tmp_path): uploader = _make_single_part_blob_uploader(tmp_path) _configure_single_part_create(mock_requests) @@ -876,7 +864,7 @@ def test_single_part_uploader_retries_on_http_error(mock_requests, _, tmp_path, @mock.patch("time.sleep", return_value=None) @mock.patch("lightning_sdk.api.utils.requests") -def test_single_part_uploader_retries_on_transient_status(mock_requests, _, tmp_path, upload_auth): +def test_single_part_uploader_retries_on_transient_status(mock_requests, _, tmp_path): uploader = _make_single_part_blob_uploader(tmp_path) _configure_single_part_create(mock_requests) @@ -893,7 +881,7 @@ def test_single_part_uploader_retries_on_transient_status(mock_requests, _, tmp_ @mock.patch("time.sleep", return_value=None) @mock.patch("lightning_sdk.api.utils.requests") -def test_single_part_uploader_retries_auth_errors_with_fresh_url(mock_requests, _, tmp_path, upload_auth): +def test_single_part_uploader_retries_auth_errors_with_fresh_url(mock_requests, _, tmp_path): uploader = _make_single_part_blob_uploader(tmp_path) _configure_single_part_create(mock_requests) @@ -917,9 +905,7 @@ def test_single_part_uploader_retries_auth_errors_with_fresh_url(mock_requests, @pytest.mark.parametrize("content", [b"hello world", b""]) @mock.patch("time.sleep", return_value=None) @mock.patch("lightning_sdk.api.utils.requests") -def test_single_part_uploader_does_not_retry_client_error( - mock_requests, _, tmp_path, content, status_code, upload_auth -): +def test_single_part_uploader_does_not_retry_client_error(mock_requests, _, tmp_path, content, status_code): uploader = _make_single_part_blob_uploader(tmp_path, content=content) _configure_single_part_create(mock_requests) diff --git a/python/tests/core/filesystem/test_filesystem_copy.py b/python/tests/core/filesystem/test_filesystem_copy.py index d8045c80e..55e8eabd4 100644 --- a/python/tests/core/filesystem/test_filesystem_copy.py +++ b/python/tests/core/filesystem/test_filesystem_copy.py @@ -6,7 +6,6 @@ from lightning_sdk.filesystem import Filesystem TEAMSPACE_ID = "ts-123" -FAKE_AUTH_HEADERS = {"Authorization": "Bearer fake-token"} HOST = "https://lightning.ai" REMOTE_PATH = "data/model.ckpt" LIT_URL = f"lit://my-org/my-teamspace/{REMOTE_PATH}" @@ -31,12 +30,9 @@ def fake_path_result(): @mock.patch("lightning_sdk.api.filesystem_api.requests.get") @mock.patch("lightning_sdk.api.utils.LightningClient") -@mock.patch("lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers") @mock.patch("lightning_sdk.cli.utils.filesystem.resolve_teamspace_ref") @mock.patch("lightning_sdk.cli.utils.filesystem.parse_lit_url") -def test_copy_download_file( - mock_parse_lit_url, mock_resolve, mock_authenticate, mock_client_cls, mock_get, fake_teamspace, tmp_path -): +def test_copy_download_file(mock_parse_lit_url, mock_resolve, mock_client_cls, mock_get, fake_teamspace, tmp_path): remote_path = "data/model.ckpt" mock_parse_lit_url.return_value = { "teamspace": "my-teamspace", @@ -44,7 +40,6 @@ def test_copy_download_file( "destination": remote_path, } mock_resolve.return_value = fake_teamspace - mock_authenticate.return_value = FAKE_AUTH_HEADERS mock_client_cls.return_value.api_client.configuration.host = HOST def fake_get(url, **kwargs): @@ -73,13 +68,11 @@ def fake_get(url, **kwargs): @mock.patch("lightning_sdk.api.filesystem_api.ThreadPoolExecutor") @mock.patch("lightning_sdk.api.filesystem_api.requests.get") @mock.patch("lightning_sdk.api.utils.LightningClient") -@mock.patch("lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers") @mock.patch("lightning_sdk.cli.utils.filesystem.resolve_teamspace_ref") @mock.patch("lightning_sdk.cli.utils.filesystem.parse_lit_url") def test_copy_download_folder( mock_parse_lit_url, mock_resolve, - mock_authenticate, mock_client_cls, mock_get, mock_executor, @@ -93,7 +86,6 @@ def test_copy_download_folder( "destination": "data/mydir", } mock_resolve.return_value = fake_teamspace - mock_authenticate.return_value = FAKE_AUTH_HEADERS mock_client_cls.return_value.api_client.configuration.host = HOST mock_collect.return_value = None @@ -120,11 +112,10 @@ def fake_get(url, **kwargs): @mock.patch("lightning_sdk.api.filesystem_api.requests.get") @mock.patch("lightning_sdk.api.utils.LightningClient") -@mock.patch("lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers", return_value=FAKE_AUTH_HEADERS) @mock.patch("lightning_sdk.cli.utils.filesystem.resolve_teamspace_ref") @mock.patch("lightning_sdk.cli.utils.filesystem.parse_lit_url") def test_copy_raises_if_directory_without_recursive( - mock_parse_lit_url, mock_resolve, _mock_authenticate, mock_client_cls, mock_get, fake_teamspace + mock_parse_lit_url, mock_resolve, mock_client_cls, mock_get, fake_teamspace ): mock_parse_lit_url.return_value = { "teamspace": "my-teamspace", @@ -143,11 +134,10 @@ def test_copy_raises_if_directory_without_recursive( @mock.patch("lightning_sdk.api.filesystem_api.requests.get") @mock.patch("lightning_sdk.api.utils.LightningClient") -@mock.patch("lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers", return_value=FAKE_AUTH_HEADERS) @mock.patch("lightning_sdk.cli.utils.filesystem.resolve_teamspace_ref") @mock.patch("lightning_sdk.cli.utils.filesystem.parse_lit_url") def test_copy_raises_if_remote_file_not_found( - mock_parse_lit_url, mock_resolve, _mock_authenticate, mock_client_cls, mock_get, fake_teamspace, fake_path_result + mock_parse_lit_url, mock_resolve, mock_client_cls, mock_get, fake_teamspace, fake_path_result ): mock_parse_lit_url.return_value = fake_path_result mock_resolve.return_value = fake_teamspace diff --git a/python/tests/core/filesystem/test_filesystem_listdir.py b/python/tests/core/filesystem/test_filesystem_listdir.py index a292e4cfc..d782e0df7 100644 --- a/python/tests/core/filesystem/test_filesystem_listdir.py +++ b/python/tests/core/filesystem/test_filesystem_listdir.py @@ -5,7 +5,6 @@ from lightning_sdk.filesystem import Filesystem TEAMSPACE_ID = "ts-123" -FAKE_AUTH_HEADERS = {"Authorization": "Bearer fake-token"} HOST = "https://lightning.ai" REMOTE_PATH = "data/test1.txt" LIT_URL = f"lit://my-org/my-teamspace/{REMOTE_PATH}" @@ -30,10 +29,6 @@ def fake_path_result(): @mock.patch("lightning_sdk.api.filesystem_api.requests.get") @mock.patch("lightning_sdk.api.utils.LightningClient") -@mock.patch( - "lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers", - new=mock.MagicMock(return_value=FAKE_AUTH_HEADERS), -) @mock.patch("lightning_sdk.cli.utils.filesystem.resolve_teamspace_ref") @mock.patch("lightning_sdk.cli.utils.filesystem.parse_lit_url") def test_listdir_returns_files( @@ -57,10 +52,6 @@ def test_listdir_returns_files( @mock.patch("lightning_sdk.api.filesystem_api.requests.get") @mock.patch("lightning_sdk.api.utils.LightningClient") -@mock.patch( - "lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers", - new=mock.MagicMock(return_value=FAKE_AUTH_HEADERS), -) @mock.patch("lightning_sdk.cli.utils.filesystem.resolve_teamspace_ref") @mock.patch("lightning_sdk.cli.utils.filesystem.parse_lit_url") def test_listdir_passes_correct_teamspace_id(mock_parse_lit_url, mock_resolve, mock_client_cls, mock_get): @@ -81,10 +72,6 @@ def test_listdir_passes_correct_teamspace_id(mock_parse_lit_url, mock_resolve, m @mock.patch("lightning_sdk.api.filesystem_api.requests.get") @mock.patch("lightning_sdk.api.utils.LightningClient") -@mock.patch( - "lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers", - new=mock.MagicMock(return_value=FAKE_AUTH_HEADERS), -) @mock.patch("lightning_sdk.cli.utils.filesystem.resolve_teamspace_ref") @mock.patch("lightning_sdk.cli.utils.filesystem.parse_lit_url") def test_listdir_non_recursive( diff --git a/python/tests/core/filesystem/test_filesystem_teamspace_lookup.py b/python/tests/core/filesystem/test_filesystem_teamspace_lookup.py index 742ff8dc7..0d1d130a5 100644 --- a/python/tests/core/filesystem/test_filesystem_teamspace_lookup.py +++ b/python/tests/core/filesystem/test_filesystem_teamspace_lookup.py @@ -7,7 +7,6 @@ from lightning_sdk.filesystem import Filesystem from lightning_sdk.lightning_cloud.openapi.rest import ApiException -FAKE_AUTH_HEADERS = {"Authorization": "Bearer fake-token"} HOST = "https://lightning.ai" ORG_ID = "org-1" TEAMSPACE_ID = "ts-123" @@ -35,10 +34,6 @@ def client(): @pytest.fixture() def tree(): with ( - mock.patch( - "lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers", - return_value=FAKE_AUTH_HEADERS, - ), mock.patch("lightning_sdk.api.filesystem_api.requests.get") as get, ): get.return_value.status_code = 200 diff --git a/python/tests/core/filesystem/test_filesystem_walk.py b/python/tests/core/filesystem/test_filesystem_walk.py index 5d52b7bb2..c57b6554d 100644 --- a/python/tests/core/filesystem/test_filesystem_walk.py +++ b/python/tests/core/filesystem/test_filesystem_walk.py @@ -5,7 +5,6 @@ from lightning_sdk.filesystem import Filesystem TEAMSPACE_ID = "ts-123" -FAKE_AUTH_HEADERS = {"Authorization": "Bearer fake-token"} HOST = "https://lightning.ai" REMOTE_PATH = "data/test1.txt" LIT_URL = f"lit://my-org/my-teamspace/{REMOTE_PATH}" @@ -30,10 +29,6 @@ def fake_path_result(): @mock.patch("lightning_sdk.api.filesystem_api.requests.get") @mock.patch("lightning_sdk.api.utils.LightningClient") -@mock.patch( - "lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers", - new=mock.MagicMock(return_value=FAKE_AUTH_HEADERS), -) @mock.patch("lightning_sdk.cli.utils.filesystem.resolve_teamspace_ref") @mock.patch("lightning_sdk.cli.utils.filesystem.parse_lit_url") def test_walk_yields_os_walk_style_tuples( @@ -63,10 +58,6 @@ def test_walk_yields_os_walk_style_tuples( @mock.patch("lightning_sdk.api.filesystem_api.requests.get") @mock.patch("lightning_sdk.api.utils.LightningClient") -@mock.patch( - "lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers", - new=mock.MagicMock(return_value=FAKE_AUTH_HEADERS), -) @mock.patch("lightning_sdk.cli.utils.filesystem.resolve_teamspace_ref") @mock.patch("lightning_sdk.cli.utils.filesystem.parse_lit_url") def test_walk_flat_directory( @@ -93,10 +84,6 @@ def test_walk_flat_directory( @mock.patch("lightning_sdk.api.filesystem_api.requests.get") @mock.patch("lightning_sdk.api.utils.LightningClient") -@mock.patch( - "lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers", - new=mock.MagicMock(return_value=FAKE_AUTH_HEADERS), -) @mock.patch("lightning_sdk.cli.utils.filesystem.resolve_teamspace_ref") @mock.patch("lightning_sdk.cli.utils.filesystem.parse_lit_url") def test_walk_empty(mock_parse_lit_url, mock_resolve, mock_client_cls, mock_get, fake_teamspace, fake_path_result): @@ -114,10 +101,6 @@ def test_walk_empty(mock_parse_lit_url, mock_resolve, mock_client_cls, mock_get, @mock.patch("lightning_sdk.api.filesystem_api.requests.get") @mock.patch("lightning_sdk.api.utils.LightningClient") -@mock.patch( - "lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers", - new=mock.MagicMock(return_value=FAKE_AUTH_HEADERS), -) @mock.patch("lightning_sdk.cli.utils.filesystem.resolve_teamspace_ref") @mock.patch("lightning_sdk.cli.utils.filesystem.parse_lit_url") def test_walk_is_recursive( @@ -138,10 +121,6 @@ def test_walk_is_recursive( @mock.patch("lightning_sdk.api.filesystem_api.requests.get") @mock.patch("lightning_sdk.api.utils.LightningClient") -@mock.patch( - "lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers", - new=mock.MagicMock(return_value=FAKE_AUTH_HEADERS), -) @mock.patch("lightning_sdk.cli.utils.filesystem.resolve_teamspace_ref") @mock.patch("lightning_sdk.cli.utils.filesystem.parse_lit_url") def test_walk_passes_correct_teamspace_id(mock_parse_lit_url, mock_resolve, mock_client_cls, mock_get): From 70588e0a72b6765c8d23810448677fa7eb2fc1b4 Mon Sep 17 00:00:00 2001 From: fnwoke-lightning Date: Wed, 4 Nov 2026 19:48:16 -0500 Subject: [PATCH 06/27] feat(cli): route login and auth commands through the credential source The interim sess_ login stays the default. The new login runs only when LIGHTNING_CLI_OIDC_LOGIN is set, and the SDK no longer patches Auth. Co-authored-by: Cursor --- python/lightning_sdk/cli/auth/status.py | 95 ++-- python/lightning_sdk/cli/auth/token.py | 46 +- python/lightning_sdk/cli/entrypoint.py | 59 +-- python/lightning_sdk/cli/legacy/entrypoint.py | 30 +- python/lightning_sdk/cli/utils/auth.py | 477 ++++-------------- python/tests/cli/auth/test_status_token.py | 109 ---- python/tests/cli/test_entrypoint.py | 135 +---- python/tests/cli/test_login_commands.py | 348 +++++++++++++ python/tests/cli/utils/test_auth.py | 450 +---------------- 9 files changed, 601 insertions(+), 1148 deletions(-) delete mode 100644 python/tests/cli/auth/test_status_token.py create mode 100644 python/tests/cli/test_login_commands.py diff --git a/python/lightning_sdk/cli/auth/status.py b/python/lightning_sdk/cli/auth/status.py index c0f3a1698..0fe61b066 100644 --- a/python/lightning_sdk/cli/auth/status.py +++ b/python/lightning_sdk/cli/auth/status.py @@ -1,5 +1,6 @@ -"""Show the saved CLI login.""" +"""Show the credential requests send, and what a saved login reaches.""" +import os from datetime import datetime from typing import Any, Dict, Optional @@ -9,10 +10,11 @@ from lightning_sdk.api.utils import cached_lightning_client from lightning_sdk.cli.utils.logging import LightningCommand from lightning_sdk.cli.utils.richt_print import rich_to_str -from lightning_sdk.lightning_cloud import cli_login -from lightning_sdk.lightning_cloud.login import Auth +from lightning_sdk.lightning_cloud import cli_login, env +from lightning_sdk.lightning_cloud.credentials import Credential, CredentialKind, CredentialSource -_NOT_LOGGED_IN = "You are not logged in. Run `lightning login` first." +_NOT_LOGGED_IN = "You are not logged in. Run `lightning login`." +_UNKNOWN_REACH = "not known locally" def _stamp(epoch: float) -> str: @@ -47,15 +49,15 @@ def _teamspace_names(client: Any) -> Dict[str, str]: return {member.project_id: member.name for member in memberships if getattr(member, "project_id", None)} -def _grant_lines(grant: Dict[str, Any], orgs: Dict[str, str], teamspaces: Dict[str, str]) -> tuple[str, str]: +def _grant_lines(grant: Dict[str, Any], orgs: Dict[str, str], teamspaces: Dict[str, str]) -> tuple: scopes = grant.get("orgs") if scopes == "all": return "all", "all" if not isinstance(scopes, list) or not scopes: return "none", "none" - shown_orgs: list[str] = [] - shown_teamspaces: list[str] = [] + shown_orgs = [] + shown_teamspaces = [] for scope in scopes: if not isinstance(scope, dict): continue @@ -67,49 +69,56 @@ def _grant_lines(grant: Dict[str, Any], orgs: Dict[str, str], teamspaces: Dict[s return ", ".join(shown_orgs) or "none", ", ".join(shown_teamspaces) or "none" -def _legacy_session() -> bool: - auth = Auth() - return bool(auth.load() and auth.auth_token) +def _source(credential: Credential) -> str: + if credential.kind is CredentialKind.TOKEN_FILE: + return f"token file {os.environ.get(env.LIGHTNING_AUTH_TOKEN_FILE, '')}" + if credential.kind is CredentialKind.ENV_TOKEN: + return "LIGHTNING_AUTH_TOKEN" + if credential.kind is CredentialKind.LOGIN: + login = cli_login.current_login() + storage = login.storage if login is not None else "login" + return f"login ({storage})" + return "interim session login" + + +def _add_user(table: Table, user: Dict[str, Optional[str]]) -> None: + table.add_row("User", user.get("username") or user.get("email") or user.get("id") or "unknown") + if user.get("email"): + table.add_row("Email", user["email"]) + if user.get("id"): + table.add_row("User ID", user["id"]) @click.command("status", cls=LightningCommand) def status() -> None: - """Show who is logged in, the login's reach, its permissions and when it expires.""" - login = cli_login.current_login() + """Show the credential requests send, who it belongs to, and what a saved login reaches.""" + credential = CredentialSource().resolve() + if credential is None: + raise click.ClickException(_NOT_LOGGED_IN) + client = cached_lightning_client() user = _user(client) - - if login is None: - if not _legacy_session(): - raise click.ClickException(_NOT_LOGGED_IN) - table = Table(show_header=False, box=None) - table.add_column("Field", style="bold") - table.add_column("Value") - table.add_row("User", user.get("username") or user.get("email") or user.get("id") or "unknown") - table.add_row("Login", "session token saved by an older SDK") - table.add_row("Expiry", "managed by the server") - table.add_row("Hint", "Run `lightning login` to use the new login.") - click.echo(rich_to_str(table), color=True) - return - - if login.grant.get("orgs") == "all": - orgs, teamspaces = "all", "all" - else: - orgs, teamspaces = _grant_lines(login.grant, _org_names(client), _teamspace_names(client)) - permissions = ", ".join(login.grant.get("permissions") or []) or "none" - table = Table(show_header=False, box=None) table.add_column("Field", style="bold") table.add_column("Value") - table.add_row("User", user.get("username") or user.get("email") or user.get("id") or "unknown") - if user.get("email"): - table.add_row("Email", user["email"]) - if user.get("id"): - table.add_row("User ID", user["id"]) - table.add_row("Storage", login.storage) - table.add_row("Orgs", orgs) - table.add_row("Teamspaces", teamspaces) - table.add_row("Permissions", permissions) - table.add_row("Access expires", _stamp(login.access_expires_at)) - table.add_row("Login expires", _stamp(login.refresh_expires_at)) + _add_user(table, user) + table.add_row("Source", _source(credential)) + + login = cli_login.current_login() if credential.kind is CredentialKind.LOGIN else None + if login is None: + table.add_row("Orgs", _UNKNOWN_REACH) + table.add_row("Teamspaces", _UNKNOWN_REACH) + table.add_row("Permissions", _UNKNOWN_REACH) + table.add_row("Expiry", _UNKNOWN_REACH) + else: + if login.grant.get("orgs") == "all": + orgs, teamspaces = "all", "all" + else: + orgs, teamspaces = _grant_lines(login.grant, _org_names(client), _teamspace_names(client)) + permissions = ", ".join(login.grant.get("permissions") or []) or "none" + table.add_row("Orgs", orgs) + table.add_row("Teamspaces", teamspaces) + table.add_row("Permissions", permissions) + table.add_row("Access expires", _stamp(login.access_expires_at)) + table.add_row("Login expires", _stamp(login.refresh_expires_at)) click.echo(rich_to_str(table), color=True) diff --git a/python/lightning_sdk/cli/auth/token.py b/python/lightning_sdk/cli/auth/token.py index afa4d239e..c0c87fbc7 100644 --- a/python/lightning_sdk/cli/auth/token.py +++ b/python/lightning_sdk/cli/auth/token.py @@ -1,26 +1,38 @@ -"""Print an access token for curl and other callers.""" +"""Print the credential requests send, for curl and other callers.""" + +import os import rich_click as click from lightning_sdk.cli.utils.logging import LightningCommand -from lightning_sdk.lightning_cloud import cli_login -from lightning_sdk.lightning_cloud.login import Auth +from lightning_sdk.lightning_cloud import cli_login, env +from lightning_sdk.lightning_cloud.credentials import Credential, CredentialKind, CredentialSource + +_NOT_LOGGED_IN = "You are not logged in. Run `lightning login`." + + +def _note(credential: Credential) -> str: + if credential.kind is CredentialKind.TOKEN_FILE: + path = os.environ.get(env.LIGHTNING_AUTH_TOKEN_FILE, "") + origin = f"the token file {path}" + elif credential.kind is CredentialKind.ENV_TOKEN: + origin = "LIGHTNING_AUTH_TOKEN" + else: + origin = "the interim session login" + return f"This token is from {origin}. It is what requests send." @click.command("token", cls=LightningCommand) def token() -> None: - """Print an access token with at least 30 minutes left, refreshing it first when needed.""" - try: - click.echo(cli_login.token()) - return - except cli_login.LoginError as err: - failure = err - - try: - header = Auth().auth_header - except (ValueError, ConnectionError): - header = None - if header and header.startswith("Bearer "): - click.echo(header.removeprefix("Bearer ")) + """Print the credential requests send. A saved login has at least 30 minutes left.""" + credential = CredentialSource().resolve() + if credential is None: + raise click.ClickException(_NOT_LOGGED_IN) + if credential.kind is CredentialKind.LOGIN: + try: + click.echo(cli_login.token()) + except cli_login.LoginError as err: + raise click.ClickException(str(err)) from err return - raise click.ClickException(str(failure)) from failure + click.echo(_note(credential), err=True) + click.echo(credential.token) diff --git a/python/lightning_sdk/cli/entrypoint.py b/python/lightning_sdk/cli/entrypoint.py index 3cb4d8ee7..42d432f3a 100644 --- a/python/lightning_sdk/cli/entrypoint.py +++ b/python/lightning_sdk/cli/entrypoint.py @@ -6,7 +6,6 @@ import rich_click as click from lightning_sdk import __version__ -from lightning_sdk.api.studio_api import _cloud_url from lightning_sdk.cli import groups as cli_groups from lightning_sdk.cli.completion import completion @@ -45,15 +44,8 @@ from lightning_sdk.cli.lockdown import apply_lockdown from lightning_sdk.cli.ls import ls from lightning_sdk.cli.rm import rm -from lightning_sdk.cli.utils.auth import ( - browser_authentication, - login_with_pasted_code, - logout_session, - oidc_login_enabled, -) +from lightning_sdk.cli.utils.auth import browser_authentication, run_login, run_logout from lightning_sdk.cli.utils.logging import CommandLoggingGroup, logging_excepthook -from lightning_sdk.lightning_cloud.login import Auth -from lightning_sdk.utils.resolve import _get_authed_user, in_studio click.rich_click.COMMAND_GROUPS = { "lightning": [ @@ -104,53 +96,24 @@ def main_cli(ctx: click.Context) -> None: @click.option( "--no-browser", is_flag=True, - hidden=True, help="Sign in without a browser on this machine: approve on another device, then paste the code it shows.", ) -def login(no_browser: bool) -> None: +@click.option( + "--insecure-storage", + is_flag=True, + help="Keep the login in a file instead of the OS keychain.", +) +@click.option("--org", default=None, help="Comma-separated orgs this login should be limited to.") +@click.option("--project", default=None, help="Teamspace this login should be limited to.") +def login(no_browser: bool, insecure_storage: bool, org: str, project: str) -> None: """Sign in to Lightning AI.""" - if no_browser and not oidc_login_enabled(): - raise click.UsageError("`lightning login --no-browser` is not available yet.") - auth = Auth() - if auth.api_key or auth.auth_token: - click.echo('"lightning login" is not required within a Studio or when credentials are set in the environment') - return - - if in_studio(): - # this is unexpected, as we automatically auth within a Studio - raise RuntimeError("Unable to login within a Studio. Did you change your shell setup?") from None - - # always start a fresh session: a saved one may have expired - auth.clear() - try: - if no_browser: - login_with_pasted_code(auth) - else: - with browser_authentication(True): - auth.authenticate() - except ConnectionError: - raise RuntimeError(f"Unable to connect to {_cloud_url()}. Please check your internet connection.") from None - except ValueError: - # the browser flow ended without the callback saving a token (see the warning it logs) - raise RuntimeError( - "Login failed: the browser sign-in finished without returning a token. Try running `lightning login` again" - ) from None - - try: - click.echo(f'Successfully logged in as "{_get_authed_user().name}".') - except Exception: - # credentials are saved; only the lookup of the user's name failed - click.echo("Successfully logged in.") + run_login(no_browser=no_browser, insecure_storage=insecure_storage, orgs=org, project=project) @main_cli.command def logout() -> None: """Sign out of the current session.""" - try: - logout_session(Auth()) - except Exception as e: - raise RuntimeError(f"Logout failed: {e}") from e - click.echo("Successfully logged out.") + run_logout() # Add new command groups diff --git a/python/lightning_sdk/cli/legacy/entrypoint.py b/python/lightning_sdk/cli/legacy/entrypoint.py index f36108b72..c85a842a5 100644 --- a/python/lightning_sdk/cli/legacy/entrypoint.py +++ b/python/lightning_sdk/cli/legacy/entrypoint.py @@ -10,7 +10,6 @@ from rich.text import Text from lightning_sdk import __version__ -from lightning_sdk.api.studio_api import _cloud_url from lightning_sdk.cli.legacy.configure import configure from lightning_sdk.cli.legacy.connect import connect from lightning_sdk.cli.legacy.create import create @@ -27,10 +26,9 @@ from lightning_sdk.cli.legacy.stop import stop from lightning_sdk.cli.legacy.switch import switch from lightning_sdk.cli.legacy.upload import upload -from lightning_sdk.cli.utils.auth import browser_authentication +from lightning_sdk.cli.utils.auth import browser_authentication, run_login, run_logout from lightning_sdk.cli.utils.coloring import CustomHelpFormatter from lightning_sdk.constants import _LIGHTNING_DEBUG -from lightning_sdk.lightning_cloud.login import Auth def _notify_exception(exception_type: Type[BaseException], value: BaseException, tb: Optional[TracebackType]) -> None: @@ -71,23 +69,27 @@ def main_cli(ctx: click.Context) -> None: @main_cli.command -def login() -> None: +@click.option( + "--no-browser", + is_flag=True, + help="Sign in without a browser on this machine: approve on another device, then paste the code it shows.", +) +@click.option( + "--insecure-storage", + is_flag=True, + help="Keep the login in a file instead of the OS keychain.", +) +@click.option("--org", default=None, help="Comma-separated orgs this login should be limited to.") +@click.option("--project", default=None, help="Teamspace this login should be limited to.") +def login(no_browser: bool, insecure_storage: bool, org: str, project: str) -> None: """Login to Lightning AI Studios.""" - auth = Auth() - auth.clear() - - try: - with browser_authentication(True): - auth.authenticate() - except ConnectionError: - raise RuntimeError(f"Unable to connect to {_cloud_url()}. Please check your internet connection.") from None + run_login(no_browser=no_browser, insecure_storage=insecure_storage, orgs=org, project=project) @main_cli.command def logout() -> None: """Logout from Lightning AI Studios.""" - auth = Auth() - auth.clear() + run_logout() # additional commands diff --git a/python/lightning_sdk/cli/utils/auth.py b/python/lightning_sdk/cli/utils/auth.py index 5de4c68fd..b031d4f1e 100644 --- a/python/lightning_sdk/cli/utils/auth.py +++ b/python/lightning_sdk/cli/utils/auth.py @@ -1,384 +1,129 @@ -"""Non-interactive authentication helpers for CLI commands.""" +"""CLI authentication commands, on top of the vendored credential source.""" -import base64 -import hashlib -import hmac import json import os -import secrets -from contextlib import contextmanager, suppress -from contextvars import ContextVar -from pathlib import Path -from typing import Generator, Optional, Tuple -from urllib.parse import urlencode +from typing import Any, Dict, Optional, Sequence -import requests -import rich_click as click +import click -from lightning_sdk.lightning_cloud import env -from lightning_sdk.lightning_cloud.login import Auth, AuthServer, logger +from lightning_sdk.lightning_cloud import cli_login, env +from lightning_sdk.lightning_cloud.credentials import CredentialKind, LoginUnavailable, logout, start_login +from lightning_sdk.lightning_cloud.credentials import browser_login as browser_authentication +from lightning_sdk.lightning_cloud.login import Auth +from lightning_sdk.utils.resolve import _get_authed_user, in_studio -_BROWSER_AUTH_ALLOWED = ContextVar("lightning_browser_auth_allowed", default=True) -_ORIGINAL_AUTH_SAVE = Auth.save +__all__ = ["browser_authentication", "require_auth_header", "run_login", "run_logout"] -_CLIENT_ID = "lightning-cli" -_PASTE_ATTEMPTS = 3 +_REVOKE_FAILED = { + CredentialKind.LOGIN: ("The saved login could not be revoked on the server and may work until it expires."), + CredentialKind.SESSION: ("The saved session could not be revoked on the server and may work until it expires."), +} -# The authorization-code login needs GRID_AUTH_CLI_LOGIN_ENABLED on the control plane, which production -# doesn't set yet. Until it does, `lightning login` uses the session-token login, and this opts in to test it. -OIDC_LOGIN_ENV = "LIGHTNING_CLI_OIDC_LOGIN" - -def oidc_login_enabled() -> bool: - return env.truthy(os.getenv(OIDC_LOGIN_ENV)) - - -def _generate_state() -> str: - return base64.urlsafe_b64encode(os.urandom(16)).rstrip(b"=").decode("ascii") - - -def _origin() -> str: - """Platform origin without a trailing slash, so paths are not `//cli-login`.""" - return env.LIGHTNING_CLOUD_URL.rstrip("/") - - -def _pkce_pair() -> Tuple[str, str]: - """Return a PKCE `(verifier, S256 challenge)` pair.""" - verifier = secrets.token_urlsafe(64) - challenge = base64.urlsafe_b64encode(hashlib.sha256(verifier.encode("ascii")).digest()).rstrip(b"=") - return verifier, challenge.decode("ascii") - - -def _authorize_url(redirect_uri: str, state: str, challenge: str) -> str: - query = urlencode( - { - "response_type": "code", - "client_id": _CLIENT_ID, - "redirect_uri": redirect_uri, - "scope": "openid", - "state": state, - "code_challenge": challenge, - "code_challenge_method": "S256", - } - ) - return f"{_origin()}/oidc/authorize?{query}" - - -class _CodeRejectedError(Exception): - """The token endpoint refused the code; `str(e)` is its OAuth error.""" - - -def _exchange_code(code: str, verifier: str, redirect_uri: str) -> str: - """Trade an approved code and its PKCE verifier for a CLI session token.""" +def require_auth_header() -> str: + """The credential requests send, or a usage error that tells the user to log in.""" try: - response = requests.post( - f"{_origin()}/oidc/token", - data={ - "grant_type": "authorization_code", - "client_id": _CLIENT_ID, - "code": code, - "code_verifier": verifier, - "redirect_uri": redirect_uri, - }, - timeout=10, - ) - except requests.ConnectionError as e: - raise ConnectionError(str(e)) from e - body = _json_body(response) - if response.ok and body.get("access_token"): - return body["access_token"] - raise _CodeRejectedError(body.get("error") or f"HTTP {response.status_code}") - - -def _save_with_session_token( - self: Auth, - token: str = "", - user_id: str = "", - api_key: str = "", - auth_token: str = "", + header = Auth().auth_header + except (ValueError, ConnectionError) as err: + raise click.UsageError(str(err)) from err + if not header: + raise click.UsageError("No Lightning credentials are available. Run `lightning login` first.") + return header + + +def _env_set(name: str) -> bool: + return bool(os.environ.get(name, "").strip()) + + +def _org_list(orgs: Optional[str]) -> Optional[Sequence[str]]: + if not orgs: + return None + names = [part.strip() for part in orgs.split(",") if part.strip()] + return names or None + + +def _grant_reaches_further(grant: Dict[str, Any], orgs: Optional[Sequence[str]], project: Optional[str]) -> bool: + """Whether ``grant`` covers more than the orgs and project the user asked for.""" + scopes = grant.get("orgs") + asked = {org for org in orgs or [] if org} + if scopes == "all": + return bool(asked or project) + if not isinstance(scopes, list): + return False + for scope in scopes: + if not isinstance(scope, dict): + continue + org = str(scope.get("org", "")) + projects = scope.get("projects") + if asked and org not in asked: + return True + if not project: + continue + if not isinstance(projects, list) or not projects: + return True + named = [str(item) for item in projects] + if project not in named or len(named) > 1: + return True + return False + + +def run_login( + no_browser: bool = False, + insecure_storage: bool = False, + orgs: Optional[str] = None, + project: Optional[str] = None, ) -> None: - """Persist session tokens (`token`) as `auth_token` so the SDK/CLI sends `Bearer sess_...`.""" - if not auth_token and token and (not api_key or token.startswith("sess_")): - auth_token = token - if token.startswith("sess_"): - api_key = "" - _ORIGINAL_AUTH_SAVE( - self, - token=token, - user_id=user_id, - api_key=api_key, - auth_token=auth_token, - ) - with suppress(OSError): - os.chmod(self.secrets_file, 0o600) - - -def _login_with_authorization_code(self: AuthServer, auth: Auth) -> None: - """Sign in with an authorization code and PKCE, delivered to a callback on 127.0.0.1.""" - import webbrowser - - import uvicorn - from fastapi import FastAPI - from starlette.background import BackgroundTask - from starlette.responses import RedirectResponse - - from lightning_sdk.lightning_cloud.utils.network import find_free_network_port - - expected_state = _generate_state() - verifier, challenge = _pkce_pair() - app = FastAPI() - port = find_free_network_port() - redirect_uri = f"http://127.0.0.1:{port}/callback" - url = _authorize_url(redirect_uri, expected_state, challenge) - - try: - requests.head(url, timeout=10) - except requests.ConnectionError as e: - raise requests.ConnectionError( - f"No internet connection available. Please connect to a stable internet connection \n{e}" - ) from e - except requests.RequestException as e: - raise requests.RequestException( - f"An error occurred with the request. Please report this issue to Lightning Team \n{e}" - ) from e - - logger.info(f"login started for lightning.ai, opening {url}") - if not webbrowser.open(url): - if os.environ.get("LIGHTNING_DEPLOYMENT_ID"): - raise RuntimeError( - "Failed to authenticate to Lightning. Ensure that you have selected 'Include SDK credentials' in the " - "'Environment' section of the deployment settings." - ) - click.echo("No browser could be opened on this machine.") - login_with_pasted_code(auth) - return - - @app.get("/callback") - def save_token(code: str = "", state: str = "", error: str = "") -> RedirectResponse: - failure = "" - if not state or not hmac.compare_digest(state, expected_state): - failure = "state mismatch" - elif error or not code: - failure = error or "no code in the callback" - else: - try: - auth.save(auth_token=_exchange_code(code, verifier, redirect_uri)) - except (_CodeRejectedError, ConnectionError) as e: - failure = str(e) - - if failure: - logger.warning(f"Authentication Failed ({failure}). Run `lightning login` again, or update the CLI.") - target = f"{_origin()}/cli-login-failed" - else: - logger.info("Authentication Successful") - target = f"{_origin()}/cli-login-successful" - - return RedirectResponse( - url=target, - background=BackgroundTask(stop_server), + """Run the login the switch selects, without dropping the saved login first.""" + if _env_set(env.LIGHTNING_AUTH_TOKEN_FILE): + raise click.UsageError( + "LIGHTNING_AUTH_TOKEN_FILE is set, so a new login would not be what requests send. " + "Unset it and run `lightning login` again." ) - - def stop_server() -> None: - server.should_exit = True - - server = uvicorn.Server(config=uvicorn.Config(app, host="127.0.0.1", port=port, log_level="error")) - server.run() - - -def _login_with_session_token(self: AuthServer, auth: Auth) -> None: - """Sign in through the sign-in page, which hands a session token to a callback on localhost.""" - import webbrowser - - import uvicorn - from fastapi import FastAPI, Query - from starlette.background import BackgroundTask - from starlette.responses import RedirectResponse - - from lightning_sdk.lightning_cloud.utils.network import find_free_network_port - - expected_state = _generate_state() - app = FastAPI() - port = find_free_network_port() - base_url = self.get_auth_url(port) - redirect_uri = f"http://localhost:{port}/login-complete" - if urlencode({"redirectTo": redirect_uri}) in base_url: - stateful_redirect = f"{redirect_uri}?{urlencode({'state': expected_state})}" - url = base_url.replace(urlencode({"redirectTo": redirect_uri}), urlencode({"redirectTo": stateful_redirect}), 1) - else: - url = base_url - - try: - requests.head(url, timeout=10) - except requests.ConnectionError as e: - raise requests.ConnectionError( - f"No internet connection available. Please connect to a stable internet connection \n{e}" - ) from e - except requests.RequestException as e: - raise requests.RequestException( - f"An error occurred with the request. Please report this issue to Lightning Team \n{e}" - ) from e - - logger.info(f"login started for lightning.ai, opening {url}") - if not webbrowser.open(url): - deployment_id = os.environ.get("LIGHTNING_DEPLOYMENT_ID") - if deployment_id: - raise RuntimeError( - "Failed to authenticate to Lightning. Ensure that you have selected 'Include SDK credentials' in the " - "'Environment' section of the deployment settings." - ) - raise RuntimeError( - "Failed to authenticate to Lightning. When running without access to a browser, " - "'LIGHTNING_API_KEY' should be exported." + if _env_set("LIGHTNING_AUTH_TOKEN"): + raise click.UsageError( + "LIGHTNING_AUTH_TOKEN is set, so a new login would not be what requests send. " + "Unset it and run `lightning login` again." ) - - @app.get("/login-complete") - async def save_token( - token: str = "", - key: str = "", - state: str = "", - user_id: str = Query("", alias="userID"), - ) -> RedirectResponse: - state_valid = bool(state) and hmac.compare_digest(state, expected_state) - if token and state_valid: - auth.save(token=token, user_id=user_id, api_key=key) - logger.info("Authentication Successful") - target = f"{env.LIGHTNING_CLOUD_URL}/cli-login-successful" - else: - reason = "no token in the callback" if not token else "state mismatch" - logger.warning( - f"Authentication Failed ({reason}). If the browser was already signed in, the sign-in page " - "could not mint a CLI token. Run `lightning login` again, or update the CLI." - ) - target = f"{env.LIGHTNING_CLOUD_URL}/cli-login-failed" - - return RedirectResponse( - url=target, - background=BackgroundTask(stop_server), + if _env_set("LIGHTNING_API_KEY"): + click.echo( + "LIGHTNING_API_KEY is ignored until service accounts. This login continues without it.", + err=True, ) + if in_studio(): + raise RuntimeError("Unable to login within a Studio. Did you change your shell setup?") from None - def stop_server() -> None: - server.should_exit = True - - server = uvicorn.Server(config=uvicorn.Config(app, port=port, log_level="error")) - server.run() - - -def _login_with_browser(self: AuthServer, auth: Auth) -> None: - if oidc_login_enabled(): - _login_with_authorization_code(self, auth) - else: - _login_with_session_token(self, auth) - - -Auth.save = _save_with_session_token -AuthServer.login_with_browser = _login_with_browser - - -def login_with_pasted_code(auth: Auth) -> None: - """Sign in on a machine without a browser: approve on any device, then paste the code shown there.""" - redirect_uri = f"{_origin()}/cli-login" - verifier, challenge = _pkce_pair() - click.echo( - "To sign in, open this address in a browser on any device and approve the request:\n\n" - f" {_authorize_url(redirect_uri, _generate_state(), challenge)}\n\n" - "The page then shows a code. Only paste a code from a page you opened this way; " - "never give it to anyone, as it gives full access to your Lightning account." - ) - for attempt in range(_PASTE_ATTEMPTS): - code = click.prompt("Code").strip() - try: - auth.save(auth_token=_exchange_code(code, verifier, redirect_uri)) - return - except _CodeRejectedError as e: - if str(e) != "invalid_grant": - raise RuntimeError(f"Login failed: {e}.") from None - if attempt < _PASTE_ATTEMPTS - 1: - click.echo("That code was not accepted. It works once and expires after a minute. Try again.") - raise RuntimeError("Login failed: the code was not accepted. Run `lightning login --no-browser` again.") - - -def _json_body(response: requests.Response) -> dict: + asked_orgs = _org_list(orgs) + asked_project = project.strip() if isinstance(project, str) and project.strip() else None + if insecure_storage and cli_login.oidc_login_enabled(): + click.echo(f"This login is kept in a file at {env.LIGHTNING_LOGIN_PATH}.", err=True) try: - body = response.json() - except ValueError: - return {} - return body if isinstance(body, dict) else {} - - -def logout_session(auth: Auth) -> None: - """Revoke the active CLI session token on the server and clear local credentials.""" - token = getattr(auth, "auth_token", None) - if not isinstance(token, str) or not token: - secrets_file = getattr(auth, "secrets_file", None) - if isinstance(secrets_file, Path) and secrets_file.exists(): - try: - with secrets_file.open() as creds: - data = json.load(creds) - loaded = data.get("auth_token") - if isinstance(loaded, str) and loaded: - token = loaded - except Exception: - token = None - if isinstance(token, str) and token: - with suppress(requests.RequestException): - requests.post( - f"{_origin()}/v1/auth/logout", - headers={"Authorization": f"Bearer {token}"}, - timeout=10, - ) - auth.clear() - - -@contextmanager -def browser_authentication(allowed: bool) -> Generator[None, None, None]: - """Temporarily control whether missing credentials may start browser authentication.""" - token = _BROWSER_AUTH_ALLOWED.set(allowed) - authenticate = Auth.authenticate - - def guarded_authenticate(auth: Auth) -> Optional[str]: - if _BROWSER_AUTH_ALLOWED.get(): - return authenticate(auth) - - has_own_run_server = "_run_server" in auth.__dict__ - run_server = auth.__dict__.get("_run_server") - - def blocked_run_server() -> None: - raise ValueError("No Lightning credentials are available. Run `lightning login` first.") - - auth.__dict__["_run_server"] = blocked_run_server - try: - return authenticate(auth) - finally: - if has_own_run_server: - auth.__dict__["_run_server"] = run_server - else: - auth.__dict__.pop("_run_server", None) - - Auth.authenticate = guarded_authenticate - try: - yield - finally: - Auth.authenticate = authenticate - _BROWSER_AUTH_ALLOWED.reset(token) - + done = start_login( + no_browser=no_browser, + insecure_storage=insecure_storage, + orgs=asked_orgs, + project=asked_project, + ) + except LoginUnavailable as err: + raise click.UsageError(str(err)) from err -def require_auth_header() -> str: - """Return available credentials without falling back to browser authentication.""" - auth = Auth() + wider = done is not None and (asked_orgs or asked_project) + if wider and _grant_reaches_further(done.grant, asked_orgs, asked_project): + grant = json.dumps(done.grant, sort_keys=True) + click.echo(f"This login reaches more than you asked for. The grant is {grant}.") try: - header = auth.auth_header - except ValueError: - header = None - - if header: - return header - - if auth.load(): - try: - header = auth.auth_header - except ValueError: - header = None - if header: - return header - - raise click.UsageError("No Lightning credentials are available. Run `lightning login` first.") + click.echo(f'Successfully logged in as "{_get_authed_user().name}".') + except Exception: + click.echo("Successfully logged in.") + + +def run_logout() -> None: + """Forget every saved login on this machine, and say which server revoke failed.""" + failed = logout() + click.echo("Logged out on this machine.") + for kind in failed: + click.echo( + _REVOKE_FAILED.get( + kind, + f"The {kind.value} credential could not be revoked on the server and may work until it expires.", + ) + ) diff --git a/python/tests/cli/auth/test_status_token.py b/python/tests/cli/auth/test_status_token.py deleted file mode 100644 index ae2dd477a..000000000 --- a/python/tests/cli/auth/test_status_token.py +++ /dev/null @@ -1,109 +0,0 @@ -import time -from types import SimpleNamespace -from unittest.mock import MagicMock, PropertyMock - -from click.testing import CliRunner - -from lightning_sdk.cli.entrypoint import main_cli -from lightning_sdk.lightning_cloud.cli_login import LoginError, LoginStatus - - -def _status(grant=None) -> LoginStatus: - now = time.time() - return LoginStatus( - url="https://platform.lightning.ai", - storage="keychain", - access_expires_at=now + 3600, - refresh_expires_at=now + 90 * 24 * 3600, - grant=grant or {"orgs": "all", "permissions": ["read", "write"]}, - ) - - -def _client() -> MagicMock: - client = MagicMock() - client.auth_service_get_user.return_value = SimpleNamespace(id="user-1", username="ada", email="ada@example.com") - client.organizations_service_list_organizations.return_value = SimpleNamespace( - organizations=[SimpleNamespace(id="org-1", name="acme"), SimpleNamespace(id="org-2", name="beta")] - ) - client.projects_service_list_memberships.return_value = SimpleNamespace( - memberships=[SimpleNamespace(project_id="proj-1", name="research")] - ) - return client - - -def test_auth_status_shows_the_login_without_a_token(monkeypatch) -> None: - monkeypatch.setattr("lightning_sdk.cli.auth.status.cli_login.current_login", lambda: _status()) - monkeypatch.setattr("lightning_sdk.cli.auth.status.cached_lightning_client", _client) - - result = CliRunner().invoke(main_cli, ["auth", "status"]) - - assert result.exit_code == 0, result.output - assert "ada" in result.output - assert "ada@example.com" in result.output - assert "keychain" in result.output - assert "all" in result.output - assert "read, write" in result.output - assert "lit_ua_" not in result.output - - -def test_auth_status_resolves_a_narrowed_grant(monkeypatch) -> None: - grant = {"orgs": [{"org": "org-1", "projects": ["proj-1"]}], "permissions": ["read"]} - monkeypatch.setattr("lightning_sdk.cli.auth.status.cli_login.current_login", lambda: _status(grant)) - monkeypatch.setattr("lightning_sdk.cli.auth.status.cached_lightning_client", _client) - - result = CliRunner().invoke(main_cli, ["auth", "status"]) - - assert result.exit_code == 0, result.output - assert "acme" in result.output - assert "research" in result.output - assert "read" in result.output - - -def test_auth_status_instructs_login_without_a_login(monkeypatch) -> None: - monkeypatch.setattr("lightning_sdk.cli.auth.status.cli_login.current_login", lambda: None) - auth = MagicMock() - auth.load.return_value = False - monkeypatch.setattr("lightning_sdk.cli.auth.status.Auth", lambda: auth) - - result = CliRunner().invoke(main_cli, ["auth", "status"]) - - assert result.exit_code == 1 - assert "Run `lightning login` first" in result.output - - -def test_auth_token_prints_only_the_token(monkeypatch) -> None: - monkeypatch.setattr("lightning_sdk.cli.auth.token.cli_login.token", lambda: "lit_ua_access") - - result = CliRunner().invoke(main_cli, ["auth", "token"]) - - assert result.exit_code == 0, result.output - assert result.output == "lit_ua_access\n" - - -def test_auth_token_falls_back_to_a_saved_session(monkeypatch) -> None: - monkeypatch.setattr( - "lightning_sdk.cli.auth.token.cli_login.token", - MagicMock(side_effect=LoginError("You are not logged in. Run `lightning login`.")), - ) - auth = MagicMock(auth_header="Bearer sess_saved") - monkeypatch.setattr("lightning_sdk.cli.auth.token.Auth", lambda: auth) - - result = CliRunner().invoke(main_cli, ["auth", "token"]) - - assert result.exit_code == 0, result.output - assert result.output == "sess_saved\n" - - -def test_auth_token_reports_a_missing_login(monkeypatch) -> None: - monkeypatch.setattr( - "lightning_sdk.cli.auth.token.cli_login.token", - MagicMock(side_effect=LoginError("You are not logged in. Run `lightning login`.")), - ) - auth = MagicMock() - type(auth).auth_header = PropertyMock(side_effect=ValueError("missing")) - monkeypatch.setattr("lightning_sdk.cli.auth.token.Auth", lambda: auth) - - result = CliRunner().invoke(main_cli, ["auth", "token"]) - - assert result.exit_code == 1 - assert "lightning login" in result.output diff --git a/python/tests/cli/test_entrypoint.py b/python/tests/cli/test_entrypoint.py index ba8e6c302..56d166399 100644 --- a/python/tests/cli/test_entrypoint.py +++ b/python/tests/cli/test_entrypoint.py @@ -1,7 +1,7 @@ import os import re import subprocess -from unittest.mock import MagicMock, patch +from unittest.mock import patch import pytest import rich_click as click @@ -9,8 +9,8 @@ from lightning_sdk.cli.entrypoint import login, main_cli from lightning_sdk.cli.sandbox.commands import run_sandbox_command -from lightning_sdk.cli.utils.auth import OIDC_LOGIN_ENV from lightning_sdk.cli.vm.ssh import ssh_vm +from lightning_sdk.lightning_cloud.cli_login import OIDC_LOGIN_ENV from lightning_sdk.lightning_cloud.login import Auth from tests.cli.help import assert_help_contains, command_text, mock_command_logging, run_cli @@ -47,100 +47,27 @@ def test_resource_commands_never_start_browser_authentication(monkeypatch, tmp_p assert "Run `lightning login`" in str(result.exception) -def test_explicit_login_can_start_browser_authentication(monkeypatch, tmp_path) -> None: - _clear_auth(monkeypatch, tmp_path) - browser_auth_started = False - - def _complete_browser_auth(auth: Auth) -> None: - nonlocal browser_auth_started - browser_auth_started = True - auth.save(auth_token="sess_user-1.secret") - - monkeypatch.setattr(Auth, "_run_server", _complete_browser_auth) - - result = CliRunner().invoke(main_cli, ["login"]) - - assert result.exit_code == 0, result.output - assert browser_auth_started - - -def test_login_reports_success_and_failure(monkeypatch, tmp_path) -> None: - _clear_auth(monkeypatch, tmp_path) - user = MagicMock() - user.name = "Test User" - monkeypatch.setattr("lightning_sdk.cli.entrypoint._get_authed_user", lambda: user) - - def _save_token(auth: Auth) -> None: - auth.save(token="sess_user-1.secret", user_id="user-1") - - monkeypatch.setattr(Auth, "_run_server", _save_token) - ok = CliRunner().invoke(main_cli, ["login"]) - assert ok.exit_code == 0, ok.output - assert 'Successfully logged in as "Test User".' in ok.output - - # the browser callback finished without saving a token - _clear_auth(monkeypatch, tmp_path / "second") - monkeypatch.setattr(Auth, "_run_server", lambda _auth: None) - failed = CliRunner().invoke(main_cli, ["login"]) - assert failed.exit_code == 1 - assert "Login failed" in str(failed.exception) - - -def test_login_no_browser_starts_a_fresh_paste_login_instead_of_the_browser(monkeypatch, tmp_path) -> None: - _clear_auth(monkeypatch, tmp_path) - monkeypatch.setenv(OIDC_LOGIN_ENV, "1") - Auth().save(auth_token="sess_expired") - user = MagicMock() - user.name = "Test User" - monkeypatch.setattr("lightning_sdk.cli.entrypoint._get_authed_user", lambda: user) - monkeypatch.setattr( - Auth, - "_run_server", - lambda _self: (_ for _ in ()).throw(AssertionError("browser auth reached")), - ) - - def _approve(auth: Auth) -> None: - assert not Auth.secrets_file.exists() - auth.save(auth_token="sess_new") - - monkeypatch.setattr("lightning_sdk.cli.entrypoint.login_with_pasted_code", _approve) - - result = CliRunner().invoke(main_cli, ["login", "--no-browser"]) - - assert result.exit_code == 0, result.output - assert 'Successfully logged in as "Test User".' in result.output - auth = Auth() - assert auth.load() - assert auth.auth_token == "sess_new" - - def test_login_no_browser_is_refused_without_the_opt_in(monkeypatch, tmp_path) -> None: _clear_auth(monkeypatch, tmp_path) Auth().save(auth_token="sess_kept") - monkeypatch.setattr( - "lightning_sdk.cli.entrypoint.login_with_pasted_code", - lambda _auth: (_ for _ in ()).throw(AssertionError("paste login reached")), - ) result = CliRunner().invoke(main_cli, ["login", "--no-browser"]) assert result.exit_code == 2 - assert "not available yet" in result.output + assert "LIGHTNING_CLI_OIDC_LOGIN" in result.output auth = Auth() assert auth.load() assert auth.auth_token == "sess_kept" -def test_logout_reports_success_and_failure(monkeypatch, tmp_path) -> None: +def test_logout_forgets_the_local_login(monkeypatch, tmp_path) -> None: _clear_auth(monkeypatch, tmp_path) - ok = CliRunner().invoke(main_cli, ["logout"]) - assert ok.exit_code == 0, ok.output - assert "Successfully logged out." in ok.output - monkeypatch.setattr("lightning_sdk.cli.entrypoint.logout_session", MagicMock(side_effect=OSError("denied"))) - failed = CliRunner().invoke(main_cli, ["logout"]) - assert failed.exit_code == 1 - assert "Logout failed: denied" in str(failed.exception) + result = CliRunner().invoke(main_cli, ["logout"]) + + assert result.exit_code == 0, result.output + assert "Logged out on this machine." in result.output + assert "Logout failed" not in result.output def test_cli_browser_auth_policy_is_reset_after_command(monkeypatch, tmp_path) -> None: @@ -317,7 +244,7 @@ def test_login_help(): result_text = assert_help_contains("lightning login --help", "Usage: lightning login [OPTIONS]") assert "Sign in to Lightning AI." in result_text - assert "--no-browser" not in result_text + assert "--no-browser" in result_text @mock_command_logging @@ -327,50 +254,10 @@ def test_logout_help(): assert "Sign out of the current session." in result_text -def test_login_replaces_existing_session(monkeypatch, tmp_path): - _clear_auth(monkeypatch, tmp_path) - Auth().save(auth_token="sess_expired") - - def _complete_browser_auth(auth: Auth) -> None: - assert not Auth.secrets_file.exists() - auth.save(auth_token="sess_new") - - monkeypatch.setattr(Auth, "_run_server", _complete_browser_auth) - - result = CliRunner().invoke(main_cli, ["login"]) - - assert result.exit_code == 0, result.output - auth = Auth() - assert auth.load() - assert auth.auth_token == "sess_new" - - -@pytest.mark.parametrize("env_var", ["LIGHTNING_AUTH_TOKEN", "LIGHTNING_API_KEY"]) -def test_login_not_required_with_credentials_in_env(monkeypatch, tmp_path, env_var): - _clear_auth(monkeypatch, tmp_path) - monkeypatch.setenv(env_var, "from-env") - Auth().save(auth_token="sess_saved") - - def _fail_browser_auth(auth: Auth) -> None: - raise AssertionError("browser authentication should not start") - - monkeypatch.setattr(Auth, "_run_server", _fail_browser_auth) - - result = CliRunner().invoke(main_cli, ["login"]) - - assert result.exit_code == 0, result.output - assert "credentials are set in the environment" in result.output - assert Auth.secrets_file.exists() - - @patch.dict(os.environ, {"LIGHTNING_INTERACTIVE": "true", "LIGHTNING_CLOUD_SPACE_ID": "cloud-space-id"}, clear=True) @mock_command_logging -def test_cannot_login_within_studio_without_credentials(monkeypatch): - mock_auth_instance = MagicMock(api_key=None, auth_token=None) - monkeypatch.setattr("lightning_sdk.cli.entrypoint.Auth", MagicMock(return_value=mock_auth_instance)) - +def test_cannot_login_within_studio_without_credentials(): result = CliRunner().invoke(login) assert result.exit_code == 1 assert "Unable to login within a Studio. Did you change your shell setup?" in str(result.exception) - mock_auth_instance.authenticate.assert_not_called() diff --git a/python/tests/cli/test_login_commands.py b/python/tests/cli/test_login_commands.py new file mode 100644 index 000000000..fdab1c44e --- /dev/null +++ b/python/tests/cli/test_login_commands.py @@ -0,0 +1,348 @@ +"""Login, logout, auth status and auth token, against the fake platform.""" + +import getpass +import json +import sys +import threading +import webbrowser +from types import SimpleNamespace + +import keyring +import pytest +from click.testing import CliRunner + +from lightning_sdk.api.utils import cached_lightning_client +from lightning_sdk.cli.entrypoint import main_cli +from lightning_sdk.cli.legacy.entrypoint import main_cli as legacy_cli +from lightning_sdk.lightning_cloud import env +from lightning_sdk.lightning_cloud.credentials import CredentialSource, logout +from lightning_sdk.lightning_cloud.login import Auth +from tests.lightning_cloud.utils.fake_platform import FakeBrowser, FakePlatform, OtherDevice +from tests.lightning_cloud.utils.keyrings import MemoryKeyring + +_REAL_THREAD = threading.Thread +_NOT_LOGGED_IN = "You are not logged in. Run `lightning login`." + + +def _text(result) -> str: + return result.output + (getattr(result, "stderr", None) or "") + + +@pytest.fixture() +def isolated(monkeypatch, tmp_path): + """A private credential store, no ambient token, and real threads for the fake servers.""" + # Assigning, not monkeypatch: restoring the suite's Thread mock after the test + # makes the timeout timer fail its Thread.__init__ check. + threading.Thread = _REAL_THREAD + for name in ( + "LIGHTNING_AUTH_TOKEN", + "LIGHTNING_AUTH_TOKEN_FILE", + "LIGHTNING_API_KEY", + "LIGHTNING_USER_ID", + "LIGHTNING_CLI_OIDC_LOGIN", + ): + monkeypatch.delenv(name, raising=False) + credentials = tmp_path / "credentials.json" + monkeypatch.setattr(Auth, "secrets_file", credentials) + monkeypatch.setenv("LIGHTNING_CREDENTIAL_PATH", str(credentials)) + login_path = tmp_path / "login.json" + monkeypatch.setattr(env, "LIGHTNING_LOGIN_PATH", str(login_path)) + monkeypatch.setenv("LIGHTNING_LOGIN_PATH", str(login_path)) + monkeypatch.setattr("lightning_sdk.cli.utils.auth._get_authed_user", lambda: SimpleNamespace(name="ada")) + cached_lightning_client.cache_clear() + previous = keyring.get_keyring() + memory = MemoryKeyring() + keyring.set_keyring(memory) + try: + yield tmp_path + finally: + keyring.set_keyring(previous) + cached_lightning_client.cache_clear() + + +@pytest.fixture() +def platform(isolated, monkeypatch): + fake = FakePlatform() + monkeypatch.setattr(env, "LIGHTNING_CLOUD_URL", fake.url) + monkeypatch.setenv("LIGHTNING_CLOUD_URL", fake.url) + cached_lightning_client.cache_clear() + yield fake + fake.stop() + + +@pytest.fixture() +def browser(platform, monkeypatch): + fake = FakeBrowser(platform) + monkeypatch.setattr(webbrowser, "open", fake.open) + yield fake + fake.finish() + + +def _invoke(cli, args): + return CliRunner().invoke(cli, args) + + +def _saved_session() -> str: + data = json.loads(Auth.secrets_file.read_text()) + return data["auth_token"] + + +@pytest.mark.parametrize("cli", [main_cli, legacy_cli]) +def test_login_uses_the_interim_flow_when_the_switch_is_unset(platform, browser, cli) -> None: + result = _invoke(cli, ["login"]) + + assert result.exit_code == 0, _text(result) + assert _saved_session().startswith("sess_") + assert platform.sign_ins + assert not platform.authorizations + assert 'Successfully logged in as "ada".' in result.output + + +@pytest.mark.parametrize("cli", [main_cli, legacy_cli]) +@pytest.mark.parametrize( + "args", + [["--no-browser"], ["--insecure-storage"], ["--org", "acme"], ["--project", "research"]], +) +def test_interim_login_refuses_flags_the_new_login_needs(isolated, cli, args) -> None: + result = _invoke(cli, ["login", *args]) + + assert result.exit_code == 2, _text(result) + assert "LIGHTNING_CLI_OIDC_LOGIN" in _text(result) + assert not Auth.secrets_file.exists() + + +@pytest.mark.parametrize("cli", [main_cli, legacy_cli]) +def test_login_refuses_a_token_file_and_an_environment_token(isolated, monkeypatch, cli) -> None: + token_file = isolated / "token" + token_file.write_text("from-file") + monkeypatch.setenv("LIGHTNING_AUTH_TOKEN_FILE", str(token_file)) + refused = _invoke(cli, ["login"]) + assert refused.exit_code == 2 + assert "LIGHTNING_AUTH_TOKEN_FILE" in _text(refused) + + monkeypatch.delenv("LIGHTNING_AUTH_TOKEN_FILE") + monkeypatch.setenv("LIGHTNING_AUTH_TOKEN", "from-env") + refused = _invoke(cli, ["login"]) + assert refused.exit_code == 2 + assert "LIGHTNING_AUTH_TOKEN" in _text(refused) + assert not Auth.secrets_file.exists() + + +def test_login_warns_that_an_api_key_is_ignored_and_continues(platform, browser, monkeypatch) -> None: + monkeypatch.setenv("LIGHTNING_API_KEY", "lit_sa_ignored") + + result = _invoke(main_cli, ["login"]) + + assert result.exit_code == 0, _text(result) + assert "LIGHTNING_API_KEY is ignored until service accounts" in _text(result) + assert _saved_session().startswith("sess_") + assert "lit_sa_ignored" not in _text(result) + + +def test_login_keeps_the_previous_session_until_the_new_one_is_saved(platform, browser) -> None: + previous = platform.mint_session("user-1") + Auth().save(auth_token=previous) + + result = _invoke(main_cli, ["login"]) + + assert result.exit_code == 0, _text(result) + assert previous in platform.logouts() + assert _saved_session() != previous + assert _saved_session().startswith("sess_") + + +@pytest.mark.parametrize("cli", [main_cli, legacy_cli]) +def test_login_with_the_switch_uses_the_new_login(platform, browser, monkeypatch, cli) -> None: + monkeypatch.setenv("LIGHTNING_CLI_OIDC_LOGIN", "1") + + result = _invoke(cli, ["login"]) + + assert result.exit_code == 0, _text(result) + assert platform.authorizations + assert CredentialSource().resolve().token.startswith("lit_ua_") + assert not platform.sign_ins + + +def test_insecure_storage_warns_and_names_the_login_file(platform, browser, monkeypatch) -> None: + monkeypatch.setenv("LIGHTNING_CLI_OIDC_LOGIN", "1") + + result = _invoke(main_cli, ["login", "--insecure-storage"]) + + assert result.exit_code == 0, _text(result) + assert env.LIGHTNING_LOGIN_PATH in _text(result) + assert "file" in _text(result) + assert CredentialSource().resolve().kind.value == "login" + + +@pytest.mark.parametrize("args", [["--org", "acme"], ["--no-browser", "--org", "acme"]]) +def test_login_names_a_grant_that_reaches_further_than_the_hints(platform, browser, monkeypatch, args) -> None: + monkeypatch.setenv("LIGHTNING_CLI_OIDC_LOGIN", "1") + if "--no-browser" in args: + + def read_terminal(): + sys.stderr.flush() + captured = getattr(sys.stderr, "buffer", sys.stderr) + raw = captured.getvalue() if hasattr(captured, "getvalue") else "" + return raw.decode() if isinstance(raw, bytes) else raw + + device = OtherDevice(platform, read_terminal, ["CODE"]) + monkeypatch.setattr(getpass, "getpass", device.getpass) + + result = _invoke(main_cli, ["login", *args]) + + assert result.exit_code == 0, _text(result) + assert "reaches more than you asked for" in _text(result) + assert '"orgs": "all"' in _text(result) + + +def test_login_stays_quiet_when_the_grant_matches_the_hints(platform, browser, monkeypatch) -> None: + monkeypatch.setenv("LIGHTNING_CLI_OIDC_LOGIN", "1") + original = platform._issue + + def matching(family): + body = original(family) + body["lightning_grant"] = {"orgs": [{"org": "acme", "projects": ["research"]}], "permissions": ["read"]} + return body + + platform._issue = matching + + result = _invoke(main_cli, ["login", "--org", "acme", "--project", "research"]) + + assert result.exit_code == 0, _text(result) + assert "reaches more than you asked for" not in _text(result) + + +def test_a_studio_without_credentials_still_refuses_to_log_in(monkeypatch) -> None: + monkeypatch.delenv("LIGHTNING_AUTH_TOKEN", raising=False) + monkeypatch.delenv("LIGHTNING_API_KEY", raising=False) + monkeypatch.delenv("LIGHTNING_AUTH_TOKEN_FILE", raising=False) + monkeypatch.setenv("LIGHTNING_INTERACTIVE", "true") + monkeypatch.setenv("LIGHTNING_CLOUD_SPACE_ID", "studio-1") + + result = _invoke(main_cli, ["login"]) + + assert result.exit_code != 0 + assert "Unable to login within a Studio" in str(result.exception) + + +@pytest.mark.parametrize("cli", [main_cli, legacy_cli]) +def test_logout_revokes_both_kinds_and_names_the_one_the_server_refused(platform, browser, monkeypatch, cli) -> None: + monkeypatch.setenv("LIGHTNING_CLI_OIDC_LOGIN", "1") + assert _invoke(cli, ["login"]).exit_code == 0 + session = platform.mint_session("user-1") + Auth().save(auth_token=session) + platform.revoke = lambda _form: (500, {"error": "unavailable"}) + + result = _invoke(cli, ["logout"]) + + text = _text(result) + assert result.exit_code == 0, text + assert "Logged out on this machine." in text + assert "could not be revoked on the server" in text + assert "may work until it expires" in text + assert "Logout failed" not in text + assert not platform.session_active(session) + assert CredentialSource().resolve() is None + + +def test_auth_status_names_each_source_and_never_a_token(platform, browser, monkeypatch, tmp_path) -> None: + monkeypatch.setenv("LIGHTNING_CLI_OIDC_LOGIN", "1") + assert _invoke(main_cli, ["login"]).exit_code == 0 + login_token = CredentialSource().resolve().token + + logged_in = _invoke(main_cli, ["auth", "status"]) + assert logged_in.exit_code == 0, _text(logged_in) + assert "login (keychain)" in _text(logged_in) + assert "all" in _text(logged_in) + assert login_token not in _text(logged_in) + assert f"Bearer {login_token}" in platform.authorizations_sent("/v1/auth/user") + + logout() + token_path = tmp_path / "workload" + token_path.write_text("file-token") + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("LIGHTNING_AUTH_TOKEN_FILE", "workload") + from_file = _invoke(main_cli, ["auth", "status"]) + assert "token file workload" in _text(from_file) + assert "not known locally" in _text(from_file) + assert "file-token" not in _text(from_file) + assert "not logged in" not in _text(from_file).lower() + + monkeypatch.delenv("LIGHTNING_AUTH_TOKEN_FILE") + monkeypatch.setenv("LIGHTNING_AUTH_TOKEN", "env-token") + from_env = _invoke(main_cli, ["auth", "status"]) + assert "LIGHTNING_AUTH_TOKEN" in _text(from_env) + assert "not known locally" in _text(from_env) + assert "env-token" not in _text(from_env) + assert "not logged in" not in _text(from_env).lower() + + monkeypatch.delenv("LIGHTNING_AUTH_TOKEN") + Auth().save(auth_token="sess_saved") + from_session = _invoke(main_cli, ["auth", "status"]) + assert "interim session login" in _text(from_session) + assert "not known locally" in _text(from_session) + assert "sess_saved" not in _text(from_session) + + Auth.secrets_file.unlink() + missing = _invoke(main_cli, ["auth", "status"]) + assert missing.exit_code != 0 + assert _NOT_LOGGED_IN in _text(missing) + + +def test_auth_token_prints_what_requests_send_and_refuses_a_short_login(platform, monkeypatch) -> None: + monkeypatch.setenv("LIGHTNING_AUTH_TOKEN", "env-token") + printed = _invoke(main_cli, ["auth", "token"]) + assert printed.exit_code == 0, _text(printed) + assert printed.stdout == "env-token\n" + assert "LIGHTNING_AUTH_TOKEN" in printed.stderr + assert "what requests send" in printed.stderr + + monkeypatch.delenv("LIGHTNING_AUTH_TOKEN") + monkeypatch.setenv("LIGHTNING_API_KEY", "lit_sa_secret") + refused = _invoke(main_cli, ["auth", "token"]) + assert refused.exit_code != 0 + assert "lit_sa_secret" not in _text(refused) + assert _NOT_LOGGED_IN in _text(refused) + + monkeypatch.delenv("LIGHTNING_API_KEY") + Auth().save(auth_token="sess_saved") + session = _invoke(main_cli, ["auth", "token"]) + assert session.stdout == "sess_saved\n" + assert "interim session login" in session.stderr + assert "what requests send" in session.stderr + + +def test_auth_token_refuses_a_short_lived_login_without_a_fallback(monkeypatch, tmp_path) -> None: + threading.Thread = _REAL_THREAD + for name in ("LIGHTNING_AUTH_TOKEN", "LIGHTNING_API_KEY", "LIGHTNING_AUTH_TOKEN_FILE"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("LIGHTNING_CLI_OIDC_LOGIN", "1") + credentials = tmp_path / "credentials.json" + monkeypatch.setattr(Auth, "secrets_file", credentials) + monkeypatch.setattr(env, "LIGHTNING_LOGIN_PATH", str(tmp_path / "login.json")) + previous = keyring.get_keyring() + keyring.set_keyring(MemoryKeyring()) + fake = FakePlatform(access_ttl=60) + monkeypatch.setattr(env, "LIGHTNING_CLOUD_URL", fake.url) + browser = FakeBrowser(fake) + monkeypatch.setattr(webbrowser, "open", browser.open) + monkeypatch.setattr("lightning_sdk.cli.utils.auth._get_authed_user", lambda: SimpleNamespace(name="ada")) + try: + assert _invoke(main_cli, ["login"]).exit_code == 0 + access = CredentialSource().resolve().token + original = fake.token + + def refuse_refresh(form): + if form.get("grant_type") == "refresh_token": + return 500, {"error": "unavailable"} + return original(form) + + fake.token = refuse_refresh + result = _invoke(main_cli, ["auth", "token"]) + assert result.exit_code != 0 + assert "Couldn't refresh" in _text(result) + assert access not in _text(result) + finally: + browser.finish() + fake.stop() + keyring.set_keyring(previous) diff --git a/python/tests/cli/utils/test_auth.py b/python/tests/cli/utils/test_auth.py index 1292fec5e..0084420d0 100644 --- a/python/tests/cli/utils/test_auth.py +++ b/python/tests/cli/utils/test_auth.py @@ -1,442 +1,38 @@ -import base64 -import hashlib -import json -from unittest.mock import MagicMock, PropertyMock, patch -from urllib.parse import parse_qs, urlparse +"""require_auth_header reads the credential source and nothing patches the vendored login.""" +import click import pytest -import rich_click as click -from lightning_sdk.cli.legacy.deploy._auth import _AuthLitServe, _AuthMode -from lightning_sdk.cli.utils.auth import ( - OIDC_LOGIN_ENV, - browser_authentication, - login_with_pasted_code, - logout_session, - require_auth_header, -) -from lightning_sdk.lightning_cloud.login import Auth +from lightning_sdk.cli.utils.auth import require_auth_header +from lightning_sdk.lightning_cloud.login import Auth, AuthServer -def _response(status_code: int, body=None) -> MagicMock: - response = MagicMock(status_code=status_code, ok=status_code < 400) - response.json.return_value = body - return response +def test_require_auth_header_returns_the_environment_token(monkeypatch) -> None: + monkeypatch.setenv("LIGHTNING_AUTH_TOKEN", "env-token") + assert require_auth_header() == "Bearer env-token" -def _login_env(monkeypatch, tmp_path): - """Point credentials at tmp_path and return the secrets file.""" - secrets_file = tmp_path / "credentials.json" - monkeypatch.setattr(Auth, "secrets_file", secrets_file) - for key in ( - "LIGHTNING_USER_ID", - "LIGHTNING_API_KEY", - "LIGHTNING_AUTH_TOKEN", - "LIGHTNING_DEPLOYMENT_ID", - OIDC_LOGIN_ENV, - ): - monkeypatch.delenv(key, raising=False) - return secrets_file - -def _challenge(verifier: str) -> str: - return base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest()).rstrip(b"=").decode() - - -def _capture_browser_login(monkeypatch, tmp_path, token_response) -> dict: - """Run the loopback login with a fake server and browser; return what was captured.""" - captured: dict = {"secrets_file": _login_env(monkeypatch, tmp_path)} - monkeypatch.setenv(OIDC_LOGIN_ENV, "1") - - class FakeServer: - def __init__(self, config) -> None: - captured["config"] = config - captured["route"] = next(r for r in config.app.routes if getattr(r, "path", None) == "/callback") - - def run(self) -> None: - captured["run"](captured["route"].endpoint) - - post = MagicMock(side_effect=token_response) - monkeypatch.setattr("lightning_sdk.cli.utils.auth._generate_state", lambda: "fixed-state") - monkeypatch.setattr("lightning_sdk.cli.utils.auth.requests.head", MagicMock()) - monkeypatch.setattr("lightning_sdk.cli.utils.auth.requests.post", post) - monkeypatch.setattr("webbrowser.open", lambda url: captured.setdefault("auth_url", url) and True) - monkeypatch.setattr("uvicorn.Server", FakeServer) - captured["post"] = post - return captured - - -def test_require_auth_header_uses_environment_credentials() -> None: - """An environment header must be used without loading or authenticating.""" - auth = MagicMock(auth_header="Basic encoded") - - with patch("lightning_sdk.cli.utils.auth.Auth", return_value=auth): - assert require_auth_header() == "Basic encoded" - - auth.load.assert_not_called() - auth.authenticate.assert_not_called() - - -def test_require_auth_header_loads_saved_credentials() -> None: - """Saved credentials must work without starting browser authentication.""" - auth = MagicMock() - type(auth).auth_header = PropertyMock(side_effect=[ValueError("missing"), "Bearer token"]) - auth.load.return_value = True - - with patch("lightning_sdk.cli.utils.auth.Auth", return_value=auth): - assert require_auth_header() == "Bearer token" - - auth.authenticate.assert_not_called() - - -def test_require_auth_header_instructs_login_without_opening_browser() -> None: - """Missing credentials must stop ordinary commands and direct users to login.""" - auth = MagicMock() - type(auth).auth_header = PropertyMock(side_effect=ValueError("missing")) - auth.load.return_value = False - - with patch("lightning_sdk.cli.utils.auth.Auth", return_value=auth), pytest.raises( - click.UsageError, match="lightning login" - ): - require_auth_header() - - auth.authenticate.assert_not_called() - - -def test_browser_authentication_blocks_subclass_callback_server(monkeypatch, tmp_path) -> None: - """The guard must also block SDK auth subclasses with custom callback servers.""" - monkeypatch.setattr(Auth, "secrets_file", tmp_path / "credentials.json") - monkeypatch.delenv("LIGHTNING_AUTH_TOKEN", raising=False) - - class CustomAuth(Auth): - browser_started = False - - def load(self) -> bool: - return False - - def _run_server(self) -> None: - self.browser_started = True - - auth = CustomAuth(api_key="", auth_token="") - - with browser_authentication(False), pytest.raises(ValueError, match="lightning login"): - auth.authenticate() - - assert not auth.browser_started - - -def test_litserve_auth_uses_saved_credentials_without_starting_callback_server() -> None: - """Deployment authentication must not start callback auth when credentials exist.""" - auth = _AuthLitServe(_AuthMode.DEPLOY) - auth.api_key = "key" - - with patch.object( - auth, - "_run_server", - side_effect=AssertionError("callback server must not start"), - ): - assert auth.authenticate().startswith("Bearer ") - - -def test_litserve_auth_instructs_login_when_callback_auth_would_be_required() -> None: - """Deployment callback auth must tell the user to run the explicit login command.""" - with pytest.raises(click.UsageError, match="lightning login"): - _AuthLitServe(_AuthMode.DEPLOY)._run_server() - - -def test_auth_save_stores_session_token_as_auth_token(monkeypatch, tmp_path) -> None: - """Browser login callback passes the session token as `token`; it must be saved as `auth_token`.""" - secrets_file = tmp_path / "credentials.json" - monkeypatch.setattr(Auth, "secrets_file", secrets_file) - monkeypatch.delenv("LIGHTNING_USER_ID", raising=False) - monkeypatch.delenv("LIGHTNING_API_KEY", raising=False) +def test_require_auth_header_turns_a_missing_credential_into_a_usage_error(monkeypatch) -> None: monkeypatch.delenv("LIGHTNING_AUTH_TOKEN", raising=False) - - auth = Auth() - auth.save(token="sess_user-1.secret", user_id="user-1", api_key="") - - assert auth.auth_token == "sess_user-1.secret" - assert auth.api_key == "" - assert auth.auth_header == "Bearer sess_user-1.secret" - assert json.loads(secrets_file.read_text()) == { - "user_id": "user-1", - "api_key": "", - "auth_token": "sess_user-1.secret", - } - assert (secrets_file.stat().st_mode & 0o777) == 0o600 - - -def test_logout_session_revokes_server_session_and_clears_credentials(monkeypatch, tmp_path) -> None: - """`lightning logout` revokes the saved session token on the server and removes local credentials.""" - secrets_file = tmp_path / "credentials.json" - monkeypatch.setattr(Auth, "secrets_file", secrets_file) - monkeypatch.delenv("LIGHTNING_USER_ID", raising=False) monkeypatch.delenv("LIGHTNING_API_KEY", raising=False) - monkeypatch.delenv("LIGHTNING_AUTH_TOKEN", raising=False) - - auth = Auth() - auth.save(token="sess_user-1.secret", user_id="user-1") - - with patch("lightning_sdk.cli.utils.auth.requests.post") as post_mock: - logout_session(Auth()) - - post_mock.assert_called_once() - assert post_mock.call_args.kwargs["headers"] == {"Authorization": "Bearer sess_user-1.secret"} - assert not secrets_file.exists() - - -def test_login_with_browser_exchanges_the_code_with_pkce_on_loopback(monkeypatch, tmp_path) -> None: - """The callback trades the code and the verifier matching the sent challenge for a session.""" - from lightning_sdk.lightning_cloud.login import AuthServer - - captured = _capture_browser_login( - monkeypatch, tmp_path, [_response(200, {"access_token": "sess_user-1.secret", "token_type": "Bearer"})] - ) - captured["run"] = lambda endpoint: captured.setdefault("location", endpoint(code="the-code", state="fixed-state")) - - auth = Auth() - AuthServer().login_with_browser(auth) - - url = urlparse(captured["auth_url"]) - query = {k: v[0] for k, v in parse_qs(url.query).items()} - assert url.path == "/oidc/authorize" - assert query["client_id"] == "lightning-cli" - assert query["response_type"] == "code" - assert query["scope"] == "openid" - assert query["state"] == "fixed-state" - assert query["code_challenge_method"] == "S256" - assert query["redirect_uri"] == f"http://127.0.0.1:{captured['config'].port}/callback" - assert captured["config"].host == "127.0.0.1" - - data = captured["post"].call_args.kwargs["data"] - assert captured["post"].call_args.args[0].endswith("/oidc/token") - assert data["grant_type"] == "authorization_code" - assert data["code"] == "the-code" - assert data["redirect_uri"] == query["redirect_uri"] - assert _challenge(data["code_verifier"]) == query["code_challenge"] - - assert captured["location"].headers["location"].endswith("/cli-login-successful") - assert auth.auth_token == "sess_user-1.secret" - assert json.loads(captured["secrets_file"].read_text())["auth_token"] == "sess_user-1.secret" - assert (captured["secrets_file"].stat().st_mode & 0o777) == 0o600 - - -@pytest.mark.parametrize( - "params", - [ - {"code": "the-code", "state": ""}, - {"code": "the-code", "state": "forged-state"}, - {"code": "", "state": "fixed-state", "error": "access_denied"}, - {"code": "", "state": "fixed-state"}, - ], - ids=["no state", "forged state", "denied", "no code"], -) -def test_login_with_browser_refuses_a_bad_callback_without_exchanging(monkeypatch, tmp_path, params) -> None: - from lightning_sdk.lightning_cloud.login import AuthServer - - captured = _capture_browser_login(monkeypatch, tmp_path, []) - captured["run"] = lambda endpoint: captured.setdefault("location", endpoint(**params)) - - auth = Auth() - AuthServer().login_with_browser(auth) + monkeypatch.delenv("LIGHTNING_AUTH_TOKEN_FILE", raising=False) + opened = [] + monkeypatch.setattr("webbrowser.open", lambda *_args, **_kwargs: opened.append(True)) - assert captured["location"].headers["location"].endswith("/cli-login-failed") - captured["post"].assert_not_called() - assert not captured["secrets_file"].exists() - assert not auth.auth_token - - -def test_login_with_browser_reports_a_refused_exchange(monkeypatch, tmp_path) -> None: - from lightning_sdk.lightning_cloud.login import AuthServer - - captured = _capture_browser_login(monkeypatch, tmp_path, [_response(400, {"error": "invalid_grant"})]) - captured["run"] = lambda endpoint: captured.setdefault("location", endpoint(code="used", state="fixed-state")) - - AuthServer().login_with_browser(Auth()) - - assert captured["location"].headers["location"].endswith("/cli-login-failed") - assert not captured["secrets_file"].exists() - - -def test_login_with_browser_without_a_browser_falls_back_to_paste(monkeypatch, tmp_path) -> None: - from lightning_sdk.lightning_cloud.login import AuthServer - - secrets_file = _login_env(monkeypatch, tmp_path) - monkeypatch.setenv(OIDC_LOGIN_ENV, "1") - monkeypatch.setattr("lightning_sdk.cli.utils.auth.requests.head", MagicMock()) - monkeypatch.setattr("webbrowser.open", lambda url: False) - monkeypatch.setattr("uvicorn.Server", lambda config: (_ for _ in ()).throw(AssertionError("server started"))) - monkeypatch.setattr( - "lightning_sdk.cli.utils.auth.login_with_pasted_code", lambda auth: auth.save(auth_token="sess_x") - ) - - AuthServer().login_with_browser(Auth()) - - assert json.loads(secrets_file.read_text())["auth_token"] == "sess_x" - - -def test_login_with_browser_in_a_deployment_without_a_browser_raises(monkeypatch, tmp_path) -> None: - from lightning_sdk.lightning_cloud.login import AuthServer - - _login_env(monkeypatch, tmp_path) - monkeypatch.setenv("LIGHTNING_DEPLOYMENT_ID", "dep") - monkeypatch.setattr("lightning_sdk.cli.utils.auth.requests.head", MagicMock()) - monkeypatch.setattr("webbrowser.open", lambda url: False) - - with pytest.raises(RuntimeError, match="Include SDK credentials"): - AuthServer().login_with_browser(Auth()) - - -def test_login_with_browser_defaults_to_the_session_token_login(monkeypatch, tmp_path) -> None: - """Without the opt-in, the sign-in page hands a session token to `/login-complete`, and the callback - redirects to `/cli-login-successful` without forwarding it in the URL.""" - import asyncio - - from lightning_sdk.lightning_cloud.login import AuthServer - - _login_env(monkeypatch, tmp_path) - - captured: dict = {} - - class FakeServer: - def __init__(self, config) -> None: - captured["app"] = config.app - - def run(self) -> None: - route = next(r for r in captured["app"].routes if getattr(r, "path", None) == "/login-complete") - resp = asyncio.run( - route.endpoint(token="sess_user-1.secret", key="", state="fixed-state", user_id="user-1") - ) - captured["location"] = resp.headers["location"] - - monkeypatch.setattr("lightning_sdk.cli.utils.auth._generate_state", lambda: "fixed-state") - monkeypatch.setattr("lightning_sdk.cli.utils.auth.requests.head", MagicMock()) - monkeypatch.setattr("webbrowser.open", lambda url: captured.setdefault("auth_url", url) or True) - monkeypatch.setattr("uvicorn.Server", FakeServer) - - auth = Auth() - AuthServer().login_with_browser(auth) - - assert "/sign-in?" in captured["auth_url"] - assert "state=fixed-state" in captured["auth_url"] or "state%3Dfixed-state" in captured["auth_url"] - assert captured["location"].endswith("/cli-login-successful") - assert "token=" not in captured["location"] - assert auth.auth_token == "sess_user-1.secret" - - -def test_session_token_login_rejects_callback_without_matching_state(monkeypatch, tmp_path) -> None: - """Omitting or forging `state` must not store the callback token or key.""" - import asyncio - - from lightning_sdk.lightning_cloud.login import AuthServer - - secrets_file = _login_env(monkeypatch, tmp_path) - - captured: dict = {} - - class FakeServer: - def __init__(self, config) -> None: - captured["app"] = config.app - - def run(self) -> None: - captured["route"] = next(r for r in captured["app"].routes if getattr(r, "path", None) == "/login-complete") - - monkeypatch.setattr("lightning_sdk.cli.utils.auth._generate_state", lambda: "fixed-state") - monkeypatch.setattr("lightning_sdk.cli.utils.auth.requests.head", MagicMock()) - monkeypatch.setattr("webbrowser.open", lambda url: True) - monkeypatch.setattr("uvicorn.Server", FakeServer) - - auth = Auth() - AuthServer().login_with_browser(auth) - - for state in ("", "forged-state"): - resp = asyncio.run( - captured["route"].endpoint( - token="sess_attacker.secret", key="attacker-key", state=state, user_id="attacker" - ) - ) - assert resp.headers["location"].endswith("/cli-login-failed") - - assert not secrets_file.exists() - assert not auth.auth_token - assert not auth.api_key - - -def _paste_login(monkeypatch, tmp_path, codes: list, responses: list) -> tuple: - secrets_file = _login_env(monkeypatch, tmp_path) - post = MagicMock(side_effect=responses) - monkeypatch.setattr("lightning_sdk.cli.utils.auth.requests.post", post) - monkeypatch.setattr("lightning_sdk.cli.utils.auth.click.prompt", MagicMock(side_effect=codes)) - return secrets_file, post - - -def test_login_with_pasted_code_strips_a_trailing_slash_from_the_origin(monkeypatch, tmp_path, capsys) -> None: - monkeypatch.setattr("lightning_sdk.lightning_cloud.env.LIGHTNING_CLOUD_URL", "http://localhost:9800/") - _paste_login(monkeypatch, tmp_path, ["the-code"], [_response(200, {"access_token": "sess_user-1.secret"})]) - - login_with_pasted_code(Auth()) - - url = next(word for word in capsys.readouterr().out.split() if "/oidc/authorize" in word) - query = {k: v[0] for k, v in parse_qs(urlparse(url).query).items()} - assert url.startswith("http://localhost:9800/oidc/authorize?") - assert query["redirect_uri"] == "http://localhost:9800/cli-login" - assert query["scope"] == "openid" - - -def test_login_with_pasted_code_prints_the_authorize_url_and_saves_the_session(monkeypatch, tmp_path, capsys) -> None: - secrets_file, post = _paste_login( - monkeypatch, tmp_path, [" the-code \n"], [_response(200, {"access_token": "sess_user-1.secret"})] - ) - - login_with_pasted_code(Auth()) - - output = capsys.readouterr().out - url = next(word for word in output.split() if "/oidc/authorize" in word) - query = {k: v[0] for k, v in parse_qs(urlparse(url).query).items()} - assert query["client_id"] == "lightning-cli" - assert query["scope"] == "openid" - assert query["redirect_uri"].endswith("/cli-login") - assert "//cli-login" not in query["redirect_uri"] - assert query["code_challenge_method"] == "S256" - assert "never give it to anyone" in output - - data = post.call_args.kwargs["data"] - assert data["code"] == "the-code" - assert data["redirect_uri"] == query["redirect_uri"] - assert _challenge(data["code_verifier"]) == query["code_challenge"] - assert json.loads(secrets_file.read_text())["auth_token"] == "sess_user-1.secret" - - -def test_login_with_pasted_code_prompts_again_on_invalid_grant(monkeypatch, tmp_path) -> None: - secrets_file, post = _paste_login( - monkeypatch, - tmp_path, - ["typo", "the-code"], - [_response(400, {"error": "invalid_grant"}), _response(200, {"access_token": "sess_user-1.secret"})], - ) - - login_with_pasted_code(Auth()) - - assert [c.kwargs["data"]["code"] for c in post.call_args_list] == ["typo", "the-code"] - assert json.loads(secrets_file.read_text())["auth_token"] == "sess_user-1.secret" - - -def test_login_with_pasted_code_gives_up_after_three_bad_codes(monkeypatch, tmp_path) -> None: - refused = [_response(400, {"error": "invalid_grant"})] * 3 - secrets_file, post = _paste_login(monkeypatch, tmp_path, ["a", "b", "c"], refused) - - with pytest.raises(RuntimeError, match="not accepted"): - login_with_pasted_code(Auth()) - - assert post.call_count == 3 - assert not secrets_file.exists() + with pytest.raises(click.UsageError, match="lightning login"): + require_auth_header() + assert opened == [] -def test_login_with_pasted_code_stops_on_any_other_error(monkeypatch, tmp_path) -> None: - secrets_file, post = _paste_login(monkeypatch, tmp_path, ["a"], [_response(404, {})]) - with pytest.raises(RuntimeError, match="HTTP 404"): - login_with_pasted_code(Auth()) +def test_importing_the_cli_leaves_the_vendored_login_in_place() -> None: + import lightning_sdk.api.llm_api + import lightning_sdk.cli.entrypoint + import lightning_sdk.cli.legacy.deploy._auth - assert post.call_count == 1 - assert not secrets_file.exists() + assert lightning_sdk.cli.entrypoint.main_cli is not None + assert lightning_sdk.api.llm_api.authenticate is not None + assert lightning_sdk.cli.legacy.deploy._auth.authenticate is not None + assert Auth.save.__module__ == "lightning_sdk.lightning_cloud.login" + assert AuthServer.login_with_browser.__module__ == "lightning_sdk.lightning_cloud.login" From 458305ece070af43a7d2e31d361052783e64dfb1 Mon Sep 17 00:00:00 2001 From: fnwoke-lightning Date: Wed, 4 Nov 2026 20:04:13 -0500 Subject: [PATCH 07/27] fix(lightning_cloud): send the injected legacy key until the new login is switched on Co-authored-by: Cursor --- .../lightning_cloud/credentials.py | 53 +++++++--- python/lightning_sdk/lightning_cloud/login.py | 16 ++- .../lightning_cloud/rest_client.py | 3 +- python/tests/cli/utils/test_auth.py | 1 + .../tests/lightning_cloud/test_credentials.py | 100 +++++++++++++++++- 5 files changed, 149 insertions(+), 24 deletions(-) diff --git a/python/lightning_sdk/lightning_cloud/credentials.py b/python/lightning_sdk/lightning_cloud/credentials.py index 71e08c7c4..ba9578440 100644 --- a/python/lightning_sdk/lightning_cloud/credentials.py +++ b/python/lightning_sdk/lightning_cloud/credentials.py @@ -2,6 +2,7 @@ # Do not edit it directly; changes will be overwritten. # Make changes in the internal upstream repository instead. +import base64 import contextlib import contextvars import dataclasses @@ -29,12 +30,21 @@ class CredentialKind(str, enum.Enum): ENV_TOKEN = "env_token" LOGIN = "login" SESSION = "session" + LEGACY_KEY = "legacy_key" @dataclasses.dataclass(frozen=True) class Credential: kind: CredentialKind token: str = dataclasses.field(repr=False) + user_id: Optional[str] = None + + @property + def header(self) -> str: + if self.kind is CredentialKind.LEGACY_KEY and self.user_id: + raw = f"{self.user_id}:{self.token}".encode("ascii") + return f"Basic {base64.b64encode(raw).decode('ascii')}" + return f"Bearer {self.token}" class LoginUnavailable(ValueError): @@ -75,13 +85,21 @@ def _from_env_token() -> Optional[str]: return _present(os.environ.get(Keys.AUTH_TOKEN.value)) -# The order is the precedence. LIGHTNING_API_KEY and a saved api_key have no -# row, because the control plane accepts a token, never a legacy key. +def _from_legacy_key() -> Optional[str]: + # Studios and jobs still inject this key, and the control plane accepts the + # header main sends. The switch removes the row. + if cli_login.oidc_login_enabled(): + return None + return _present(os.environ.get(Keys.API_KEY.value)) + + +# The order is the precedence. A saved api_key has no row in either state. _READERS = ( (CredentialKind.TOKEN_FILE, _from_token_file), (CredentialKind.ENV_TOKEN, _from_env_token), (CredentialKind.LOGIN, cli_login._request_token), (CredentialKind.SESSION, cli_login._saved_session), + (CredentialKind.LEGACY_KEY, _from_legacy_key), ) @@ -104,12 +122,14 @@ def resolve(self) -> Optional[Credential]: """The credential a request would send now. Never opens a browser.""" for kind, read in _READERS: token = read() - if token: - return Credential(kind, token) + if not token: + continue + user_id = _present(os.environ.get(Keys.USER_ID.value)) if kind is CredentialKind.LEGACY_KEY else None + return Credential(kind, token, user_id) return None def header(self) -> str: - """``Bearer `` for a request, logging in once when nothing is saved.""" + """The authorization header for a request, logging in once when nothing is saved.""" credential = self._resolve() if credential is None and self._refusal is None: if not _browser_login_allowed.get(): @@ -123,17 +143,26 @@ def header(self) -> str: credential = self._resolve() if credential is None: raise ValueError(self._refusal) - return f"Bearer {credential.token}" + return credential.header - def refresh(self, token: str) -> bool: - """After the server refused ``token``, whether a different one is now available.""" + def refresh(self, sent: str) -> bool: + """Whether the request that sent ``sent`` should be retried. + + ``sent`` is the Authorization value (``Bearer `` or ``Basic ``), + or a bare token, which counts as ``Bearer ``. True when a different + credential is now available, or when the login refreshed the token it sent. + """ + header = sent if sent.startswith(("Bearer ", "Basic ")) else f"Bearer {sent}" current = self._resolve() if current is None: return False - if current.kind is CredentialKind.LOGIN and current.token == token: - fresh = self._noting_refusal(lambda: cli_login._request_token(token)) - return fresh is not None and fresh != token - return current.token != token + if current.header != header: + return True + if current.kind is not CredentialKind.LOGIN: + return False + token = current.token + fresh = self._noting_refusal(lambda: cli_login._request_token(token)) + return fresh is not None and fresh != token def _resolve(self) -> Optional[Credential]: return self._noting_refusal(self.resolve) diff --git a/python/lightning_sdk/lightning_cloud/login.py b/python/lightning_sdk/lightning_cloud/login.py index ee77cca76..c366f01bb 100644 --- a/python/lightning_sdk/lightning_cloud/login.py +++ b/python/lightning_sdk/lightning_cloud/login.py @@ -112,18 +112,22 @@ def clear(cls) -> None: logger.debug("credentials removed successfully") def get_auth_header(self, override: Optional[AuthOverride] = None) -> Optional[str]: - """Get a Bearer authentication header, with an optional override. + """Get an authorization header, with an optional override. By default, returns the credential every request would send: the file named by ``LIGHTNING_AUTH_TOKEN_FILE``, then ``LIGHTNING_AUTH_TOKEN``, - then the saved login. It never opens a browser. + then the saved login, then ``LIGHTNING_API_KEY`` while + ``LIGHTNING_CLI_OIDC_LOGIN`` is unset. That key is + ``Basic base64(:)`` when + ``LIGHTNING_USER_ID`` is set, and Bearer otherwise. A saved ``api_key`` + is never used. This never opens a browser. Parameters ---------- override : AuthOverride, optional Use a token held by this instance instead: - - "auth_token": this instance's auth token - - "api_key": this instance's API key + - "auth_token": this instance's auth token, as Bearer + - "api_key": this instance's API key, as Bearer - None: the credential every request would send (default) Returns @@ -180,7 +184,9 @@ def authenticate(self) -> Optional[str]: ---------- authorization header to use when authentication completes. With no credential, this logs in through ``_run_server``, unless browser login - is turned off. + is turned off. While ``LIGHTNING_CLI_OIDC_LOGIN`` is unset, an + ``LIGHTNING_API_KEY`` with nothing ahead of it is that header: Basic + when ``LIGHTNING_USER_ID`` is set, Bearer otherwise. """ from lightning_sdk.lightning_cloud.credentials import CredentialSource diff --git a/python/lightning_sdk/lightning_cloud/rest_client.py b/python/lightning_sdk/lightning_cloud/rest_client.py index 78568b316..4411abb42 100644 --- a/python/lightning_sdk/lightning_cloud/rest_client.py +++ b/python/lightning_sdk/lightning_cloud/rest_client.py @@ -294,8 +294,7 @@ def wrap(*args, **kwargs): raise if supplied: raise ConnectionError(_LOGIN_MESSAGE) - sent = headers["Authorization"][len("Bearer "):] - if not source.refresh(sent): + if not source.refresh(headers["Authorization"]): raise ConnectionError(_LOGIN_MESSAGE) headers["Authorization"] = source.header() try: diff --git a/python/tests/cli/utils/test_auth.py b/python/tests/cli/utils/test_auth.py index 1292fec5e..91b6d4659 100644 --- a/python/tests/cli/utils/test_auth.py +++ b/python/tests/cli/utils/test_auth.py @@ -107,6 +107,7 @@ def test_browser_authentication_blocks_subclass_callback_server(monkeypatch, tmp """The guard must also block SDK auth subclasses with custom callback servers.""" monkeypatch.setattr(Auth, "secrets_file", tmp_path / "credentials.json") monkeypatch.delenv("LIGHTNING_AUTH_TOKEN", raising=False) + monkeypatch.delenv("LIGHTNING_API_KEY", raising=False) class CustomAuth(Auth): browser_started = False diff --git a/python/tests/lightning_cloud/test_credentials.py b/python/tests/lightning_cloud/test_credentials.py index 44ff5faea..2091fd61a 100644 --- a/python/tests/lightning_cloud/test_credentials.py +++ b/python/tests/lightning_cloud/test_credentials.py @@ -1,3 +1,4 @@ +import base64 import importlib import importlib.metadata import json @@ -128,7 +129,9 @@ def request(method, url, headers=None, **kwargs): assert seen == ["Bearer sess_one", "Bearer sess_two"] -def test_no_request_sends_basic_whatever_user_id_and_the_saved_login_hold(platform, isolated_credentials, monkeypatch): +@_SWITCH +def test_a_saved_login_is_bearer_not_basic_in_either_switch_state(platform, isolated_credentials, monkeypatch, switch): + _set_switch(monkeypatch, switch) session = platform.mint_session("user-1") isolated_credentials.write_text(json.dumps({"user_id": "user-1", "api_key": "key-1", "auth_token": session})) monkeypatch.setenv("LIGHTNING_USER_ID", "user-1") @@ -147,7 +150,8 @@ def test_no_request_sends_basic_whatever_user_id_and_the_saved_login_hold(platfo assert _sent_to_the_api(platform) == [f"Bearer {session}"] -def test_no_request_sends_lightning_api_key_or_a_saved_api_key(platform, isolated_credentials, monkeypatch): +def test_with_the_switch_set_an_env_key_and_a_saved_api_key_are_not_sent(platform, isolated_credentials, monkeypatch): + _set_switch(monkeypatch, "oidc") isolated_credentials.write_text(json.dumps({"user_id": "user-1", "api_key": "key-saved"})) monkeypatch.setenv("LIGHTNING_USER_ID", "user-1") monkeypatch.setenv("LIGHTNING_API_KEY", "key-env") @@ -163,6 +167,89 @@ def test_no_request_sends_lightning_api_key_or_a_saved_api_key(platform, isolate assert platform.requests == [] +def test_with_the_switch_unset_a_user_id_and_api_key_are_sent_as_main_sends_them( + platform, isolated_credentials, monkeypatch +): + user_id = "user-1" + api_key = "key-1" + isolated_credentials.write_text(json.dumps({"user_id": "saved-user", "api_key": "key-saved"})) + monkeypatch.setenv("LIGHTNING_USER_ID", f" {user_id} ") + monkeypatch.setenv("LIGHTNING_API_KEY", f" {api_key}\n") + expected = f"Basic {base64.b64encode(f'{user_id}:{api_key}'.encode('ascii')).decode('ascii')}" + + with pytest.raises(ConnectionError): + LightningClient().auth_service_get_user() + + assert _sent_to_the_api(platform) == [expected] + assert Auth().auth_header == expected + + +def test_with_the_switch_unset_an_api_key_alone_is_sent_as_bearer(platform, monkeypatch): + monkeypatch.setenv("LIGHTNING_USER_ID", " \n") + monkeypatch.setenv("LIGHTNING_API_KEY", " key-only\n") + + with pytest.raises(ConnectionError): + LightningClient().auth_service_get_user() + + assert _sent_to_the_api(platform) == ["Bearer key-only"] + assert Auth().auth_header == "Bearer key-only" + + +def test_a_token_file_an_env_token_a_new_login_and_a_saved_session_each_beat_the_legacy_key( + platform, browser, isolated_credentials, monkeypatch, tmp_path +): + monkeypatch.setenv("LIGHTNING_USER_ID", "user-key") + monkeypatch.setenv("LIGHTNING_API_KEY", "key-env") + file_session = platform.mint_session("user-file") + env_session = platform.mint_session("user-env") + token_file = tmp_path / "token" + token_file.write_text(file_session + "\n") + monkeypatch.setenv(env.LIGHTNING_AUTH_TOKEN_FILE, str(token_file)) + monkeypatch.setenv("LIGHTNING_AUTH_TOKEN", env_session) + made = cli_login.login() + browser.finish() + + def sent(header): + platform.requests.clear() + user = LightningClient().auth_service_get_user().id + assert _sent_to_the_api(platform) == [header] + assert Auth().auth_header == header + return user + + assert sent(f"Bearer {file_session}") == "user-file" + token_file.unlink() + assert sent(f"Bearer {env_session}") == "user-env" + monkeypatch.delenv("LIGHTNING_AUTH_TOKEN") + assert sent(f"Bearer {made.access_token}") == "user-1" + logout() + saved = platform.mint_session("user-saved") + isolated_credentials.write_text(json.dumps({"user_id": "user-saved", "api_key": "key-saved", "auth_token": saved})) + assert sent(f"Bearer {saved}") == "user-saved" + + +@_SWITCH +def test_a_saved_api_key_is_never_sent(platform, isolated_credentials, monkeypatch, switch): + _set_switch(monkeypatch, switch) + isolated_credentials.write_text(json.dumps({"user_id": "user-1", "api_key": "key-saved"})) + + with browser_login(False), pytest.raises(ValueError, match=_exactly(_LOGIN_FIRST)): + LightningClient().auth_service_get_user() + with pytest.raises(ValueError, match=_exactly(_LOGIN_FIRST)): + _ = Auth().auth_header + assert platform.requests == [] + + +def test_a_401_on_the_legacy_key_says_run_lightning_login_once(platform, monkeypatch): + monkeypatch.setenv("LIGHTNING_USER_ID", "user-1") + monkeypatch.setenv("LIGHTNING_API_KEY", "key-1") + + with pytest.raises(ConnectionError) as refused: + LightningClient().auth_service_get_user() + + assert str(refused.value) == _LOGIN_FAILED + assert _sent_to_the_api(platform) == [f"Basic {base64.b64encode(b'user-1:key-1').decode('ascii')}"] + + @_SWITCH def test_with_no_credential_the_source_logs_in_once_with_the_login_the_switch_selects( platform, browser, monkeypatch, switch @@ -461,14 +548,17 @@ def test_token_file_env_and_session_cannot_refresh(isolated_credentials, monkeyp isolated_credentials.write_text(json.dumps({"auth_token": "sess_saved"})) source = CredentialSource() assert source.refresh("sess_saved") is False + assert source.refresh("Bearer sess_saved") is False monkeypatch.setenv("LIGHTNING_AUTH_TOKEN", "sess_env") assert source.refresh("sess_env") is False + assert source.refresh("Bearer sess_env") is False token_file = tmp_path / "token" token_file.write_text("lit_w_file") monkeypatch.setenv(env.LIGHTNING_AUTH_TOKEN_FILE, str(token_file)) assert source.refresh("lit_w_file") is False + assert source.refresh("Bearer lit_w_file") is False def test_a_token_file_rotated_since_the_refused_request_counts_as_refreshed(monkeypatch, tmp_path): @@ -503,7 +593,7 @@ def request(method, url, headers=None, **kwargs): client = GridRestClient(api_client=_Api(request), with_auth=True) assert client.api_client.request("GET", "http://example") == "ok" assert seen == ["Bearer old", "Bearer new"] - assert source.refused == ["old"] + assert source.refused == ["Bearer old"] def test_second_401_after_refresh_raises_once(monkeypatch): @@ -520,7 +610,7 @@ def request(method, url, headers=None, **kwargs): client.api_client.request("GET", "http://example") assert str(exc.value) == _LOGIN_FAILED assert len(seen) == 2 - assert source.refused == ["same"] + assert source.refused == ["Bearer same"] def test_non_refreshable_source_raises_on_the_first_401(monkeypatch): @@ -537,7 +627,7 @@ def request(method, url, headers=None, **kwargs): client.api_client.request("GET", "http://example") assert str(exc.value) == _LOGIN_FAILED assert seen == [1] - assert source.refused == ["static"] + assert source.refused == ["Bearer static"] def test_non_401_is_unchanged(monkeypatch): From 146d838da12aa475ede96e4a76838252acde2f28 Mon Sep 17 00:00:00 2001 From: fnwoke-lightning Date: Wed, 4 Nov 2026 20:10:04 -0500 Subject: [PATCH 08/27] fix(cli): keep main's behaviour for an injected legacy key until the switch Co-authored-by: Cursor --- python/lightning_sdk/api/utils.py | 3 +- python/lightning_sdk/cli/auth/status.py | 2 + python/lightning_sdk/cli/auth/token.py | 2 + python/lightning_sdk/cli/utils/auth.py | 5 ++ python/tests/api/test_credential_requests.py | 45 ++++++++++++- python/tests/api/test_library_login.py | 34 ++++++++++ python/tests/cli/test_login_commands.py | 69 ++++++++++++++++++-- 7 files changed, 152 insertions(+), 8 deletions(-) diff --git a/python/lightning_sdk/api/utils.py b/python/lightning_sdk/api/utils.py index 825b076e7..60fde2665 100644 --- a/python/lightning_sdk/api/utils.py +++ b/python/lightning_sdk/api/utils.py @@ -1281,12 +1281,11 @@ def __init__(self, source: Optional[CredentialSource] = None) -> None: def __call__(self, request: requests.PreparedRequest) -> requests.PreparedRequest: header = self._source.header() request.headers["Authorization"] = header - sent = header[len("Bearer ") :] if header.startswith("Bearer ") else header def handle(response: requests.Response, **kwargs: Any) -> requests.Response: if getattr(handle, "retried", False) or response.status_code != 401: return response - if not self._source.refresh(sent): + if not self._source.refresh(header): return response handle.retried = True # The connection stays checked out until the body is read, so the retry can reuse it. diff --git a/python/lightning_sdk/cli/auth/status.py b/python/lightning_sdk/cli/auth/status.py index 0fe61b066..74cf510c8 100644 --- a/python/lightning_sdk/cli/auth/status.py +++ b/python/lightning_sdk/cli/auth/status.py @@ -78,6 +78,8 @@ def _source(credential: Credential) -> str: login = cli_login.current_login() storage = login.storage if login is not None else "login" return f"login ({storage})" + if credential.kind is CredentialKind.LEGACY_KEY: + return "LIGHTNING_API_KEY (legacy key, until the new login is switched on)" return "interim session login" diff --git a/python/lightning_sdk/cli/auth/token.py b/python/lightning_sdk/cli/auth/token.py index c0c87fbc7..a18781ea9 100644 --- a/python/lightning_sdk/cli/auth/token.py +++ b/python/lightning_sdk/cli/auth/token.py @@ -28,6 +28,8 @@ def token() -> None: credential = CredentialSource().resolve() if credential is None: raise click.ClickException(_NOT_LOGGED_IN) + if credential.kind is CredentialKind.LEGACY_KEY: + raise click.ClickException("A legacy key is never printed. Run `lightning login`.") if credential.kind is CredentialKind.LOGIN: try: click.echo(cli_login.token()) diff --git a/python/lightning_sdk/cli/utils/auth.py b/python/lightning_sdk/cli/utils/auth.py index b031d4f1e..612b5278b 100644 --- a/python/lightning_sdk/cli/utils/auth.py +++ b/python/lightning_sdk/cli/utils/auth.py @@ -85,6 +85,11 @@ def run_login( "Unset it and run `lightning login` again." ) if _env_set("LIGHTNING_API_KEY"): + if not cli_login.oidc_login_enabled(): + click.echo( + '"lightning login" is not required within a Studio or when credentials are set in the environment' + ) + return click.echo( "LIGHTNING_API_KEY is ignored until service accounts. This login continues without it.", err=True, diff --git a/python/tests/api/test_credential_requests.py b/python/tests/api/test_credential_requests.py index a44f5a4c6..0b18b74f2 100644 --- a/python/tests/api/test_credential_requests.py +++ b/python/tests/api/test_credential_requests.py @@ -1,17 +1,19 @@ """Raw requests take a credential handler, and a 401 is retried once with a fresh token.""" +import base64 import threading import webbrowser from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer import keyring import pytest +import requests from lightning_sdk.api.billing_api import BillingApi from lightning_sdk.api.filesystem_api import FilesystemApi from lightning_sdk.api.studio_api import StudioApi from lightning_sdk.api.teamspace_api import TeamspaceApi -from lightning_sdk.api.utils import _BlobUploader +from lightning_sdk.api.utils import CredentialAuth, _BlobUploader from lightning_sdk.lightning_cloud import env from lightning_sdk.lightning_cloud.credentials import CredentialSource, start_login from lightning_sdk.lightning_cloud.login import Auth @@ -107,6 +109,47 @@ def test_billing_csv_retries_a_401(refreshed_token): assert refreshed_token.seen == ["old-token", "new-token"] +def test_a_401_on_a_basic_legacy_key_is_not_retried(monkeypatch): + threading.Thread = _REAL_THREAD + for name in ("LIGHTNING_AUTH_TOKEN", "LIGHTNING_AUTH_TOKEN_FILE", "LIGHTNING_CLI_OIDC_LOGIN"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("LIGHTNING_USER_ID", "user-1") + monkeypatch.setenv("LIGHTNING_API_KEY", "key-1") + seen = [] + expected = f"Basic {base64.b64encode(b'user-1:key-1').decode('ascii')}" + + class Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_GET(self): # noqa: N802 + seen.append(self.headers.get("Authorization")) + body = b'{"error":"unauthenticated"}' + self.send_response(401) + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, fmt, *args): + return + + httpd = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread = _REAL_THREAD(target=httpd.serve_forever, daemon=True) + thread.start() + try: + response = requests.get( + f"http://127.0.0.1:{httpd.server_address[1]}/v1/auth/user", + auth=CredentialAuth(), + timeout=10, + ) + finally: + httpd.shutdown() + thread.join(timeout=5) + httpd.server_close() + + assert response.status_code == 401 + assert seen == [expected] + + def test_blob_upload_retries_a_401(refreshed_token, tmp_path): local = tmp_path / "blob.bin" local.write_bytes(b"hi") diff --git a/python/tests/api/test_library_login.py b/python/tests/api/test_library_login.py index e7e048821..5fad87f38 100644 --- a/python/tests/api/test_library_login.py +++ b/python/tests/api/test_library_login.py @@ -69,6 +69,40 @@ def test_litai_puts_okbhrt_and_litaimodel_on_the_interim_sign_in(monkeypatch, tm platform.stop() +def test_litai_treats_a_legacy_key_as_an_existing_credential(monkeypatch): + monkeypatch.delenv("LIGHTNING_AUTH_TOKEN", raising=False) + monkeypatch.delenv("LIGHTNING_AUTH_TOKEN_FILE", raising=False) + monkeypatch.delenv("LIGHTNING_CLI_OIDC_LOGIN", raising=False) + monkeypatch.setenv("LIGHTNING_USER_ID", "user-1") + monkeypatch.setenv("LIGHTNING_API_KEY", "key-1") + asked = [] + monkeypatch.setattr( + "lightning_sdk.api.llm_api.Confirm.ask", + lambda *_args, **_kwargs: asked.append(True) or True, + ) + opened = [] + monkeypatch.setattr(webbrowser, "open", lambda *_args, **_kwargs: opened.append(True)) + + litai_authenticate("story-model") + + assert asked == [] + assert opened == [] + + +def test_litserve_treats_a_legacy_key_as_an_existing_credential(monkeypatch): + monkeypatch.delenv("LIGHTNING_AUTH_TOKEN", raising=False) + monkeypatch.delenv("LIGHTNING_AUTH_TOKEN_FILE", raising=False) + monkeypatch.delenv("LIGHTNING_CLI_OIDC_LOGIN", raising=False) + monkeypatch.setenv("LIGHTNING_USER_ID", "user-1") + monkeypatch.setenv("LIGHTNING_API_KEY", "key-1") + opened = [] + monkeypatch.setattr(webbrowser, "open", lambda *_args, **_kwargs: opened.append(True)) + + litserve_authenticate(_AuthMode.DEPLOY) + + assert opened == [] + + def test_litai_does_nothing_when_a_credential_already_exists(monkeypatch): monkeypatch.setenv("LIGHTNING_AUTH_TOKEN", "env-token") asked = [] diff --git a/python/tests/cli/test_login_commands.py b/python/tests/cli/test_login_commands.py index fdab1c44e..161168bb9 100644 --- a/python/tests/cli/test_login_commands.py +++ b/python/tests/cli/test_login_commands.py @@ -2,6 +2,7 @@ import getpass import json +import re import sys import threading import webbrowser @@ -28,6 +29,10 @@ def _text(result) -> str: return result.output + (getattr(result, "stderr", None) or "") +def _plain(result) -> str: + return " ".join(re.sub(r"\x1b\[[0-9;]*m", "", _text(result)).split()) + + @pytest.fixture() def isolated(monkeypatch, tmp_path): """A private credential store, no ambient token, and real threads for the fake servers.""" @@ -128,14 +133,36 @@ def test_login_refuses_a_token_file_and_an_environment_token(isolated, monkeypat assert not Auth.secrets_file.exists() -def test_login_warns_that_an_api_key_is_ignored_and_continues(platform, browser, monkeypatch) -> None: +_LOGIN_NOT_REQUIRED = '"lightning login" is not required within a Studio or when credentials are set in the environment' + + +@pytest.mark.parametrize("cli", [main_cli, legacy_cli]) +def test_with_the_switch_unset_login_says_an_api_key_needs_no_login(platform, browser, monkeypatch, cli) -> None: + monkeypatch.setenv("LIGHTNING_USER_ID", "user-1") + monkeypatch.setenv("LIGHTNING_API_KEY", "key-1") + + result = _invoke(cli, ["login"]) + + assert result.exit_code == 0, _text(result) + assert result.stdout == _LOGIN_NOT_REQUIRED + "\n" + assert "Successfully logged in" not in _text(result) + assert not platform.sign_ins + assert not platform.authorizations + assert not Auth.secrets_file.exists() + + +@pytest.mark.parametrize("cli", [main_cli, legacy_cli]) +def test_with_the_switch_set_login_warns_that_an_api_key_is_ignored_and_continues( + platform, browser, monkeypatch, cli +) -> None: + monkeypatch.setenv("LIGHTNING_CLI_OIDC_LOGIN", "1") monkeypatch.setenv("LIGHTNING_API_KEY", "lit_sa_ignored") - result = _invoke(main_cli, ["login"]) + result = _invoke(cli, ["login"]) assert result.exit_code == 0, _text(result) - assert "LIGHTNING_API_KEY is ignored until service accounts" in _text(result) - assert _saved_session().startswith("sess_") + assert "LIGHTNING_API_KEY is ignored until service accounts" in result.stderr + assert CredentialSource().resolve().token.startswith("lit_ua_") assert "lit_sa_ignored" not in _text(result) @@ -288,6 +315,26 @@ def test_auth_status_names_each_source_and_never_a_token(platform, browser, monk assert missing.exit_code != 0 assert _NOT_LOGGED_IN in _text(missing) + monkeypatch.setenv("LIGHTNING_USER_ID", "user-1") + monkeypatch.setenv("LIGHTNING_API_KEY", "key-1") + monkeypatch.delenv("LIGHTNING_CLI_OIDC_LOGIN", raising=False) + legacy = _invoke(main_cli, ["auth", "status"]) + shown = _plain(legacy) + assert legacy.exit_code == 0, _text(legacy) + assert "LIGHTNING_API_KEY (legacy key, until the new login is switched on)" in shown + assert "Orgs" in shown + assert "Teamspaces" in shown + assert "Permissions" in shown + assert "Expiry" in shown + assert shown.count("not known locally") == 4 + assert "key-1" not in shown + + monkeypatch.setenv("LIGHTNING_CLI_OIDC_LOGIN", "1") + hidden = _invoke(main_cli, ["auth", "status"]) + assert hidden.exit_code != 0 + assert _NOT_LOGGED_IN in _text(hidden) + assert "key-1" not in _text(hidden) + def test_auth_token_prints_what_requests_send_and_refuses_a_short_login(platform, monkeypatch) -> None: monkeypatch.setenv("LIGHTNING_AUTH_TOKEN", "env-token") @@ -298,13 +345,25 @@ def test_auth_token_prints_what_requests_send_and_refuses_a_short_login(platform assert "what requests send" in printed.stderr monkeypatch.delenv("LIGHTNING_AUTH_TOKEN") + monkeypatch.delenv("LIGHTNING_CLI_OIDC_LOGIN", raising=False) + monkeypatch.setenv("LIGHTNING_USER_ID", "user-1") monkeypatch.setenv("LIGHTNING_API_KEY", "lit_sa_secret") refused = _invoke(main_cli, ["auth", "token"]) assert refused.exit_code != 0 + assert refused.stdout == "" assert "lit_sa_secret" not in _text(refused) - assert _NOT_LOGGED_IN in _text(refused) + assert "A legacy key is never printed. Run `lightning login`." in _text(refused) + + monkeypatch.setenv("LIGHTNING_CLI_OIDC_LOGIN", "1") + hidden = _invoke(main_cli, ["auth", "token"]) + assert hidden.exit_code != 0 + assert hidden.stdout == "" + assert "lit_sa_secret" not in _text(hidden) + assert _NOT_LOGGED_IN in _text(hidden) + monkeypatch.delenv("LIGHTNING_CLI_OIDC_LOGIN") monkeypatch.delenv("LIGHTNING_API_KEY") + monkeypatch.delenv("LIGHTNING_USER_ID") Auth().save(auth_token="sess_saved") session = _invoke(main_cli, ["auth", "token"]) assert session.stdout == "sess_saved\n" From 47750b665c10f8c29746821f16807899e4a82e48 Mon Sep 17 00:00:00 2001 From: fnwoke-lightning Date: Wed, 4 Nov 2026 20:20:54 -0500 Subject: [PATCH 09/27] fix(cli): let login replace a saved session that outranks an injected key With the switch unset, a saved sess_ is sent ahead of LIGHTNING_API_KEY. When it went stale, requests said to run lightning login and login refused because the key was set. Login now returns early only when the key is what requests send. Co-authored-by: Cursor --- python/lightning_sdk/cli/utils/auth.py | 24 +++++++++++++++++------- python/tests/cli/test_login_commands.py | 17 +++++++++++++++++ 2 files changed, 34 insertions(+), 7 deletions(-) diff --git a/python/lightning_sdk/cli/utils/auth.py b/python/lightning_sdk/cli/utils/auth.py index 612b5278b..8fa459843 100644 --- a/python/lightning_sdk/cli/utils/auth.py +++ b/python/lightning_sdk/cli/utils/auth.py @@ -7,7 +7,13 @@ import click from lightning_sdk.lightning_cloud import cli_login, env -from lightning_sdk.lightning_cloud.credentials import CredentialKind, LoginUnavailable, logout, start_login +from lightning_sdk.lightning_cloud.credentials import ( + CredentialKind, + CredentialSource, + LoginUnavailable, + logout, + start_login, +) from lightning_sdk.lightning_cloud.credentials import browser_login as browser_authentication from lightning_sdk.lightning_cloud.login import Auth from lightning_sdk.utils.resolve import _get_authed_user, in_studio @@ -86,14 +92,18 @@ def run_login( ) if _env_set("LIGHTNING_API_KEY"): if not cli_login.oidc_login_enabled(): + # A saved login outranks the key, so a stale one can only be replaced by logging in. + sent = CredentialSource().resolve() + if sent is not None and sent.kind is CredentialKind.LEGACY_KEY: + click.echo( + '"lightning login" is not required within a Studio or when credentials are set in the environment' + ) + return + else: click.echo( - '"lightning login" is not required within a Studio or when credentials are set in the environment' + "LIGHTNING_API_KEY is ignored until service accounts. This login continues without it.", + err=True, ) - return - click.echo( - "LIGHTNING_API_KEY is ignored until service accounts. This login continues without it.", - err=True, - ) if in_studio(): raise RuntimeError("Unable to login within a Studio. Did you change your shell setup?") from None diff --git a/python/tests/cli/test_login_commands.py b/python/tests/cli/test_login_commands.py index 161168bb9..873d42a31 100644 --- a/python/tests/cli/test_login_commands.py +++ b/python/tests/cli/test_login_commands.py @@ -151,6 +151,23 @@ def test_with_the_switch_unset_login_says_an_api_key_needs_no_login(platform, br assert not Auth.secrets_file.exists() +def test_with_the_switch_unset_login_replaces_a_saved_session_that_outranks_an_api_key( + platform, browser, monkeypatch +) -> None: + monkeypatch.setenv("LIGHTNING_USER_ID", "user-1") + monkeypatch.setenv("LIGHTNING_API_KEY", "key-1") + previous = platform.mint_session("user-1") + Auth().save(auth_token=previous) + + result = _invoke(main_cli, ["login"]) + + assert result.exit_code == 0, _text(result) + assert _LOGIN_NOT_REQUIRED not in _text(result) + assert platform.sign_ins + assert previous in platform.logouts() + assert _saved_session() not in ("", previous) + + @pytest.mark.parametrize("cli", [main_cli, legacy_cli]) def test_with_the_switch_set_login_warns_that_an_api_key_is_ignored_and_continues( platform, browser, monkeypatch, cli From 8b8552db82bda6c179f42194271619462c507d74 Mon Sep 17 00:00:00 2001 From: fnwoke-lightning Date: Wed, 4 Nov 2026 22:16:50 -0500 Subject: [PATCH 10/27] fix(lightning_cloud): drop LIGHTNING_API_KEY from the credential source Co-authored-by: Cursor --- python/lightning_sdk/cli/entrypoint.py | 2 +- .../lightning_sdk/cli/resource_completion.py | 4 +- .../cli/utils/lockdown_status.py | 2 +- python/lightning_sdk/cli/utils/logging.py | 2 +- .../lightning_cloud/credentials.py | 28 +---- python/lightning_sdk/lightning_cloud/login.py | 31 ++--- python/tests/cli/test_entrypoint.py | 25 +++- python/tests/cli/utils/test_auth.py | 7 +- .../tests/lightning_cloud/test_credentials.py | 107 +++--------------- .../lightning_cloud/test_interim_login.py | 4 +- 10 files changed, 59 insertions(+), 153 deletions(-) diff --git a/python/lightning_sdk/cli/entrypoint.py b/python/lightning_sdk/cli/entrypoint.py index 3cb4d8ee7..4c4a5cb97 100644 --- a/python/lightning_sdk/cli/entrypoint.py +++ b/python/lightning_sdk/cli/entrypoint.py @@ -112,7 +112,7 @@ def login(no_browser: bool) -> None: if no_browser and not oidc_login_enabled(): raise click.UsageError("`lightning login --no-browser` is not available yet.") auth = Auth() - if auth.api_key or auth.auth_token: + if auth.auth_token: click.echo('"lightning login" is not required within a Studio or when credentials are set in the environment') return diff --git a/python/lightning_sdk/cli/resource_completion.py b/python/lightning_sdk/cli/resource_completion.py index ac613cb25..08cbf2d10 100644 --- a/python/lightning_sdk/cli/resource_completion.py +++ b/python/lightning_sdk/cli/resource_completion.py @@ -14,11 +14,11 @@ def has_credentials() -> bool: """Check for existing credentials without starting the browser login flow.""" auth = Auth() - if auth.api_key or auth.auth_token: + if auth.auth_token: return True if not auth.load(): return False - return bool(auth.api_key or auth.auth_token) + return bool(auth.auth_token) def accessible_teamspaces() -> dict[str, dict[str, str]]: diff --git a/python/lightning_sdk/cli/utils/lockdown_status.py b/python/lightning_sdk/cli/utils/lockdown_status.py index 178096db4..2ca6b4e3d 100644 --- a/python/lightning_sdk/cli/utils/lockdown_status.py +++ b/python/lightning_sdk/cli/utils/lockdown_status.py @@ -85,7 +85,7 @@ def _auth_header() -> Optional[str]: Reads the saved credentials unless the environment sets some, as the rest of the CLI does. """ auth = Auth() - if not (auth.api_key or auth.auth_token): + if not auth.auth_token: auth.load() try: return auth.get_auth_header() diff --git a/python/lightning_sdk/cli/utils/logging.py b/python/lightning_sdk/cli/utils/logging.py index 5158be6fe..b2c07fb7b 100644 --- a/python/lightning_sdk/cli/utils/logging.py +++ b/python/lightning_sdk/cli/utils/logging.py @@ -34,7 +34,7 @@ def _auth_header_without_browser() -> Optional[str]: auth = Auth() - if not (auth.api_key or auth.auth_token): + if not auth.auth_token: try: if not auth.load(): return None diff --git a/python/lightning_sdk/lightning_cloud/credentials.py b/python/lightning_sdk/lightning_cloud/credentials.py index ba9578440..d0eafab3f 100644 --- a/python/lightning_sdk/lightning_cloud/credentials.py +++ b/python/lightning_sdk/lightning_cloud/credentials.py @@ -2,7 +2,6 @@ # Do not edit it directly; changes will be overwritten. # Make changes in the internal upstream repository instead. -import base64 import contextlib import contextvars import dataclasses @@ -30,20 +29,15 @@ class CredentialKind(str, enum.Enum): ENV_TOKEN = "env_token" LOGIN = "login" SESSION = "session" - LEGACY_KEY = "legacy_key" @dataclasses.dataclass(frozen=True) class Credential: kind: CredentialKind token: str = dataclasses.field(repr=False) - user_id: Optional[str] = None @property def header(self) -> str: - if self.kind is CredentialKind.LEGACY_KEY and self.user_id: - raw = f"{self.user_id}:{self.token}".encode("ascii") - return f"Basic {base64.b64encode(raw).decode('ascii')}" return f"Bearer {self.token}" @@ -85,21 +79,12 @@ def _from_env_token() -> Optional[str]: return _present(os.environ.get(Keys.AUTH_TOKEN.value)) -def _from_legacy_key() -> Optional[str]: - # Studios and jobs still inject this key, and the control plane accepts the - # header main sends. The switch removes the row. - if cli_login.oidc_login_enabled(): - return None - return _present(os.environ.get(Keys.API_KEY.value)) - - -# The order is the precedence. A saved api_key has no row in either state. +# The order is the precedence. _READERS = ( (CredentialKind.TOKEN_FILE, _from_token_file), (CredentialKind.ENV_TOKEN, _from_env_token), (CredentialKind.LOGIN, cli_login._request_token), (CredentialKind.SESSION, cli_login._saved_session), - (CredentialKind.LEGACY_KEY, _from_legacy_key), ) @@ -124,8 +109,7 @@ def resolve(self) -> Optional[Credential]: token = read() if not token: continue - user_id = _present(os.environ.get(Keys.USER_ID.value)) if kind is CredentialKind.LEGACY_KEY else None - return Credential(kind, token, user_id) + return Credential(kind, token) return None def header(self) -> str: @@ -148,11 +132,11 @@ def header(self) -> str: def refresh(self, sent: str) -> bool: """Whether the request that sent ``sent`` should be retried. - ``sent`` is the Authorization value (``Bearer `` or ``Basic ``), - or a bare token, which counts as ``Bearer ``. True when a different - credential is now available, or when the login refreshed the token it sent. + ``sent`` is the Authorization value (``Bearer ``), or a bare token, + which counts as ``Bearer ``. True when a different credential is + now available, or when the login refreshed the token it sent. """ - header = sent if sent.startswith(("Bearer ", "Basic ")) else f"Bearer {sent}" + header = sent if sent.startswith("Bearer ") else f"Bearer {sent}" current = self._resolve() if current is None: return False diff --git a/python/lightning_sdk/lightning_cloud/login.py b/python/lightning_sdk/lightning_cloud/login.py index c366f01bb..4fd217ca2 100644 --- a/python/lightning_sdk/lightning_cloud/login.py +++ b/python/lightning_sdk/lightning_cloud/login.py @@ -29,12 +29,11 @@ logger = logging.getLogger(__name__) # Authentication override types -AuthOverride = Literal["auth_token", "api_key"] +AuthOverride = Literal["auth_token"] class Keys(Enum): # USERNAME = "LIGHTNING_USERNAME" USER_ID = "LIGHTNING_USER_ID" - API_KEY = "LIGHTNING_API_KEY" AUTH_TOKEN = "LIGHTNING_AUTH_TOKEN" @property @@ -46,7 +45,6 @@ def suffix(self): class Auth: # username: Optional[str] = None user_id: Optional[str] = None - api_key: Optional[str] = None auth_token: Optional[str] = None secrets_file = pathlib.Path(env.LIGHTNING_CREDENTIAL_PATH) @@ -81,10 +79,12 @@ def save( auth_token: str = "", # username: str = "", ) -> None: - """Replace the saved login with the given fields. ``token`` is ignored.""" + """Replace the saved login. ``token`` and ``api_key`` are ignored. + + The CLI on main still calls ``save(token=..., api_key=...)``. + """ stored = { Keys.USER_ID.suffix: user_id, - Keys.API_KEY.suffix: api_key, Keys.AUTH_TOKEN.suffix: auth_token, } @@ -116,18 +116,14 @@ def get_auth_header(self, override: Optional[AuthOverride] = None) -> Optional[s By default, returns the credential every request would send: the file named by ``LIGHTNING_AUTH_TOKEN_FILE``, then ``LIGHTNING_AUTH_TOKEN``, - then the saved login, then ``LIGHTNING_API_KEY`` while - ``LIGHTNING_CLI_OIDC_LOGIN`` is unset. That key is - ``Basic base64(:)`` when - ``LIGHTNING_USER_ID`` is set, and Bearer otherwise. A saved ``api_key`` - is never used. This never opens a browser. + then the saved login, then the interim session. This never opens a + browser. Parameters ---------- override : AuthOverride, optional Use a token held by this instance instead: - "auth_token": this instance's auth token, as Bearer - - "api_key": this instance's API key, as Bearer - None: the credential every request would send (default) Returns @@ -148,15 +144,6 @@ def get_auth_header(self, override: Optional[AuthOverride] = None) -> Optional[s ) return f"Bearer {self.auth_token}" - elif override == "api_key": - if not self.api_key: - raise ValueError( - "API key override requested but no API key available. " - "Please set LIGHTNING_API_KEY environment variable " - "or use authenticate() method." - ) - return f"Bearer {self.api_key}" - elif override is None: from lightning_sdk.lightning_cloud.credentials import CredentialSource, browser_login @@ -184,9 +171,7 @@ def authenticate(self) -> Optional[str]: ---------- authorization header to use when authentication completes. With no credential, this logs in through ``_run_server``, unless browser login - is turned off. While ``LIGHTNING_CLI_OIDC_LOGIN`` is unset, an - ``LIGHTNING_API_KEY`` with nothing ahead of it is that header: Basic - when ``LIGHTNING_USER_ID`` is set, Bearer otherwise. + is turned off. """ from lightning_sdk.lightning_cloud.credentials import CredentialSource diff --git a/python/tests/cli/test_entrypoint.py b/python/tests/cli/test_entrypoint.py index ba8e6c302..f5c0ef4aa 100644 --- a/python/tests/cli/test_entrypoint.py +++ b/python/tests/cli/test_entrypoint.py @@ -345,10 +345,9 @@ def _complete_browser_auth(auth: Auth) -> None: assert auth.auth_token == "sess_new" -@pytest.mark.parametrize("env_var", ["LIGHTNING_AUTH_TOKEN", "LIGHTNING_API_KEY"]) -def test_login_not_required_with_credentials_in_env(monkeypatch, tmp_path, env_var): +def test_login_not_required_with_credentials_in_env(monkeypatch, tmp_path): _clear_auth(monkeypatch, tmp_path) - monkeypatch.setenv(env_var, "from-env") + monkeypatch.setenv("LIGHTNING_AUTH_TOKEN", "from-env") Auth().save(auth_token="sess_saved") def _fail_browser_auth(auth: Auth) -> None: @@ -363,6 +362,26 @@ def _fail_browser_auth(auth: Auth) -> None: assert Auth.secrets_file.exists() +def test_an_env_api_key_does_not_skip_login(monkeypatch, tmp_path): + _clear_auth(monkeypatch, tmp_path) + monkeypatch.setenv("LIGHTNING_API_KEY", "from-env") + Auth().save(auth_token="sess_saved") + + def _complete_browser_auth(auth: Auth) -> None: + assert not Auth.secrets_file.exists() + auth.save(auth_token="sess_new") + + monkeypatch.setattr(Auth, "_run_server", _complete_browser_auth) + + result = CliRunner().invoke(main_cli, ["login"]) + + assert result.exit_code == 0, result.output + assert "credentials are set in the environment" not in result.output + auth = Auth() + assert auth.load() + assert auth.auth_token == "sess_new" + + @patch.dict(os.environ, {"LIGHTNING_INTERACTIVE": "true", "LIGHTNING_CLOUD_SPACE_ID": "cloud-space-id"}, clear=True) @mock_command_logging def test_cannot_login_within_studio_without_credentials(monkeypatch): diff --git a/python/tests/cli/utils/test_auth.py b/python/tests/cli/utils/test_auth.py index 91b6d4659..b8dbfad84 100644 --- a/python/tests/cli/utils/test_auth.py +++ b/python/tests/cli/utils/test_auth.py @@ -107,7 +107,6 @@ def test_browser_authentication_blocks_subclass_callback_server(monkeypatch, tmp """The guard must also block SDK auth subclasses with custom callback servers.""" monkeypatch.setattr(Auth, "secrets_file", tmp_path / "credentials.json") monkeypatch.delenv("LIGHTNING_AUTH_TOKEN", raising=False) - monkeypatch.delenv("LIGHTNING_API_KEY", raising=False) class CustomAuth(Auth): browser_started = False @@ -118,7 +117,7 @@ def load(self) -> bool: def _run_server(self) -> None: self.browser_started = True - auth = CustomAuth(api_key="", auth_token="") + auth = CustomAuth(auth_token="") with browser_authentication(False), pytest.raises(ValueError, match="lightning login"): auth.authenticate() @@ -129,7 +128,6 @@ def _run_server(self) -> None: def test_litserve_auth_uses_saved_credentials_without_starting_callback_server() -> None: """Deployment authentication must not start callback auth when credentials exist.""" auth = _AuthLitServe(_AuthMode.DEPLOY) - auth.api_key = "key" with patch.object( auth, @@ -157,11 +155,9 @@ def test_auth_save_stores_session_token_as_auth_token(monkeypatch, tmp_path) -> auth.save(token="sess_user-1.secret", user_id="user-1", api_key="") assert auth.auth_token == "sess_user-1.secret" - assert auth.api_key == "" assert auth.auth_header == "Bearer sess_user-1.secret" assert json.loads(secrets_file.read_text()) == { "user_id": "user-1", - "api_key": "", "auth_token": "sess_user-1.secret", } assert (secrets_file.stat().st_mode & 0o777) == 0o600 @@ -360,7 +356,6 @@ def run(self) -> None: assert not secrets_file.exists() assert not auth.auth_token - assert not auth.api_key def _paste_login(monkeypatch, tmp_path, codes: list, responses: list) -> tuple: diff --git a/python/tests/lightning_cloud/test_credentials.py b/python/tests/lightning_cloud/test_credentials.py index 2091fd61a..b15510c7b 100644 --- a/python/tests/lightning_cloud/test_credentials.py +++ b/python/tests/lightning_cloud/test_credentials.py @@ -1,4 +1,3 @@ -import base64 import importlib import importlib.metadata import json @@ -144,17 +143,23 @@ def test_a_saved_login_is_bearer_not_basic_in_either_switch_state(platform, isol Auth().authenticate(), Auth().auth_header, Auth().get_auth_header(), - Auth().get_auth_header("api_key"), ] assert all(header.startswith("Bearer ") for header in headers), headers assert _sent_to_the_api(platform) == [f"Bearer {session}"] -def test_with_the_switch_set_an_env_key_and_a_saved_api_key_are_not_sent(platform, isolated_credentials, monkeypatch): - _set_switch(monkeypatch, "oidc") - isolated_credentials.write_text(json.dumps({"user_id": "user-1", "api_key": "key-saved"})) - monkeypatch.setenv("LIGHTNING_USER_ID", "user-1") - monkeypatch.setenv("LIGHTNING_API_KEY", "key-env") +@pytest.mark.parametrize("user_id", ["user-1", None]) +@_SWITCH +def test_an_env_api_key_and_a_saved_api_key_are_never_sent( + platform, isolated_credentials, monkeypatch, switch, user_id +): + _set_switch(monkeypatch, switch) + isolated_credentials.write_text(json.dumps({"user_id": "saved-user", "api_key": "key-saved"})) + if user_id is None: + monkeypatch.delenv("LIGHTNING_USER_ID", raising=False) + else: + monkeypatch.setenv("LIGHTNING_USER_ID", f" {user_id} ") + monkeypatch.setenv("LIGHTNING_API_KEY", " key-env\n") assert CredentialSource().resolve() is None with browser_login(False): @@ -167,89 +172,6 @@ def test_with_the_switch_set_an_env_key_and_a_saved_api_key_are_not_sent(platfor assert platform.requests == [] -def test_with_the_switch_unset_a_user_id_and_api_key_are_sent_as_main_sends_them( - platform, isolated_credentials, monkeypatch -): - user_id = "user-1" - api_key = "key-1" - isolated_credentials.write_text(json.dumps({"user_id": "saved-user", "api_key": "key-saved"})) - monkeypatch.setenv("LIGHTNING_USER_ID", f" {user_id} ") - monkeypatch.setenv("LIGHTNING_API_KEY", f" {api_key}\n") - expected = f"Basic {base64.b64encode(f'{user_id}:{api_key}'.encode('ascii')).decode('ascii')}" - - with pytest.raises(ConnectionError): - LightningClient().auth_service_get_user() - - assert _sent_to_the_api(platform) == [expected] - assert Auth().auth_header == expected - - -def test_with_the_switch_unset_an_api_key_alone_is_sent_as_bearer(platform, monkeypatch): - monkeypatch.setenv("LIGHTNING_USER_ID", " \n") - monkeypatch.setenv("LIGHTNING_API_KEY", " key-only\n") - - with pytest.raises(ConnectionError): - LightningClient().auth_service_get_user() - - assert _sent_to_the_api(platform) == ["Bearer key-only"] - assert Auth().auth_header == "Bearer key-only" - - -def test_a_token_file_an_env_token_a_new_login_and_a_saved_session_each_beat_the_legacy_key( - platform, browser, isolated_credentials, monkeypatch, tmp_path -): - monkeypatch.setenv("LIGHTNING_USER_ID", "user-key") - monkeypatch.setenv("LIGHTNING_API_KEY", "key-env") - file_session = platform.mint_session("user-file") - env_session = platform.mint_session("user-env") - token_file = tmp_path / "token" - token_file.write_text(file_session + "\n") - monkeypatch.setenv(env.LIGHTNING_AUTH_TOKEN_FILE, str(token_file)) - monkeypatch.setenv("LIGHTNING_AUTH_TOKEN", env_session) - made = cli_login.login() - browser.finish() - - def sent(header): - platform.requests.clear() - user = LightningClient().auth_service_get_user().id - assert _sent_to_the_api(platform) == [header] - assert Auth().auth_header == header - return user - - assert sent(f"Bearer {file_session}") == "user-file" - token_file.unlink() - assert sent(f"Bearer {env_session}") == "user-env" - monkeypatch.delenv("LIGHTNING_AUTH_TOKEN") - assert sent(f"Bearer {made.access_token}") == "user-1" - logout() - saved = platform.mint_session("user-saved") - isolated_credentials.write_text(json.dumps({"user_id": "user-saved", "api_key": "key-saved", "auth_token": saved})) - assert sent(f"Bearer {saved}") == "user-saved" - - -@_SWITCH -def test_a_saved_api_key_is_never_sent(platform, isolated_credentials, monkeypatch, switch): - _set_switch(monkeypatch, switch) - isolated_credentials.write_text(json.dumps({"user_id": "user-1", "api_key": "key-saved"})) - - with browser_login(False), pytest.raises(ValueError, match=_exactly(_LOGIN_FIRST)): - LightningClient().auth_service_get_user() - with pytest.raises(ValueError, match=_exactly(_LOGIN_FIRST)): - _ = Auth().auth_header - assert platform.requests == [] - - -def test_a_401_on_the_legacy_key_says_run_lightning_login_once(platform, monkeypatch): - monkeypatch.setenv("LIGHTNING_USER_ID", "user-1") - monkeypatch.setenv("LIGHTNING_API_KEY", "key-1") - - with pytest.raises(ConnectionError) as refused: - LightningClient().auth_service_get_user() - - assert str(refused.value) == _LOGIN_FAILED - assert _sent_to_the_api(platform) == [f"Basic {base64.b64encode(b'user-1:key-1').decode('ascii')}"] - - @_SWITCH def test_with_no_credential_the_source_logs_in_once_with_the_login_the_switch_selects( platform, browser, monkeypatch, switch @@ -688,9 +610,10 @@ def test_save_replaces_the_saved_login(isolated_credentials): auth = Auth() auth.save(user_id="user-1", api_key="key-1") auth.save(token="ignored", auth_token="sess_1") - assert json.loads(path.read_text()) == {"user_id": "", "api_key": "", "auth_token": "sess_1"} + assert json.loads(path.read_text()) == {"user_id": "", "auth_token": "sess_1"} + assert "api_key" not in path.read_text() + assert "key-1" not in path.read_text() assert not auth.user_id - assert not auth.api_key assert auth.auth_token == "sess_1" diff --git a/python/tests/lightning_cloud/test_interim_login.py b/python/tests/lightning_cloud/test_interim_login.py index 0ac4ef773..7e3ce7694 100644 --- a/python/tests/lightning_cloud/test_interim_login.py +++ b/python/tests/lightning_cloud/test_interim_login.py @@ -41,7 +41,7 @@ def test_the_interim_login_saves_the_sess_as_auth_token_and_the_next_request_sen [sign_in] = platform.sign_ins assert re.fullmatch(r"http://localhost:\d+/login-complete\?state=[\w-]+", sign_in["redirectTo"]) saved = _saved() - assert saved == {"user_id": "user-1", "api_key": "", "auth_token": saved["auth_token"]} + assert saved == {"user_id": "user-1", "auth_token": saved["auth_token"]} assert saved["auth_token"].startswith("sess_") assert browser.landed == f"{platform.url}/cli-login-successful" assert CredentialSource().resolve().kind is CredentialKind.SESSION @@ -118,7 +118,7 @@ def test_the_interim_login_never_saves_the_legacy_key_the_callback_carries(platf start_login() browser.finish() - assert _saved()["api_key"] == "" + assert "api_key" not in _saved() assert "key_legacy" not in Auth.secrets_file.read_text() From 41b41768e464430989a5cf20fdbc2c194dce5011 Mon Sep 17 00:00:00 2001 From: fnwoke-lightning Date: Wed, 4 Nov 2026 22:17:24 -0500 Subject: [PATCH 11/27] fix(lightning_cloud): leave login hints to M4-8b Co-authored-by: Cursor --- .../lightning_cloud/cli_login.py | 20 +++---------------- .../lightning_cloud/credentials.py | 8 +------- .../tests/lightning_cloud/test_cli_login.py | 18 ++++------------- .../lightning_cloud/test_interim_login.py | 7 ++----- 4 files changed, 10 insertions(+), 43 deletions(-) diff --git a/python/lightning_sdk/lightning_cloud/cli_login.py b/python/lightning_sdk/lightning_cloud/cli_login.py index d681ef11b..64d9aa764 100644 --- a/python/lightning_sdk/lightning_cloud/cli_login.py +++ b/python/lightning_sdk/lightning_cloud/cli_login.py @@ -18,7 +18,7 @@ import time import webbrowser from datetime import datetime -from typing import Any, Callable, Dict, Iterator, Optional, Sequence, Tuple +from typing import Any, Callable, Dict, Iterator, Optional, Tuple from urllib.parse import parse_qs, urlencode, urlsplit import requests @@ -98,8 +98,6 @@ def oidc_login_enabled() -> bool: def login( no_browser: bool = False, insecure_storage: bool = False, - orgs: Optional[Sequence[str]] = None, - project: Optional[str] = None, ) -> Login: """Log in through the browser, or by pasting the code the login page shows.""" if insecure_storage: @@ -108,7 +106,7 @@ def login( storage = _KeychainStorage.kind else: raise LoginError(_NO_KEYCHAIN) - attempt = _Attempt(_origin(), storage, orgs, project) + attempt = _Attempt(_origin(), storage) if not no_browser: done = _browser_login(attempt) if done is not None: @@ -157,22 +155,11 @@ def _request_token(rejected: Optional[str] = None) -> Optional[str]: class _Attempt: """One login: its state, its PKCE verifier and where it will be stored.""" - def __init__( - self, - origin: str, - storage: str, - orgs: Optional[Sequence[str]], - project: Optional[str], - ) -> None: + def __init__(self, origin: str, storage: str) -> None: self.origin = origin self.storage = storage self.state = secrets.token_urlsafe(32) self._verifier = secrets.token_urlsafe(48) - self._hints = {} - if orgs: - self._hints["orgs_hint"] = orgs if isinstance(orgs, str) else ",".join(orgs) - if project: - self._hints["project_hint"] = project def authorize_url(self, redirect_uri: str) -> str: digest = hashlib.sha256(self._verifier.encode()).digest() @@ -184,7 +171,6 @@ def authorize_url(self, redirect_uri: str) -> str: "state": self.state, "code_challenge": base64.urlsafe_b64encode(digest).rstrip(b"=").decode(), "code_challenge_method": "S256", - **self._hints, } return f"{self.origin}/oidc/authorize?{urlencode(query)}" diff --git a/python/lightning_sdk/lightning_cloud/credentials.py b/python/lightning_sdk/lightning_cloud/credentials.py index d0eafab3f..cb1cabb17 100644 --- a/python/lightning_sdk/lightning_cloud/credentials.py +++ b/python/lightning_sdk/lightning_cloud/credentials.py @@ -8,7 +8,7 @@ import enum import os import pathlib -from typing import Any, Callable, Dict, Iterator, Optional, Sequence, Tuple +from typing import Any, Callable, Dict, Iterator, Optional, Tuple from lightning_sdk.lightning_cloud import cli_login, env from lightning_sdk.lightning_cloud.login import Auth, AuthServer, Keys @@ -164,8 +164,6 @@ def _noting_refusal(self, read: Callable[[], Any]) -> Any: def start_login( no_browser: bool = False, insecure_storage: bool = False, - orgs: Optional[Sequence[str]] = None, - project: Optional[str] = None, sign_in_params: Optional[Dict[str, str]] = None, ) -> Optional[cli_login.Login]: """Log in with the login ``LIGHTNING_CLI_OIDC_LOGIN`` selects when this is called. @@ -178,16 +176,12 @@ def start_login( return cli_login.login( no_browser=no_browser, insecure_storage=insecure_storage, - orgs=orgs, - project=project, ) unavailable = [ flag for flag, given in ( ("--no-browser", no_browser), ("--insecure-storage", insecure_storage), - ("--org", orgs), - ("--project", project), ) if given ] diff --git a/python/tests/lightning_cloud/test_cli_login.py b/python/tests/lightning_cloud/test_cli_login.py index 0e5925ebc..d0f3ca553 100644 --- a/python/tests/lightning_cloud/test_cli_login.py +++ b/python/tests/lightning_cloud/test_cli_login.py @@ -97,22 +97,12 @@ def test_status_shows_grant_storage_and_expiries_without_tokens(platform, browse assert not hasattr(status, "refresh_token") -def test_org_and_project_hints_reach_the_consent_page(platform, browser): - cli_login.login(orgs=["acme", "beta"], project="research") - - consent = platform.authorizations[0] - assert consent["orgs_hint"] == "acme,beta" - assert consent["project_hint"] == "research" - - -def test_paste_back_sends_the_same_hints_as_the_browser(platform, monkeypatch, capsys): - _paste_codes(monkeypatch, platform, capsys, ["CODE"]) - - cli_login.login(no_browser=True, orgs=["acme", "beta"], project="research") +def test_the_authorize_url_carries_no_org_or_project_hint(platform, browser): + cli_login.login() consent = platform.authorizations[0] - assert consent["orgs_hint"] == "acme,beta" - assert consent["project_hint"] == "research" + assert "orgs_hint" not in consent + assert "project_hint" not in consent def test_the_loopback_answers_other_paths_with_404_and_keeps_waiting(platform, monkeypatch): diff --git a/python/tests/lightning_cloud/test_interim_login.py b/python/tests/lightning_cloud/test_interim_login.py index 7e3ce7694..a638d690a 100644 --- a/python/tests/lightning_cloud/test_interim_login.py +++ b/python/tests/lightning_cloud/test_interim_login.py @@ -184,8 +184,6 @@ def test_the_switch_is_read_when_each_login_starts(platform, monkeypatch): [ ({"no_browser": True}, "--no-browser"), ({"insecure_storage": True}, "--insecure-storage"), - ({"orgs": ["acme"]}, "--org"), - ({"project": "research"}, "--project"), ], ) def test_without_the_switch_an_option_only_the_new_login_has_is_refused(platform, monkeypatch, options, named): @@ -200,11 +198,10 @@ def test_without_the_switch_an_option_only_the_new_login_has_is_refused(platform assert platform.requests == [] -def test_with_the_switch_the_new_login_takes_the_hints_and_ignores_sign_in_params(platform, browser, _oidc_login): - made = start_login(orgs=["acme"], project="research", sign_in_params={"okbhrt": "code"}) +def test_with_the_switch_the_new_login_ignores_sign_in_params(platform, browser, _oidc_login): + made = start_login(sign_in_params={"okbhrt": "code"}) [consent] = platform.authorizations - assert (consent["orgs_hint"], consent["project_hint"]) == ("acme", "research") assert "okbhrt" not in consent assert made.grant == {"orgs": "all", "permissions": ["read", "write"]} From c432245ee02419b10eb7153e9912fa69e57f0152 Mon Sep 17 00:00:00 2001 From: fnwoke-lightning Date: Wed, 4 Nov 2026 22:20:27 -0500 Subject: [PATCH 12/27] fix(lightning_cloud): forget an unreadable login at logout and report failed revokes on login Co-authored-by: Cursor --- .../lightning_cloud/cli_login.py | 15 ++++++++--- .../lightning_cloud/credentials.py | 15 ++++++++--- .../tests/lightning_cloud/test_cli_login.py | 12 +++++++++ .../lightning_cloud/test_replaced_login.py | 26 +++++++++++++++++++ 4 files changed, 61 insertions(+), 7 deletions(-) diff --git a/python/lightning_sdk/lightning_cloud/cli_login.py b/python/lightning_sdk/lightning_cloud/cli_login.py index 64d9aa764..b06414b2f 100644 --- a/python/lightning_sdk/lightning_cloud/cli_login.py +++ b/python/lightning_sdk/lightning_cloud/cli_login.py @@ -639,10 +639,17 @@ def _replace_saved(save: Callable[[], Optional[str]]) -> None: replaced_login = _load(_origin(), cached=False) replaced_session = _saved_session() kept_session = save() - if replaced_login is not None: - _revoke(replaced_login.refresh_token) - if replaced_session is not None and replaced_session != kept_session: - _revoke_session(replaced_session) + # credentials imports this module, so the shared wording is imported at the call. + from lightning_sdk.lightning_cloud.credentials import REVOKE_FAILED, CredentialKind + + if replaced_login is not None and not _revoke(replaced_login.refresh_token): + _say(REVOKE_FAILED[CredentialKind.LOGIN]) + if ( + replaced_session is not None + and replaced_session != kept_session + and not _revoke_session(replaced_session) + ): + _say(REVOKE_FAILED[CredentialKind.SESSION]) def _login_from(body: Dict[str, Any], storage: str, url: str) -> Login: diff --git a/python/lightning_sdk/lightning_cloud/credentials.py b/python/lightning_sdk/lightning_cloud/credentials.py index cb1cabb17..e15285ab8 100644 --- a/python/lightning_sdk/lightning_cloud/credentials.py +++ b/python/lightning_sdk/lightning_cloud/credentials.py @@ -31,6 +31,12 @@ class CredentialKind(str, enum.Enum): SESSION = "session" +REVOKE_FAILED = { + CredentialKind.LOGIN: "The saved login could not be revoked on the server and may work until it expires.", + CredentialKind.SESSION: "The saved session could not be revoked on the server and may work until it expires.", +} + + @dataclasses.dataclass(frozen=True) class Credential: kind: CredentialKind @@ -200,13 +206,16 @@ def logout() -> Tuple[CredentialKind, ...]: Returns the kinds whose revoke failed, which may still work until they expire. """ with cli_login._locked(): - login = cli_login._load(cli_login._origin(), cached=False) + origin = cli_login._origin() + metadata = cli_login._read_metadata() + login = cli_login._load(origin, cached=False) session = cli_login._saved_session() - if login is not None: + saved_here = isinstance(metadata, dict) and metadata.get("url") == origin + if saved_here: cli_login._forget() Auth.secrets_file.unlink(missing_ok=True) failed = [] - if login is not None and not cli_login._revoke(login.refresh_token): + if saved_here and (login is None or not cli_login._revoke(login.refresh_token)): failed.append(CredentialKind.LOGIN) if session is not None and not cli_login._revoke_session(session): failed.append(CredentialKind.SESSION) diff --git a/python/tests/lightning_cloud/test_cli_login.py b/python/tests/lightning_cloud/test_cli_login.py index d0f3ca553..341f97685 100644 --- a/python/tests/lightning_cloud/test_cli_login.py +++ b/python/tests/lightning_cloud/test_cli_login.py @@ -344,6 +344,18 @@ def test_logout_that_cannot_reach_the_server_still_forgets_locally(platform, bro _nothing_saved(platform.url) +def test_logout_forgets_a_login_whose_keychain_cannot_be_read(platform, browser): + cli_login.login() + path = pathlib.Path(env.LIGHTNING_LOGIN_PATH) + assert path.is_file() + keyring.set_keyring(fail.Keyring()) + + assert logout() == (CredentialKind.LOGIN,) + + assert not path.exists() + assert platform.calls_to("/oidc/revoke") == [] + + def test_logout_with_nothing_saved_revokes_nothing(platform): assert logout() == () assert platform.requests == [] diff --git a/python/tests/lightning_cloud/test_replaced_login.py b/python/tests/lightning_cloud/test_replaced_login.py index 30043d940..28b89bca6 100644 --- a/python/tests/lightning_cloud/test_replaced_login.py +++ b/python/tests/lightning_cloud/test_replaced_login.py @@ -77,3 +77,29 @@ def test_a_login_that_fails_leaves_the_old_login_saved_and_live(platform, monkey assert platform.accepts(old.token) assert platform.calls_to("/oidc/revoke") == [] assert platform.logouts() == [] + + +def test_replacing_a_login_says_when_that_revoke_fails(platform, monkeypatch, capsys): + old = _log_in(platform, monkeypatch, "oidc", "user-a") + monkeypatch.setattr(platform, "revoke", lambda *args: (503, {})) + + _log_in(platform, monkeypatch, "oidc", "user-b") + + assert platform.family_active(old.refresh_token) + assert ( + "The saved login could not be revoked on the server and may work until it expires." + in capsys.readouterr().err + ) + + +def test_replacing_a_session_says_when_that_revoke_fails(platform, monkeypatch, capsys): + session = platform.mint_session("user-a") + Auth.secrets_file.write_text(json.dumps({"auth_token": session})) + monkeypatch.setattr(platform, "logout", lambda *args: (503, {})) + + _log_in(platform, monkeypatch, "oidc", "user-b") + + assert platform.accepts(session) + err = capsys.readouterr().err + assert "The saved session could not be revoked on the server and may work until it expires." in err + assert "The saved login could not be revoked" not in err From 7a741aba49e746358daf6ddfbc31b6900b79a0f1 Mon Sep 17 00:00:00 2001 From: fnwoke-lightning Date: Wed, 4 Nov 2026 22:35:21 -0500 Subject: [PATCH 13/27] fix(cli): drop LIGHTNING_API_KEY handling Co-authored-by: Cursor --- python/lightning_sdk/cli/auth/status.py | 2 - python/lightning_sdk/cli/auth/token.py | 2 - python/lightning_sdk/cli/utils/auth.py | 25 ++---- python/lightning_sdk/lightning_cloud/login.py | 7 +- python/tests/api/test_credential_requests.py | 45 +--------- python/tests/api/test_library_login.py | 34 -------- python/tests/cli/test_login_commands.py | 87 +++++-------------- .../tests/lightning_cloud/test_credentials.py | 5 +- 8 files changed, 34 insertions(+), 173 deletions(-) diff --git a/python/lightning_sdk/cli/auth/status.py b/python/lightning_sdk/cli/auth/status.py index 74cf510c8..0fe61b066 100644 --- a/python/lightning_sdk/cli/auth/status.py +++ b/python/lightning_sdk/cli/auth/status.py @@ -78,8 +78,6 @@ def _source(credential: Credential) -> str: login = cli_login.current_login() storage = login.storage if login is not None else "login" return f"login ({storage})" - if credential.kind is CredentialKind.LEGACY_KEY: - return "LIGHTNING_API_KEY (legacy key, until the new login is switched on)" return "interim session login" diff --git a/python/lightning_sdk/cli/auth/token.py b/python/lightning_sdk/cli/auth/token.py index a18781ea9..c0c87fbc7 100644 --- a/python/lightning_sdk/cli/auth/token.py +++ b/python/lightning_sdk/cli/auth/token.py @@ -28,8 +28,6 @@ def token() -> None: credential = CredentialSource().resolve() if credential is None: raise click.ClickException(_NOT_LOGGED_IN) - if credential.kind is CredentialKind.LEGACY_KEY: - raise click.ClickException("A legacy key is never printed. Run `lightning login`.") if credential.kind is CredentialKind.LOGIN: try: click.echo(cli_login.token()) diff --git a/python/lightning_sdk/cli/utils/auth.py b/python/lightning_sdk/cli/utils/auth.py index 8fa459843..624d16aed 100644 --- a/python/lightning_sdk/cli/utils/auth.py +++ b/python/lightning_sdk/cli/utils/auth.py @@ -9,7 +9,6 @@ from lightning_sdk.lightning_cloud import cli_login, env from lightning_sdk.lightning_cloud.credentials import ( CredentialKind, - CredentialSource, LoginUnavailable, logout, start_login, @@ -91,19 +90,10 @@ def run_login( "Unset it and run `lightning login` again." ) if _env_set("LIGHTNING_API_KEY"): - if not cli_login.oidc_login_enabled(): - # A saved login outranks the key, so a stale one can only be replaced by logging in. - sent = CredentialSource().resolve() - if sent is not None and sent.kind is CredentialKind.LEGACY_KEY: - click.echo( - '"lightning login" is not required within a Studio or when credentials are set in the environment' - ) - return - else: - click.echo( - "LIGHTNING_API_KEY is ignored until service accounts. This login continues without it.", - err=True, - ) + click.echo( + "LIGHTNING_API_KEY is ignored until service accounts. This login continues without it.", + err=True, + ) if in_studio(): raise RuntimeError("Unable to login within a Studio. Did you change your shell setup?") from None @@ -112,11 +102,14 @@ def run_login( if insecure_storage and cli_login.oidc_login_enabled(): click.echo(f"This login is kept in a file at {env.LIGHTNING_LOGIN_PATH}.", err=True) try: + if (asked_orgs or asked_project) and not cli_login.oidc_login_enabled(): + flags = [flag for flag, given in (("--org", asked_orgs), ("--project", asked_project)) if given] + raise LoginUnavailable( + f"{', '.join(flags)} needs the new login. Set {cli_login.OIDC_LOGIN_ENV}=1 to use it." + ) done = start_login( no_browser=no_browser, insecure_storage=insecure_storage, - orgs=asked_orgs, - project=asked_project, ) except LoginUnavailable as err: raise click.UsageError(str(err)) from err diff --git a/python/lightning_sdk/lightning_cloud/login.py b/python/lightning_sdk/lightning_cloud/login.py index 4fd217ca2..6cdd364d8 100644 --- a/python/lightning_sdk/lightning_cloud/login.py +++ b/python/lightning_sdk/lightning_cloud/login.py @@ -73,16 +73,11 @@ def load(self) -> bool: def save( self, - token: str = "", user_id: str = "", - api_key: str = "", auth_token: str = "", # username: str = "", ) -> None: - """Replace the saved login. ``token`` and ``api_key`` are ignored. - - The CLI on main still calls ``save(token=..., api_key=...)``. - """ + """Replace the saved login.""" stored = { Keys.USER_ID.suffix: user_id, Keys.AUTH_TOKEN.suffix: auth_token, diff --git a/python/tests/api/test_credential_requests.py b/python/tests/api/test_credential_requests.py index 0b18b74f2..a44f5a4c6 100644 --- a/python/tests/api/test_credential_requests.py +++ b/python/tests/api/test_credential_requests.py @@ -1,19 +1,17 @@ """Raw requests take a credential handler, and a 401 is retried once with a fresh token.""" -import base64 import threading import webbrowser from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer import keyring import pytest -import requests from lightning_sdk.api.billing_api import BillingApi from lightning_sdk.api.filesystem_api import FilesystemApi from lightning_sdk.api.studio_api import StudioApi from lightning_sdk.api.teamspace_api import TeamspaceApi -from lightning_sdk.api.utils import CredentialAuth, _BlobUploader +from lightning_sdk.api.utils import _BlobUploader from lightning_sdk.lightning_cloud import env from lightning_sdk.lightning_cloud.credentials import CredentialSource, start_login from lightning_sdk.lightning_cloud.login import Auth @@ -109,47 +107,6 @@ def test_billing_csv_retries_a_401(refreshed_token): assert refreshed_token.seen == ["old-token", "new-token"] -def test_a_401_on_a_basic_legacy_key_is_not_retried(monkeypatch): - threading.Thread = _REAL_THREAD - for name in ("LIGHTNING_AUTH_TOKEN", "LIGHTNING_AUTH_TOKEN_FILE", "LIGHTNING_CLI_OIDC_LOGIN"): - monkeypatch.delenv(name, raising=False) - monkeypatch.setenv("LIGHTNING_USER_ID", "user-1") - monkeypatch.setenv("LIGHTNING_API_KEY", "key-1") - seen = [] - expected = f"Basic {base64.b64encode(b'user-1:key-1').decode('ascii')}" - - class Handler(BaseHTTPRequestHandler): - protocol_version = "HTTP/1.1" - - def do_GET(self): # noqa: N802 - seen.append(self.headers.get("Authorization")) - body = b'{"error":"unauthenticated"}' - self.send_response(401) - self.send_header("Content-Length", str(len(body))) - self.end_headers() - self.wfile.write(body) - - def log_message(self, fmt, *args): - return - - httpd = ThreadingHTTPServer(("127.0.0.1", 0), Handler) - thread = _REAL_THREAD(target=httpd.serve_forever, daemon=True) - thread.start() - try: - response = requests.get( - f"http://127.0.0.1:{httpd.server_address[1]}/v1/auth/user", - auth=CredentialAuth(), - timeout=10, - ) - finally: - httpd.shutdown() - thread.join(timeout=5) - httpd.server_close() - - assert response.status_code == 401 - assert seen == [expected] - - def test_blob_upload_retries_a_401(refreshed_token, tmp_path): local = tmp_path / "blob.bin" local.write_bytes(b"hi") diff --git a/python/tests/api/test_library_login.py b/python/tests/api/test_library_login.py index 5fad87f38..e7e048821 100644 --- a/python/tests/api/test_library_login.py +++ b/python/tests/api/test_library_login.py @@ -69,40 +69,6 @@ def test_litai_puts_okbhrt_and_litaimodel_on_the_interim_sign_in(monkeypatch, tm platform.stop() -def test_litai_treats_a_legacy_key_as_an_existing_credential(monkeypatch): - monkeypatch.delenv("LIGHTNING_AUTH_TOKEN", raising=False) - monkeypatch.delenv("LIGHTNING_AUTH_TOKEN_FILE", raising=False) - monkeypatch.delenv("LIGHTNING_CLI_OIDC_LOGIN", raising=False) - monkeypatch.setenv("LIGHTNING_USER_ID", "user-1") - monkeypatch.setenv("LIGHTNING_API_KEY", "key-1") - asked = [] - monkeypatch.setattr( - "lightning_sdk.api.llm_api.Confirm.ask", - lambda *_args, **_kwargs: asked.append(True) or True, - ) - opened = [] - monkeypatch.setattr(webbrowser, "open", lambda *_args, **_kwargs: opened.append(True)) - - litai_authenticate("story-model") - - assert asked == [] - assert opened == [] - - -def test_litserve_treats_a_legacy_key_as_an_existing_credential(monkeypatch): - monkeypatch.delenv("LIGHTNING_AUTH_TOKEN", raising=False) - monkeypatch.delenv("LIGHTNING_AUTH_TOKEN_FILE", raising=False) - monkeypatch.delenv("LIGHTNING_CLI_OIDC_LOGIN", raising=False) - monkeypatch.setenv("LIGHTNING_USER_ID", "user-1") - monkeypatch.setenv("LIGHTNING_API_KEY", "key-1") - opened = [] - monkeypatch.setattr(webbrowser, "open", lambda *_args, **_kwargs: opened.append(True)) - - litserve_authenticate(_AuthMode.DEPLOY) - - assert opened == [] - - def test_litai_does_nothing_when_a_credential_already_exists(monkeypatch): monkeypatch.setenv("LIGHTNING_AUTH_TOKEN", "env-token") asked = [] diff --git a/python/tests/cli/test_login_commands.py b/python/tests/cli/test_login_commands.py index 873d42a31..6eb0d293b 100644 --- a/python/tests/cli/test_login_commands.py +++ b/python/tests/cli/test_login_commands.py @@ -2,7 +2,6 @@ import getpass import json -import re import sys import threading import webbrowser @@ -29,10 +28,6 @@ def _text(result) -> str: return result.output + (getattr(result, "stderr", None) or "") -def _plain(result) -> str: - return " ".join(re.sub(r"\x1b\[[0-9;]*m", "", _text(result)).split()) - - @pytest.fixture() def isolated(monkeypatch, tmp_path): """A private credential store, no ambient token, and real threads for the fake servers.""" @@ -133,53 +128,34 @@ def test_login_refuses_a_token_file_and_an_environment_token(isolated, monkeypat assert not Auth.secrets_file.exists() -_LOGIN_NOT_REQUIRED = '"lightning login" is not required within a Studio or when credentials are set in the environment' - - @pytest.mark.parametrize("cli", [main_cli, legacy_cli]) -def test_with_the_switch_unset_login_says_an_api_key_needs_no_login(platform, browser, monkeypatch, cli) -> None: - monkeypatch.setenv("LIGHTNING_USER_ID", "user-1") - monkeypatch.setenv("LIGHTNING_API_KEY", "key-1") +@pytest.mark.parametrize("switch", ["interim", "oidc"]) +def test_login_warns_that_an_api_key_is_ignored_and_continues(platform, browser, monkeypatch, cli, switch) -> None: + if switch == "oidc": + monkeypatch.setenv("LIGHTNING_CLI_OIDC_LOGIN", "1") + monkeypatch.setenv("LIGHTNING_API_KEY", "lit_sa_ignored") result = _invoke(cli, ["login"]) assert result.exit_code == 0, _text(result) - assert result.stdout == _LOGIN_NOT_REQUIRED + "\n" - assert "Successfully logged in" not in _text(result) - assert not platform.sign_ins - assert not platform.authorizations - assert not Auth.secrets_file.exists() - - -def test_with_the_switch_unset_login_replaces_a_saved_session_that_outranks_an_api_key( - platform, browser, monkeypatch -) -> None: - monkeypatch.setenv("LIGHTNING_USER_ID", "user-1") - monkeypatch.setenv("LIGHTNING_API_KEY", "key-1") - previous = platform.mint_session("user-1") - Auth().save(auth_token=previous) - - result = _invoke(main_cli, ["login"]) - - assert result.exit_code == 0, _text(result) - assert _LOGIN_NOT_REQUIRED not in _text(result) - assert platform.sign_ins - assert previous in platform.logouts() - assert _saved_session() not in ("", previous) + assert "LIGHTNING_API_KEY is ignored until service accounts. This login continues without it." in result.stderr + assert "lit_sa_ignored" not in _text(result) + if switch == "oidc": + assert CredentialSource().resolve().token.startswith("lit_ua_") + assert platform.authorizations + else: + assert _saved_session().startswith("sess_") + assert platform.sign_ins @pytest.mark.parametrize("cli", [main_cli, legacy_cli]) -def test_with_the_switch_set_login_warns_that_an_api_key_is_ignored_and_continues( - platform, browser, monkeypatch, cli -) -> None: - monkeypatch.setenv("LIGHTNING_CLI_OIDC_LOGIN", "1") +def test_login_no_browser_with_an_api_key_still_names_the_switch(isolated, monkeypatch, cli) -> None: monkeypatch.setenv("LIGHTNING_API_KEY", "lit_sa_ignored") - result = _invoke(cli, ["login"]) + result = _invoke(cli, ["login", "--no-browser"]) - assert result.exit_code == 0, _text(result) - assert "LIGHTNING_API_KEY is ignored until service accounts" in result.stderr - assert CredentialSource().resolve().token.startswith("lit_ua_") + assert result.exit_code == 2, _text(result) + assert "LIGHTNING_CLI_OIDC_LOGIN" in _text(result) assert "lit_sa_ignored" not in _text(result) @@ -336,21 +312,9 @@ def test_auth_status_names_each_source_and_never_a_token(platform, browser, monk monkeypatch.setenv("LIGHTNING_API_KEY", "key-1") monkeypatch.delenv("LIGHTNING_CLI_OIDC_LOGIN", raising=False) legacy = _invoke(main_cli, ["auth", "status"]) - shown = _plain(legacy) - assert legacy.exit_code == 0, _text(legacy) - assert "LIGHTNING_API_KEY (legacy key, until the new login is switched on)" in shown - assert "Orgs" in shown - assert "Teamspaces" in shown - assert "Permissions" in shown - assert "Expiry" in shown - assert shown.count("not known locally") == 4 - assert "key-1" not in shown - - monkeypatch.setenv("LIGHTNING_CLI_OIDC_LOGIN", "1") - hidden = _invoke(main_cli, ["auth", "status"]) - assert hidden.exit_code != 0 - assert _NOT_LOGGED_IN in _text(hidden) - assert "key-1" not in _text(hidden) + assert legacy.exit_code != 0 + assert _NOT_LOGGED_IN in _text(legacy) + assert "key-1" not in _text(legacy) def test_auth_token_prints_what_requests_send_and_refuses_a_short_login(platform, monkeypatch) -> None: @@ -362,23 +326,14 @@ def test_auth_token_prints_what_requests_send_and_refuses_a_short_login(platform assert "what requests send" in printed.stderr monkeypatch.delenv("LIGHTNING_AUTH_TOKEN") - monkeypatch.delenv("LIGHTNING_CLI_OIDC_LOGIN", raising=False) monkeypatch.setenv("LIGHTNING_USER_ID", "user-1") monkeypatch.setenv("LIGHTNING_API_KEY", "lit_sa_secret") refused = _invoke(main_cli, ["auth", "token"]) assert refused.exit_code != 0 assert refused.stdout == "" assert "lit_sa_secret" not in _text(refused) - assert "A legacy key is never printed. Run `lightning login`." in _text(refused) - - monkeypatch.setenv("LIGHTNING_CLI_OIDC_LOGIN", "1") - hidden = _invoke(main_cli, ["auth", "token"]) - assert hidden.exit_code != 0 - assert hidden.stdout == "" - assert "lit_sa_secret" not in _text(hidden) - assert _NOT_LOGGED_IN in _text(hidden) + assert _NOT_LOGGED_IN in _text(refused) - monkeypatch.delenv("LIGHTNING_CLI_OIDC_LOGIN") monkeypatch.delenv("LIGHTNING_API_KEY") monkeypatch.delenv("LIGHTNING_USER_ID") Auth().save(auth_token="sess_saved") diff --git a/python/tests/lightning_cloud/test_credentials.py b/python/tests/lightning_cloud/test_credentials.py index b15510c7b..cf33d9760 100644 --- a/python/tests/lightning_cloud/test_credentials.py +++ b/python/tests/lightning_cloud/test_credentials.py @@ -608,11 +608,10 @@ def request(method, url, headers=None, **kwargs): def test_save_replaces_the_saved_login(isolated_credentials): path = isolated_credentials auth = Auth() - auth.save(user_id="user-1", api_key="key-1") - auth.save(token="ignored", auth_token="sess_1") + auth.save(user_id="user-1") + auth.save(auth_token="sess_1") assert json.loads(path.read_text()) == {"user_id": "", "auth_token": "sess_1"} assert "api_key" not in path.read_text() - assert "key-1" not in path.read_text() assert not auth.user_id assert auth.auth_token == "sess_1" From d97fd3f6d99a0fbc8e5e8152e3b8cf2e4a580857 Mon Sep 17 00:00:00 2001 From: fnwoke-lightning Date: Wed, 4 Nov 2026 22:40:41 -0500 Subject: [PATCH 14/27] fix(cli): leave --org and --project to M4-8b Co-authored-by: Cursor --- python/lightning_sdk/cli/entrypoint.py | 6 +-- python/lightning_sdk/cli/legacy/entrypoint.py | 6 +-- python/lightning_sdk/cli/utils/auth.py | 49 +------------------ python/tests/cli/test_login_commands.py | 44 +---------------- 4 files changed, 7 insertions(+), 98 deletions(-) diff --git a/python/lightning_sdk/cli/entrypoint.py b/python/lightning_sdk/cli/entrypoint.py index 42d432f3a..c5f53f722 100644 --- a/python/lightning_sdk/cli/entrypoint.py +++ b/python/lightning_sdk/cli/entrypoint.py @@ -103,11 +103,9 @@ def main_cli(ctx: click.Context) -> None: is_flag=True, help="Keep the login in a file instead of the OS keychain.", ) -@click.option("--org", default=None, help="Comma-separated orgs this login should be limited to.") -@click.option("--project", default=None, help="Teamspace this login should be limited to.") -def login(no_browser: bool, insecure_storage: bool, org: str, project: str) -> None: +def login(no_browser: bool, insecure_storage: bool) -> None: """Sign in to Lightning AI.""" - run_login(no_browser=no_browser, insecure_storage=insecure_storage, orgs=org, project=project) + run_login(no_browser=no_browser, insecure_storage=insecure_storage) @main_cli.command diff --git a/python/lightning_sdk/cli/legacy/entrypoint.py b/python/lightning_sdk/cli/legacy/entrypoint.py index c85a842a5..7021e3dfd 100644 --- a/python/lightning_sdk/cli/legacy/entrypoint.py +++ b/python/lightning_sdk/cli/legacy/entrypoint.py @@ -79,11 +79,9 @@ def main_cli(ctx: click.Context) -> None: is_flag=True, help="Keep the login in a file instead of the OS keychain.", ) -@click.option("--org", default=None, help="Comma-separated orgs this login should be limited to.") -@click.option("--project", default=None, help="Teamspace this login should be limited to.") -def login(no_browser: bool, insecure_storage: bool, org: str, project: str) -> None: +def login(no_browser: bool, insecure_storage: bool) -> None: """Login to Lightning AI Studios.""" - run_login(no_browser=no_browser, insecure_storage=insecure_storage, orgs=org, project=project) + run_login(no_browser=no_browser, insecure_storage=insecure_storage) @main_cli.command diff --git a/python/lightning_sdk/cli/utils/auth.py b/python/lightning_sdk/cli/utils/auth.py index 624d16aed..df5a535bb 100644 --- a/python/lightning_sdk/cli/utils/auth.py +++ b/python/lightning_sdk/cli/utils/auth.py @@ -1,8 +1,6 @@ """CLI authentication commands, on top of the vendored credential source.""" -import json import os -from typing import Any, Dict, Optional, Sequence import click @@ -40,43 +38,9 @@ def _env_set(name: str) -> bool: return bool(os.environ.get(name, "").strip()) -def _org_list(orgs: Optional[str]) -> Optional[Sequence[str]]: - if not orgs: - return None - names = [part.strip() for part in orgs.split(",") if part.strip()] - return names or None - - -def _grant_reaches_further(grant: Dict[str, Any], orgs: Optional[Sequence[str]], project: Optional[str]) -> bool: - """Whether ``grant`` covers more than the orgs and project the user asked for.""" - scopes = grant.get("orgs") - asked = {org for org in orgs or [] if org} - if scopes == "all": - return bool(asked or project) - if not isinstance(scopes, list): - return False - for scope in scopes: - if not isinstance(scope, dict): - continue - org = str(scope.get("org", "")) - projects = scope.get("projects") - if asked and org not in asked: - return True - if not project: - continue - if not isinstance(projects, list) or not projects: - return True - named = [str(item) for item in projects] - if project not in named or len(named) > 1: - return True - return False - - def run_login( no_browser: bool = False, insecure_storage: bool = False, - orgs: Optional[str] = None, - project: Optional[str] = None, ) -> None: """Run the login the switch selects, without dropping the saved login first.""" if _env_set(env.LIGHTNING_AUTH_TOKEN_FILE): @@ -97,27 +61,16 @@ def run_login( if in_studio(): raise RuntimeError("Unable to login within a Studio. Did you change your shell setup?") from None - asked_orgs = _org_list(orgs) - asked_project = project.strip() if isinstance(project, str) and project.strip() else None if insecure_storage and cli_login.oidc_login_enabled(): click.echo(f"This login is kept in a file at {env.LIGHTNING_LOGIN_PATH}.", err=True) try: - if (asked_orgs or asked_project) and not cli_login.oidc_login_enabled(): - flags = [flag for flag, given in (("--org", asked_orgs), ("--project", asked_project)) if given] - raise LoginUnavailable( - f"{', '.join(flags)} needs the new login. Set {cli_login.OIDC_LOGIN_ENV}=1 to use it." - ) - done = start_login( + start_login( no_browser=no_browser, insecure_storage=insecure_storage, ) except LoginUnavailable as err: raise click.UsageError(str(err)) from err - wider = done is not None and (asked_orgs or asked_project) - if wider and _grant_reaches_further(done.grant, asked_orgs, asked_project): - grant = json.dumps(done.grant, sort_keys=True) - click.echo(f"This login reaches more than you asked for. The grant is {grant}.") try: click.echo(f'Successfully logged in as "{_get_authed_user().name}".') except Exception: diff --git a/python/tests/cli/test_login_commands.py b/python/tests/cli/test_login_commands.py index 6eb0d293b..b7daab052 100644 --- a/python/tests/cli/test_login_commands.py +++ b/python/tests/cli/test_login_commands.py @@ -1,8 +1,6 @@ """Login, logout, auth status and auth token, against the fake platform.""" -import getpass import json -import sys import threading import webbrowser from types import SimpleNamespace @@ -17,7 +15,7 @@ from lightning_sdk.lightning_cloud import env from lightning_sdk.lightning_cloud.credentials import CredentialSource, logout from lightning_sdk.lightning_cloud.login import Auth -from tests.lightning_cloud.utils.fake_platform import FakeBrowser, FakePlatform, OtherDevice +from tests.lightning_cloud.utils.fake_platform import FakeBrowser, FakePlatform from tests.lightning_cloud.utils.keyrings import MemoryKeyring _REAL_THREAD = threading.Thread @@ -101,7 +99,7 @@ def test_login_uses_the_interim_flow_when_the_switch_is_unset(platform, browser, @pytest.mark.parametrize("cli", [main_cli, legacy_cli]) @pytest.mark.parametrize( "args", - [["--no-browser"], ["--insecure-storage"], ["--org", "acme"], ["--project", "research"]], + [["--no-browser"], ["--insecure-storage"]], ) def test_interim_login_refuses_flags_the_new_login_needs(isolated, cli, args) -> None: result = _invoke(cli, ["login", *args]) @@ -194,44 +192,6 @@ def test_insecure_storage_warns_and_names_the_login_file(platform, browser, monk assert CredentialSource().resolve().kind.value == "login" -@pytest.mark.parametrize("args", [["--org", "acme"], ["--no-browser", "--org", "acme"]]) -def test_login_names_a_grant_that_reaches_further_than_the_hints(platform, browser, monkeypatch, args) -> None: - monkeypatch.setenv("LIGHTNING_CLI_OIDC_LOGIN", "1") - if "--no-browser" in args: - - def read_terminal(): - sys.stderr.flush() - captured = getattr(sys.stderr, "buffer", sys.stderr) - raw = captured.getvalue() if hasattr(captured, "getvalue") else "" - return raw.decode() if isinstance(raw, bytes) else raw - - device = OtherDevice(platform, read_terminal, ["CODE"]) - monkeypatch.setattr(getpass, "getpass", device.getpass) - - result = _invoke(main_cli, ["login", *args]) - - assert result.exit_code == 0, _text(result) - assert "reaches more than you asked for" in _text(result) - assert '"orgs": "all"' in _text(result) - - -def test_login_stays_quiet_when_the_grant_matches_the_hints(platform, browser, monkeypatch) -> None: - monkeypatch.setenv("LIGHTNING_CLI_OIDC_LOGIN", "1") - original = platform._issue - - def matching(family): - body = original(family) - body["lightning_grant"] = {"orgs": [{"org": "acme", "projects": ["research"]}], "permissions": ["read"]} - return body - - platform._issue = matching - - result = _invoke(main_cli, ["login", "--org", "acme", "--project", "research"]) - - assert result.exit_code == 0, _text(result) - assert "reaches more than you asked for" not in _text(result) - - def test_a_studio_without_credentials_still_refuses_to_log_in(monkeypatch) -> None: monkeypatch.delenv("LIGHTNING_AUTH_TOKEN", raising=False) monkeypatch.delenv("LIGHTNING_API_KEY", raising=False) From 2cd6a879e30122a4dd6a65acd6676b130d893683 Mon Sep 17 00:00:00 2001 From: fnwoke-lightning Date: Wed, 4 Nov 2026 22:47:32 -0500 Subject: [PATCH 15/27] fix(cli): status, token, logout and raw-request fixes from review Co-authored-by: Cursor --- python/lightning_sdk/api/utils.py | 5 + python/lightning_sdk/cli/auth/status.py | 21 ++-- python/lightning_sdk/cli/auth/token.py | 5 +- python/lightning_sdk/cli/utils/auth.py | 17 +-- python/tests/api/test_credential_requests.py | 24 ++++- python/tests/cli/test_login_commands.py | 102 ++++++++++++++++++ .../lightning_cloud/utils/fake_platform.py | 6 ++ 7 files changed, 159 insertions(+), 21 deletions(-) diff --git a/python/lightning_sdk/api/utils.py b/python/lightning_sdk/api/utils.py index 60fde2665..41b696f79 100644 --- a/python/lightning_sdk/api/utils.py +++ b/python/lightning_sdk/api/utils.py @@ -33,6 +33,7 @@ from requests.exceptions import HTTPError from tqdm.auto import tqdm +from lightning_sdk.lightning_cloud import env from lightning_sdk.lightning_cloud.credentials import CredentialSource from lightning_sdk.lightning_cloud.openapi import ( ModelsStoreCompleteMultiPartUploadBody, @@ -1281,6 +1282,10 @@ def __init__(self, source: Optional[CredentialSource] = None) -> None: def __call__(self, request: requests.PreparedRequest) -> requests.PreparedRequest: header = self._source.header() request.headers["Authorization"] = header + # requests writes its own agent before auth runs. + user_agent = request.headers.get("User-Agent") + if not user_agent or user_agent == requests.utils.default_user_agent(): + request.headers["User-Agent"] = env.USER_AGENT def handle(response: requests.Response, **kwargs: Any) -> requests.Response: if getattr(handle, "retried", False) or response.status_code != 401: diff --git a/python/lightning_sdk/cli/auth/status.py b/python/lightning_sdk/cli/auth/status.py index 0fe61b066..4667d6bd7 100644 --- a/python/lightning_sdk/cli/auth/status.py +++ b/python/lightning_sdk/cli/auth/status.py @@ -62,8 +62,15 @@ def _grant_lines(grant: Dict[str, Any], orgs: Dict[str, str], teamspaces: Dict[s if not isinstance(scope, dict): continue org_id = str(scope.get("org", "")) - shown_orgs.append(orgs.get(org_id, org_id)) - for project in scope.get("projects") or []: + org_name = orgs.get(org_id, org_id) + shown_orgs.append(org_name) + projects = scope.get("projects") + if projects == "all": + shown_teamspaces.append(f"all in {org_name}") + continue + if not isinstance(projects, list): + continue + for project in projects: project_id = str(project) shown_teamspaces.append(teamspaces.get(project_id, project_id)) return ", ".join(shown_orgs) or "none", ", ".join(shown_teamspaces) or "none" @@ -92,7 +99,10 @@ def _add_user(table: Table, user: Dict[str, Optional[str]]) -> None: @click.command("status", cls=LightningCommand) def status() -> None: """Show the credential requests send, who it belongs to, and what a saved login reaches.""" - credential = CredentialSource().resolve() + try: + credential = CredentialSource().resolve() + except (OSError, cli_login.LoginError) as err: + raise click.ClickException(str(err)) from err if credential is None: raise click.ClickException(_NOT_LOGGED_IN) @@ -111,10 +121,7 @@ def status() -> None: table.add_row("Permissions", _UNKNOWN_REACH) table.add_row("Expiry", _UNKNOWN_REACH) else: - if login.grant.get("orgs") == "all": - orgs, teamspaces = "all", "all" - else: - orgs, teamspaces = _grant_lines(login.grant, _org_names(client), _teamspace_names(client)) + orgs, teamspaces = _grant_lines(login.grant, _org_names(client), _teamspace_names(client)) permissions = ", ".join(login.grant.get("permissions") or []) or "none" table.add_row("Orgs", orgs) table.add_row("Teamspaces", teamspaces) diff --git a/python/lightning_sdk/cli/auth/token.py b/python/lightning_sdk/cli/auth/token.py index c0c87fbc7..3aa6583d9 100644 --- a/python/lightning_sdk/cli/auth/token.py +++ b/python/lightning_sdk/cli/auth/token.py @@ -25,7 +25,10 @@ def _note(credential: Credential) -> str: @click.command("token", cls=LightningCommand) def token() -> None: """Print the credential requests send. A saved login has at least 30 minutes left.""" - credential = CredentialSource().resolve() + try: + credential = CredentialSource().resolve() + except (OSError, cli_login.LoginError) as err: + raise click.ClickException(str(err)) from err if credential is None: raise click.ClickException(_NOT_LOGGED_IN) if credential.kind is CredentialKind.LOGIN: diff --git a/python/lightning_sdk/cli/utils/auth.py b/python/lightning_sdk/cli/utils/auth.py index df5a535bb..6f27a3fc2 100644 --- a/python/lightning_sdk/cli/utils/auth.py +++ b/python/lightning_sdk/cli/utils/auth.py @@ -6,7 +6,7 @@ from lightning_sdk.lightning_cloud import cli_login, env from lightning_sdk.lightning_cloud.credentials import ( - CredentialKind, + REVOKE_FAILED, LoginUnavailable, logout, start_login, @@ -17,11 +17,6 @@ __all__ = ["browser_authentication", "require_auth_header", "run_login", "run_logout"] -_REVOKE_FAILED = { - CredentialKind.LOGIN: ("The saved login could not be revoked on the server and may work until it expires."), - CredentialKind.SESSION: ("The saved session could not be revoked on the server and may work until it expires."), -} - def require_auth_header() -> str: """The credential requests send, or a usage error that tells the user to log in.""" @@ -82,9 +77,7 @@ def run_logout() -> None: failed = logout() click.echo("Logged out on this machine.") for kind in failed: - click.echo( - _REVOKE_FAILED.get( - kind, - f"The {kind.value} credential could not be revoked on the server and may work until it expires.", - ) - ) + click.echo(REVOKE_FAILED[kind]) + for name in (env.LIGHTNING_AUTH_TOKEN_FILE, "LIGHTNING_AUTH_TOKEN"): + if _env_set(name): + click.echo(f"Requests still authenticate with {name}. Unset it.") diff --git a/python/tests/api/test_credential_requests.py b/python/tests/api/test_credential_requests.py index a44f5a4c6..5d19d64a9 100644 --- a/python/tests/api/test_credential_requests.py +++ b/python/tests/api/test_credential_requests.py @@ -6,12 +6,13 @@ import keyring import pytest +import requests from lightning_sdk.api.billing_api import BillingApi from lightning_sdk.api.filesystem_api import FilesystemApi from lightning_sdk.api.studio_api import StudioApi from lightning_sdk.api.teamspace_api import TeamspaceApi -from lightning_sdk.api.utils import _BlobUploader +from lightning_sdk.api.utils import CredentialAuth, _BlobUploader from lightning_sdk.lightning_cloud import env from lightning_sdk.lightning_cloud.credentials import CredentialSource, start_login from lightning_sdk.lightning_cloud.login import Auth @@ -26,6 +27,7 @@ class _RetryServer: def __init__(self, token_path): seen = [] + user_agents = [] class Handler(BaseHTTPRequestHandler): protocol_version = "HTTP/1.1" @@ -42,6 +44,7 @@ def do_POST(self): # noqa: N802 def _reply(self): token = self.headers.get("Authorization", "").removeprefix("Bearer ").strip() seen.append(token) + user_agents.append(self.headers.get("User-Agent")) if token == "old-token": token_path.write_text("new-token") status, body = 401, b'{"error":"expired"}' @@ -61,6 +64,7 @@ def log_message(self, fmt, *args): return self.seen = seen + self.user_agents = user_agents self._httpd = ThreadingHTTPServer(("127.0.0.1", 0), Handler) self.url = f"http://127.0.0.1:{self._httpd.server_address[1]}" self._thread = _REAL_THREAD(target=self._httpd.serve_forever, daemon=True) @@ -87,6 +91,24 @@ def refreshed_token(monkeypatch, tmp_path): server.stop() +def test_a_raw_request_sends_the_sdk_user_agent(refreshed_token): + FilesystemApi().list_files("ts", "data") + assert refreshed_token.user_agents == [env.USER_AGENT, env.USER_AGENT] + + requests.get( + f"{refreshed_token.url}/custom", + headers={"User-Agent": "caller-agent"}, + auth=CredentialAuth(), + timeout=5, + ) + assert refreshed_token.user_agents[-1] == "caller-agent" + + prepared = requests.Request("GET", f"{refreshed_token.url}/bare").prepare() + assert "User-Agent" not in prepared.headers + CredentialAuth()(prepared) + assert prepared.headers["User-Agent"] == env.USER_AGENT + + def test_filesystem_list_retries_a_401(refreshed_token): FilesystemApi().list_files("ts", "data") assert refreshed_token.seen == ["old-token", "new-token"] diff --git a/python/tests/cli/test_login_commands.py b/python/tests/cli/test_login_commands.py index b7daab052..ec6dadeb6 100644 --- a/python/tests/cli/test_login_commands.py +++ b/python/tests/cli/test_login_commands.py @@ -225,6 +225,108 @@ def test_logout_revokes_both_kinds_and_names_the_one_the_server_refused(platform assert CredentialSource().resolve() is None +@pytest.mark.parametrize("cli", [main_cli, legacy_cli]) +def test_logout_says_when_an_environment_credential_still_wins(isolated, monkeypatch, cli) -> None: + monkeypatch.setenv("LIGHTNING_AUTH_TOKEN", "still-sent") + monkeypatch.setenv("LIGHTNING_AUTH_TOKEN_FILE", str(isolated / "token")) + + result = _invoke(cli, ["logout"]) + + text = _text(result) + assert result.exit_code == 0, text + assert "Requests still authenticate with LIGHTNING_AUTH_TOKEN_FILE. Unset it." in text + assert "Requests still authenticate with LIGHTNING_AUTH_TOKEN. Unset it." in text + assert "still-sent" not in text + + monkeypatch.setenv("LIGHTNING_AUTH_TOKEN", " ") + monkeypatch.setenv("LIGHTNING_AUTH_TOKEN_FILE", "") + quiet = _invoke(cli, ["logout"]) + assert "still authenticate" not in _text(quiet) + + +def _login_with_grant(platform, monkeypatch, grant) -> None: + monkeypatch.setenv("LIGHTNING_CLI_OIDC_LOGIN", "1") + original = platform._issue + + def issue(family): + body = original(family) + body["lightning_grant"] = grant + return body + + platform._issue = issue + assert _invoke(main_cli, ["login"]).exit_code == 0 + + +def test_auth_status_shows_every_teamspace_in_an_org(platform, browser, monkeypatch) -> None: + platform.orgs = [{"id": "org-1", "name": "Acme"}] + _login_with_grant( + platform, + monkeypatch, + {"orgs": [{"org": "org-1", "projects": "all"}], "permissions": ["read"]}, + ) + + text = _text(_invoke(main_cli, ["auth", "status"])) + + assert "all in Acme" in text + assert "a, l, l" not in text + + +def test_auth_status_shows_a_mixed_grant(platform, browser, monkeypatch) -> None: + platform.orgs = [{"id": "org-1", "name": "Acme"}] + platform.memberships = [{"projectId": "proj-9", "name": "research"}] + _login_with_grant( + platform, + monkeypatch, + { + "orgs": [ + {"org": "org-1", "projects": ["proj-9"]}, + {"org": "org-missing", "projects": "all"}, + ], + "permissions": ["read"], + }, + ) + + text = _text(_invoke(main_cli, ["auth", "status"])) + + assert "Acme" in text + assert "research" in text + assert "all in org-missing" in text + assert "a, l, l" not in text + + +def test_an_unreadable_token_file_is_an_error_not_a_traceback(isolated, monkeypatch, tmp_path) -> None: + directory = tmp_path / "not-a-file" + directory.mkdir() + monkeypatch.setenv("LIGHTNING_AUTH_TOKEN_FILE", str(directory)) + + for command in (["auth", "status"], ["auth", "token"]): + result = _invoke(main_cli, command) + text = _text(result) + assert result.exit_code != 0, text + assert "Is a directory" in text + assert "Traceback" not in text + + +def test_a_refused_login_is_an_error_not_a_traceback(platform, browser, monkeypatch) -> None: + platform.access_ttl = 60 + monkeypatch.setenv("LIGHTNING_CLI_OIDC_LOGIN", "1") + original = platform.token + + def refuse(form): + if form.get("grant_type") == "refresh_token": + return 400, {"error": "invalid_grant", "error_description": "invalid refresh token"} + return original(form) + + monkeypatch.setattr(platform, "token", refuse) + for command in (["auth", "status"], ["auth", "token"]): + assert _invoke(main_cli, ["login"]).exit_code == 0 + result = _invoke(main_cli, command) + text = _text(result) + assert result.exit_code != 0, text + assert "Your login has expired or was revoked. Run `lightning login`." in text + assert "Traceback" not in text + + def test_auth_status_names_each_source_and_never_a_token(platform, browser, monkeypatch, tmp_path) -> None: monkeypatch.setenv("LIGHTNING_CLI_OIDC_LOGIN", "1") assert _invoke(main_cli, ["login"]).exit_code == 0 diff --git a/python/tests/lightning_cloud/utils/fake_platform.py b/python/tests/lightning_cloud/utils/fake_platform.py index 1f9a3eb8e..02061ed6f 100644 --- a/python/tests/lightning_cloud/utils/fake_platform.py +++ b/python/tests/lightning_cloud/utils/fake_platform.py @@ -39,6 +39,8 @@ def __init__(self, access_ttl=3600): self.sign_ins = [] self.reuse_events = [] self.blocked = False + self.orgs = None + self.memberships = None self.token_keys = {} self._codes = {} self._families = {} @@ -278,6 +280,10 @@ def do_GET(self): # noqa: N802 self.send_response(302) self.send_header("Location", landing) self.end_headers() + elif parts.path == "/v1/orgs" and platform.orgs is not None: + self._json(200, {"organizations": platform.orgs}) + elif parts.path == "/v1/memberships" and platform.memberships is not None: + self._json(200, {"memberships": platform.memberships}) elif parts.path == "/v1/auth/user": bearer = self.headers.get("Authorization", "") token = bearer[len("Bearer ") :] if bearer.startswith("Bearer ") else "" From 87fe35e453032f954f62c601df2dabb85d1f143f Mon Sep 17 00:00:00 2001 From: fnwoke-lightning Date: Thu, 5 Nov 2026 09:42:15 -0500 Subject: [PATCH 16/27] fix(lightning_cloud): name an unreadable token file instead of raising OSError Co-authored-by: Cursor --- .../lightning_cloud/credentials.py | 5 +++++ .../tests/lightning_cloud/test_credentials.py | 19 +++++++++++++++++++ 2 files changed, 24 insertions(+) diff --git a/python/lightning_sdk/lightning_cloud/credentials.py b/python/lightning_sdk/lightning_cloud/credentials.py index e15285ab8..5042d89d3 100644 --- a/python/lightning_sdk/lightning_cloud/credentials.py +++ b/python/lightning_sdk/lightning_cloud/credentials.py @@ -78,6 +78,11 @@ def _from_token_file() -> Optional[str]: raw = pathlib.Path(path).read_text() except FileNotFoundError: return None + except OSError as err: + raise cli_login.LoginError( + f"{env.LIGHTNING_AUTH_TOKEN_FILE} names {path}, which can't be read " + f"({err.strerror or err}). Fix the file or unset {env.LIGHTNING_AUTH_TOKEN_FILE}." + ) from err return _present(raw) diff --git a/python/tests/lightning_cloud/test_credentials.py b/python/tests/lightning_cloud/test_credentials.py index b15510c7b..78a59a4e7 100644 --- a/python/tests/lightning_cloud/test_credentials.py +++ b/python/tests/lightning_cloud/test_credentials.py @@ -128,6 +128,25 @@ def request(method, url, headers=None, **kwargs): assert seen == ["Bearer sess_one", "Bearer sess_two"] +def test_an_unreadable_token_file_is_named_wherever_a_credential_is_read(isolated_credentials, monkeypatch, tmp_path): + isolated_credentials.write_text(json.dumps({"auth_token": "sess_saved"})) + monkeypatch.setenv(env.LIGHTNING_AUTH_TOKEN_FILE, str(tmp_path)) + named = re.escape(f"{env.LIGHTNING_AUTH_TOKEN_FILE} names {tmp_path}, which can't be read") + seen = [] + + def request(method, url, headers=None, **kwargs): + seen.append(headers) + return "ok" + + with pytest.raises(LoginError, match=named): + CredentialSource().resolve() + with pytest.raises(LoginError, match=named): + Auth().get_auth_header() + with pytest.raises(LoginError, match=named): + GridRestClient(api_client=_Api(request), with_auth=True).api_client.request("GET", "http://example") + assert seen == [] + + @_SWITCH def test_a_saved_login_is_bearer_not_basic_in_either_switch_state(platform, isolated_credentials, monkeypatch, switch): _set_switch(monkeypatch, switch) From c1a684c8e9174538a5f963267592a305b33ce1ce Mon Sep 17 00:00:00 2001 From: fnwoke-lightning Date: Thu, 5 Nov 2026 09:42:15 -0500 Subject: [PATCH 17/27] fix(lightning_cloud): trust SSL_CA_CERT on login, revoke and interim sign-in calls Also applies ruff-format to test_replaced_login.py, which failed pre-commit. Co-authored-by: Cursor --- .../lightning_cloud/cli_login.py | 2 ++ python/lightning_sdk/lightning_cloud/login.py | 1 + .../lightning_cloud/test_replaced_login.py | 32 ++++++++++++++++--- 3 files changed, 31 insertions(+), 4 deletions(-) diff --git a/python/lightning_sdk/lightning_cloud/cli_login.py b/python/lightning_sdk/lightning_cloud/cli_login.py index b06414b2f..b32bdf523 100644 --- a/python/lightning_sdk/lightning_cloud/cli_login.py +++ b/python/lightning_sdk/lightning_cloud/cli_login.py @@ -569,6 +569,7 @@ def _post(path: str, form: Dict[str, str]) -> requests.Response: data={"client_id": CLIENT_ID, **form}, headers={"User-Agent": env.USER_AGENT}, timeout=(env.LIGHTNING_CLOUD_CONNECT_TIMEOUT, env.LIGHTNING_CLOUD_READ_TIMEOUT), + verify=env.SSL_CA_CERT, ) except requests.RequestException as err: raise _Unreachable(f"Couldn't reach {origin}. Check your connection and try again.") from err @@ -610,6 +611,7 @@ def _revoke_session(token: str) -> bool: f"{_origin()}/v1/auth/logout", headers={"Authorization": f"Bearer {token}", "User-Agent": env.USER_AGENT}, timeout=(env.LIGHTNING_CLOUD_CONNECT_TIMEOUT, env.LIGHTNING_CLOUD_READ_TIMEOUT), + verify=env.SSL_CA_CERT, ) except requests.RequestException: return False diff --git a/python/lightning_sdk/lightning_cloud/login.py b/python/lightning_sdk/lightning_cloud/login.py index 4fd217ca2..b318037e0 100644 --- a/python/lightning_sdk/lightning_cloud/login.py +++ b/python/lightning_sdk/lightning_cloud/login.py @@ -404,6 +404,7 @@ def save() -> str: url, headers={"User-Agent": env.USER_AGENT}, timeout=env.LIGHTNING_CLOUD_CONNECT_TIMEOUT, + verify=env.SSL_CA_CERT, ) except requests.RequestException as err: raise cli_login.LoginError( diff --git a/python/tests/lightning_cloud/test_replaced_login.py b/python/tests/lightning_cloud/test_replaced_login.py index 28b89bca6..80e6d15df 100644 --- a/python/tests/lightning_cloud/test_replaced_login.py +++ b/python/tests/lightning_cloud/test_replaced_login.py @@ -1,12 +1,14 @@ import json import webbrowser +from urllib.parse import urlsplit import keyring import pytest +import requests -from lightning_sdk.lightning_cloud import cli_login +from lightning_sdk.lightning_cloud import cli_login, env from lightning_sdk.lightning_cloud.cli_login import LoginError -from lightning_sdk.lightning_cloud.credentials import start_login +from lightning_sdk.lightning_cloud.credentials import logout, start_login from lightning_sdk.lightning_cloud.login import Auth from lightning_sdk.lightning_cloud.rest_client import LightningClient from tests.lightning_cloud.utils.fake_platform import FakeBrowser @@ -87,8 +89,7 @@ def test_replacing_a_login_says_when_that_revoke_fails(platform, monkeypatch, ca assert platform.family_active(old.refresh_token) assert ( - "The saved login could not be revoked on the server and may work until it expires." - in capsys.readouterr().err + "The saved login could not be revoked on the server and may work until it expires." in capsys.readouterr().err ) @@ -103,3 +104,26 @@ def test_replacing_a_session_says_when_that_revoke_fails(platform, monkeypatch, err = capsys.readouterr().err assert "The saved session could not be revoked on the server and may work until it expires." in err assert "The saved login could not be revoked" not in err + + +def test_logins_revokes_and_logout_trust_the_configured_ca_bundle(platform, monkeypatch, tmp_path): + bundle = str(tmp_path / "ca.pem") + monkeypatch.setattr(env, "SSL_CA_CERT", bundle) + sent = [] + send = requests.adapters.HTTPAdapter.send + + def recording_send(adapter, request, **kwargs): + # The fake browser only sends GETs; the SDK's login calls are HEAD and POST. + if request.method in ("HEAD", "POST"): + sent.append((urlsplit(request.url).path, kwargs.get("verify"))) + return send(adapter, request, **kwargs) + + monkeypatch.setattr(requests.adapters.HTTPAdapter, "send", recording_send) + + _log_in(platform, monkeypatch, "interim", "user-a") + _log_in(platform, monkeypatch, "oidc", "user-b") + _log_in(platform, monkeypatch, "oidc", "user-c") + assert logout() == () + + assert {path for path, _ in sent} == {"/sign-in", "/oidc/token", "/oidc/revoke", "/v1/auth/logout"} + assert {verify for _, verify in sent} == {bundle} From 81155aad67c4616663976adfe4414913d24a2e25 Mon Sep 17 00:00:00 2001 From: fnwoke-lightning Date: Thu, 5 Nov 2026 10:02:42 -0500 Subject: [PATCH 18/27] fix(api): say to log in when a raw request's 401 can't be refreshed Co-authored-by: Cursor --- python/lightning_sdk/api/utils.py | 22 ++++-- python/tests/api/test_credential_requests.py | 73 +++++++++++++++----- 2 files changed, 73 insertions(+), 22 deletions(-) diff --git a/python/lightning_sdk/api/utils.py b/python/lightning_sdk/api/utils.py index 41b696f79..92fd2e185 100644 --- a/python/lightning_sdk/api/utils.py +++ b/python/lightning_sdk/api/utils.py @@ -45,7 +45,7 @@ ) from lightning_sdk.lightning_cloud.openapi.models.v1_model_version_archive import V1ModelVersionArchive from lightning_sdk.lightning_cloud.openapi.rest import ApiException -from lightning_sdk.lightning_cloud.rest_client import LightningClient +from lightning_sdk.lightning_cloud.rest_client import _LOGIN_MESSAGE, LightningClient from lightning_sdk.machine import Machine @@ -1274,7 +1274,11 @@ def to_iso_z(dt: datetime) -> str: class CredentialAuth(AuthBase): - """Sets ``Authorization`` when a request is prepared, and retries one 401 after a refresh.""" + """Sets ``Authorization`` when a request is prepared, and retries one 401 after a refresh. + + A 401 that no refresh fixes raises ``ConnectionError`` saying to run ``lightning login``, + as the generated client does. + """ def __init__(self, source: Optional[CredentialSource] = None) -> None: self._source = source or CredentialSource() @@ -1287,18 +1291,24 @@ def __call__(self, request: requests.PreparedRequest) -> requests.PreparedReques if not user_agent or user_agent == requests.utils.default_user_agent(): request.headers["User-Agent"] = env.USER_AGENT + retried = False + def handle(response: requests.Response, **kwargs: Any) -> requests.Response: - if getattr(handle, "retried", False) or response.status_code != 401: + nonlocal retried + if retried or response.status_code != 401: return response - if not self._source.refresh(header): - return response - handle.retried = True # The connection stays checked out until the body is read, so the retry can reuse it. _ = response.content response.close() + if not self._source.refresh(header): + raise ConnectionError(_LOGIN_MESSAGE) + retried = True prepared = response.request.copy() prepared.headers["Authorization"] = self._source.header() retry = response.connection.send(prepared, **kwargs) + if retry.status_code == 401: + retry.close() + raise ConnectionError(_LOGIN_MESSAGE) retry.history.append(response) retry.request = prepared return retry diff --git a/python/tests/api/test_credential_requests.py b/python/tests/api/test_credential_requests.py index 5d19d64a9..bf1b28943 100644 --- a/python/tests/api/test_credential_requests.py +++ b/python/tests/api/test_credential_requests.py @@ -1,5 +1,6 @@ """Raw requests take a credential handler, and a 401 is retried once with a fresh token.""" +import re import threading import webbrowser from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer @@ -20,12 +21,16 @@ from tests.lightning_cloud.utils.keyrings import MemoryKeyring _REAL_THREAD = threading.Thread +_LOGIN_FAILED = "Authentication failed. Please run `lightning login`." class _RetryServer: - """Answers 401 for ``old-token`` after rewriting the token file, then 200.""" + """Answers 401 for ``old-token`` after rewriting the token file, then 200. - def __init__(self, token_path): + With ``rotate=False`` the token file is left alone, so the 401 can't be refreshed. + """ + + def __init__(self, token_path, rotate=True): seen = [] user_agents = [] @@ -46,7 +51,8 @@ def _reply(self): seen.append(token) user_agents.append(self.headers.get("User-Agent")) if token == "old-token": - token_path.write_text("new-token") + if rotate: + token_path.write_text("new-token") status, body = 401, b'{"error":"expired"}' elif "usage-report" in self.path: status, body = 200, b"col\n" @@ -76,21 +82,65 @@ def stop(self): self._httpd.server_close() -@pytest.fixture() -def refreshed_token(monkeypatch, tmp_path): +def _serve_token_file(monkeypatch, tmp_path, rotate): threading.Thread = _REAL_THREAD for name in ("LIGHTNING_AUTH_TOKEN", "LIGHTNING_API_KEY"): monkeypatch.delenv(name, raising=False) token_path = tmp_path / "token" token_path.write_text("old-token") monkeypatch.setenv("LIGHTNING_AUTH_TOKEN_FILE", str(token_path)) - server = _RetryServer(token_path) + server = _RetryServer(token_path, rotate=rotate) monkeypatch.setenv("LIGHTNING_CLOUD_URL", server.url) monkeypatch.setattr(env, "LIGHTNING_CLOUD_URL", server.url) + return server + + +@pytest.fixture() +def refreshed_token(monkeypatch, tmp_path): + server = _serve_token_file(monkeypatch, tmp_path, rotate=True) + yield server + server.stop() + + +@pytest.fixture() +def stale_token(monkeypatch, tmp_path): + server = _serve_token_file(monkeypatch, tmp_path, rotate=False) yield server server.stop() +def _blob_request(url, tmp_path): + local = tmp_path / "blob.bin" + local.write_bytes(b"hi") + uploader = _BlobUploader( + client=None, + endpoint_base=f"{url}/v1/projects/ts/artifacts", + file_path=str(local), + remote_path="blob.bin", + progress_bar=False, + ) + uploader._post_blob_request("/blobs", {"path": "blob.bin"}, "request upload URLs for") + + +_RAW_PATHS = { + "filesystem list": lambda url, tmp_path: FilesystemApi().list_files("ts", "data"), + "filesystem download": lambda url, tmp_path: FilesystemApi()._download_single_file("a.bin", tmp_path / "a", "ts"), + "studio tree": lambda url, tmp_path: StudioApi().get_tree("st", "ts", "data"), + "teamspace tree": lambda url, tmp_path: TeamspaceApi().get_tree("ts", "data"), + "billing csv": lambda url, tmp_path: BillingApi()._download_activity_csv( + "/v1/billing/usage-report/download/detailed", "org" + ), + "blob upload": _blob_request, +} + + +@pytest.mark.parametrize("call", _RAW_PATHS.values(), ids=_RAW_PATHS.keys()) +def test_a_raw_request_refused_with_a_401_no_refresh_fixes_says_to_log_in(stale_token, tmp_path, call): + with pytest.raises(ConnectionError, match=re.escape(_LOGIN_FAILED)): + call(stale_token.url, tmp_path) + assert stale_token.seen == ["old-token"] + + def test_a_raw_request_sends_the_sdk_user_agent(refreshed_token): FilesystemApi().list_files("ts", "data") assert refreshed_token.user_agents == [env.USER_AGENT, env.USER_AGENT] @@ -130,16 +180,7 @@ def test_billing_csv_retries_a_401(refreshed_token): def test_blob_upload_retries_a_401(refreshed_token, tmp_path): - local = tmp_path / "blob.bin" - local.write_bytes(b"hi") - uploader = _BlobUploader( - client=None, - endpoint_base=f"{refreshed_token.url}/v1/projects/ts/artifacts", - file_path=str(local), - remote_path="blob.bin", - progress_bar=False, - ) - uploader._post_blob_request("/blobs", {"path": "blob.bin"}, "request upload URLs for") + _blob_request(refreshed_token.url, tmp_path) assert refreshed_token.seen == ["old-token", "new-token"] From 10c36b193e799da8dad1c9b7f4f37594be9b7be1 Mon Sep 17 00:00:00 2001 From: fnwoke-lightning Date: Thu, 5 Nov 2026 10:02:43 -0500 Subject: [PATCH 19/27] fix(api): stop following logs on a refused handshake instead of reconnecting forever Co-authored-by: Cursor --- python/lightning_sdk/api/logs_api.py | 15 +++++- python/tests/api/test_logs_api.py | 78 ++++++++++++++++++++++++++++ 2 files changed, 92 insertions(+), 1 deletion(-) diff --git a/python/lightning_sdk/api/logs_api.py b/python/lightning_sdk/api/logs_api.py index 322b21076..9581d6941 100644 --- a/python/lightning_sdk/api/logs_api.py +++ b/python/lightning_sdk/api/logs_api.py @@ -22,8 +22,9 @@ from urllib.parse import urlparse from lightning_sdk.api.utils import _get_cloud_url +from lightning_sdk.lightning_cloud.credentials import CredentialSource from lightning_sdk.lightning_cloud.login import Auth -from lightning_sdk.lightning_cloud.rest_client import LightningClient +from lightning_sdk.lightning_cloud.rest_client import _LOGIN_MESSAGE, LightningClient # While tailing, ``recv()`` is given this timeout so a quiet socket can be re-checked (the server # holds the connection open with heartbeats rather than closing it). It must stay below the @@ -497,6 +498,7 @@ def follow( try: import websocket from websocket import ( + WebSocketBadStatusException, WebSocketConnectionClosedException, WebSocketException, WebSocketTimeoutException, @@ -517,6 +519,7 @@ def follow( last_line_at = time.monotonic() backoff = _MIN_RECONNECT_BACKOFF + refreshed = False while True: ws = None try: @@ -530,6 +533,7 @@ def follow( on_socket(ws) # Connected cleanly: reset the backoff so the next blip recovers fast backoff = _MIN_RECONNECT_BACKOFF + refreshed = False while True: try: message = ws.recv() @@ -554,6 +558,15 @@ def follow( continue recent.add(entry) yield entry + except WebSocketBadStatusException as err: + # Reconnecting with a credential the server refused would loop forever. + if err.status_code == 401: + if refreshed or not CredentialSource().refresh(auth_header): + raise ConnectionError(_LOGIN_MESSAGE) from err + refreshed = True + continue + if err.status_code == 403 or not reconnect or idle_timeout is not None: + raise except (WebSocketException, OSError): if not reconnect or idle_timeout is not None: raise diff --git a/python/tests/api/test_logs_api.py b/python/tests/api/test_logs_api.py index 8316d97f1..f3328bece 100644 --- a/python/tests/api/test_logs_api.py +++ b/python/tests/api/test_logs_api.py @@ -1,3 +1,4 @@ +import re from datetime import datetime, timedelta, timezone from unittest import mock @@ -12,8 +13,11 @@ _websocket_url, parse_log_entries, ) +from lightning_sdk.lightning_cloud.login import Auth from lightning_sdk.lightning_cloud.openapi import V1GetLogsResponse, V1JobLogEntry +_LOGIN_FAILED = "Authentication failed. Please run `lightning login`." + def _api(*pages): """A LogsApi whose client returns ``pages`` in order.""" @@ -275,6 +279,12 @@ class _FakeClosedError(_FakeWebSocketError): pass +class _FakeBadStatusError(_FakeWebSocketError): + def __init__(self, status_code): + super().__init__(f"Handshake status {status_code}") + self.status_code = status_code + + @pytest.fixture() def fake_websocket(monkeypatch): """Stand in for the websocket-client module used by ``LogsApi.follow``.""" @@ -282,6 +292,7 @@ def fake_websocket(monkeypatch): module.WebSocketException = _FakeWebSocketError module.WebSocketTimeoutException = _FakeTimeoutError module.WebSocketConnectionClosedException = _FakeClosedError + module.WebSocketBadStatusException = _FakeBadStatusError monkeypatch.setitem(__import__("sys").modules, "websocket", module) monkeypatch.setattr("lightning_sdk.api.logs_api.Auth", mock.MagicMock()) return module @@ -377,6 +388,73 @@ def test_follow_reconnects_after_a_failed_handshake(fake_websocket) -> None: assert fake_websocket.create_connection.call_count == 3 +@pytest.fixture() +def token_file(fake_websocket, monkeypatch, tmp_path): + """Follow with the real credential source, reading a token file the test controls.""" + monkeypatch.setattr("lightning_sdk.api.logs_api.Auth", Auth) + path = tmp_path / "token" + path.write_text("old-token") + monkeypatch.setenv("LIGHTNING_AUTH_TOKEN_FILE", str(path)) + return path + + +def _sent_tokens(fake_websocket): + return [call.kwargs["header"] for call in fake_websocket.create_connection.call_args_list] + + +def test_follow_says_to_log_in_when_a_refused_token_cannot_refresh(fake_websocket, token_file) -> None: + fake_websocket.create_connection.side_effect = _FakeBadStatusError(401) + + with mock.patch("lightning_sdk.api.logs_api.time.sleep") as sleep, pytest.raises( + ConnectionError, match=re.escape(_LOGIN_FAILED) + ): + list(LogsApi(client=mock.MagicMock()).follow("wss://host/logs")) + + assert _sent_tokens(fake_websocket) == [["Authorization: Bearer old-token"]] + sleep.assert_not_called() + + +def test_follow_reconnects_once_with_a_refreshed_token(fake_websocket, token_file) -> None: + def handshake(url, header, timeout): + if header == ["Authorization: Bearer old-token"]: + token_file.write_text("new-token") + raise _FakeBadStatusError(401) + return _FakeSocket(['[{"message":"live"}]']) + + fake_websocket.create_connection.side_effect = handshake + + entries = list(LogsApi(client=mock.MagicMock()).follow("wss://host/logs", stop=lambda: True)) + + assert [e.message for e in entries] == ["live"] + assert _sent_tokens(fake_websocket) == [ + ["Authorization: Bearer old-token"], + ["Authorization: Bearer new-token"], + ] + + +def test_follow_says_to_log_in_when_the_refreshed_token_is_refused_too(fake_websocket, token_file) -> None: + def handshake(url, header, timeout): + token_file.write_text("new-token") + raise _FakeBadStatusError(401) + + fake_websocket.create_connection.side_effect = handshake + + with pytest.raises(ConnectionError, match=re.escape(_LOGIN_FAILED)): + list(LogsApi(client=mock.MagicMock()).follow("wss://host/logs")) + + assert fake_websocket.create_connection.call_count == 2 + + +def test_follow_stops_on_a_forbidden_handshake(fake_websocket, token_file) -> None: + fake_websocket.create_connection.side_effect = _FakeBadStatusError(403) + + with mock.patch("lightning_sdk.api.logs_api.time.sleep") as sleep, pytest.raises(_FakeBadStatusError): + list(LogsApi(client=mock.MagicMock()).follow("wss://host/logs")) + + assert fake_websocket.create_connection.call_count == 1 + sleep.assert_not_called() + + def test_follow_does_not_swallow_a_failed_handshake_without_reconnect(fake_websocket) -> None: fake_websocket.create_connection.side_effect = OSError("connection refused") From 6574343587addd71b10acb408d7e83849193fd36 Mon Sep 17 00:00:00 2001 From: fnwoke-lightning Date: Thu, 5 Nov 2026 10:02:43 -0500 Subject: [PATCH 20/27] fix(cli): report an unreadable token file from LitServe's authenticate Co-authored-by: Cursor --- python/lightning_sdk/cli/legacy/deploy/_auth.py | 4 ++-- python/tests/cli/test_login_commands.py | 13 +++++++++++++ 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/python/lightning_sdk/cli/legacy/deploy/_auth.py b/python/lightning_sdk/cli/legacy/deploy/_auth.py index 2adb35ec8..fd75e9eba 100644 --- a/python/lightning_sdk/cli/legacy/deploy/_auth.py +++ b/python/lightning_sdk/cli/legacy/deploy/_auth.py @@ -32,8 +32,8 @@ def authenticate(mode: _AuthMode, shall_confirm: bool = True) -> None: with browser_login(False): try: CredentialSource().header() - except ValueError as err: - raise click.UsageError("No Lightning credentials are available. Run `lightning login` first.") from err + except (ValueError, ConnectionError) as err: + raise click.UsageError(str(err)) from err def select_teamspace(teamspace: Optional[str], org: Optional[str], user: Optional[str]) -> Teamspace: diff --git a/python/tests/cli/test_login_commands.py b/python/tests/cli/test_login_commands.py index ec6dadeb6..52b4b7508 100644 --- a/python/tests/cli/test_login_commands.py +++ b/python/tests/cli/test_login_commands.py @@ -1,17 +1,23 @@ """Login, logout, auth status and auth token, against the fake platform.""" import json +import re import threading import webbrowser from types import SimpleNamespace +import click import keyring import pytest from click.testing import CliRunner from lightning_sdk.api.utils import cached_lightning_client from lightning_sdk.cli.entrypoint import main_cli +from lightning_sdk.cli.legacy.deploy._auth import _AuthMode +from lightning_sdk.cli.legacy.deploy._auth import authenticate as litserve_authenticate from lightning_sdk.cli.legacy.entrypoint import main_cli as legacy_cli +from lightning_sdk.cli.utils import lockdown_status +from lightning_sdk.cli.utils.auth import require_auth_header from lightning_sdk.lightning_cloud import env from lightning_sdk.lightning_cloud.credentials import CredentialSource, logout from lightning_sdk.lightning_cloud.login import Auth @@ -306,6 +312,13 @@ def test_an_unreadable_token_file_is_an_error_not_a_traceback(isolated, monkeypa assert "Is a directory" in text assert "Traceback" not in text + named = f"LIGHTNING_AUTH_TOKEN_FILE names {directory}, which can't be read" + assert lockdown_status._auth_header() is None + with pytest.raises(click.UsageError, match=re.escape(named)): + require_auth_header() + with pytest.raises(click.UsageError, match=re.escape(named)): + litserve_authenticate(_AuthMode.DEPLOY) + def test_a_refused_login_is_an_error_not_a_traceback(platform, browser, monkeypatch) -> None: platform.access_ttl = 60 From 68f099d2da6bce0ab87bd8b589eabf4cd786ccf8 Mon Sep 17 00:00:00 2001 From: fnwoke-lightning Date: Thu, 5 Nov 2026 11:25:37 -0500 Subject: [PATCH 21/27] ci: check the login passthrough with LIGHTNING_AUTH_TOKEN, now that the API key is ignored The interim login's headless error now points at LIGHTNING_AUTH_TOKEN too. Co-authored-by: Cursor --- .github/workflows/test-cli.yaml | 2 +- python/lightning_sdk/cli/utils/auth.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/test-cli.yaml b/.github/workflows/test-cli.yaml index 14a250db2..7ea02bff0 100644 --- a/.github/workflows/test-cli.yaml +++ b/.github/workflows/test-cli.yaml @@ -44,7 +44,7 @@ jobs: - name: Login passthrough run: lightning login # since we are providing the credentials, this should be a no-op env: - LIGHTNING_API_KEY: dummy-api-key + LIGHTNING_AUTH_TOKEN: dummy-auth-token - name: Tests working-directory: python/tests diff --git a/python/lightning_sdk/cli/utils/auth.py b/python/lightning_sdk/cli/utils/auth.py index 5de4c68fd..d8026ac66 100644 --- a/python/lightning_sdk/cli/utils/auth.py +++ b/python/lightning_sdk/cli/utils/auth.py @@ -228,7 +228,7 @@ def _login_with_session_token(self: AuthServer, auth: Auth) -> None: ) raise RuntimeError( "Failed to authenticate to Lightning. When running without access to a browser, " - "'LIGHTNING_API_KEY' should be exported." + "'LIGHTNING_AUTH_TOKEN' should be exported." ) @app.get("/login-complete") From 679e9133e487827d43902c4378d8fac019326096 Mon Sep 17 00:00:00 2001 From: fnwoke-lightning Date: Thu, 5 Nov 2026 11:27:09 -0500 Subject: [PATCH 22/27] ci: drop the login passthrough step, since login now refuses when a token is set The API key in the test env is ignored now, so it goes too. Co-authored-by: Cursor --- .github/workflows/test-cli.yaml | 7 ------- 1 file changed, 7 deletions(-) diff --git a/.github/workflows/test-cli.yaml b/.github/workflows/test-cli.yaml index 7ea02bff0..f16c9eecf 100644 --- a/.github/workflows/test-cli.yaml +++ b/.github/workflows/test-cli.yaml @@ -41,15 +41,8 @@ jobs: - name: show pip list run: pip list - - name: Login passthrough - run: lightning login # since we are providing the credentials, this should be a no-op - env: - LIGHTNING_AUTH_TOKEN: dummy-auth-token - - name: Tests working-directory: python/tests - env: # the following value is just a dummy since the actual http calls are mocked out - LIGHTNING_API_KEY: dummy-api-key run: | # NOTE: do not include coverage report here, see: https://github.com/nedbat/coveragepy/issues/1003 coverage run --source lightning_sdk -m pytest cli/ -vv -s From 019c5034465c1d312f5f4ca99595d1eb4aa7831a Mon Sep 17 00:00:00 2001 From: fnwoke-lightning Date: Thu, 5 Nov 2026 14:35:41 -0500 Subject: [PATCH 23/27] refactor(lightning_cloud): rename --insecure-storage to --no-keychain The flag keeps the login in a file only the user can read instead of the OS keychain; the new name says that without alarming users. Co-authored-by: Cursor --- python/lightning_sdk/lightning_cloud/cli_login.py | 8 ++++---- python/lightning_sdk/lightning_cloud/credentials.py | 6 +++--- python/tests/lightning_cloud/test_cli_login.py | 10 +++++----- python/tests/lightning_cloud/test_interim_login.py | 2 +- 4 files changed, 13 insertions(+), 13 deletions(-) diff --git a/python/lightning_sdk/lightning_cloud/cli_login.py b/python/lightning_sdk/lightning_cloud/cli_login.py index b32bdf523..37be33090 100644 --- a/python/lightning_sdk/lightning_cloud/cli_login.py +++ b/python/lightning_sdk/lightning_cloud/cli_login.py @@ -48,7 +48,7 @@ ) _NO_KEYCHAIN = ( "No OS keychain is available to store the login. Run " - "`lightning login --insecure-storage` to keep it in a file only you can read." + "`lightning login --no-keychain` to keep it in a file only you can read." ) _PASTE_WARNING = ( "The code gives access to your Lightning account. Paste it only into the " @@ -97,10 +97,10 @@ def oidc_login_enabled() -> bool: def login( no_browser: bool = False, - insecure_storage: bool = False, + no_keychain: bool = False, ) -> Login: """Log in through the browser, or by pasting the code the login page shows.""" - if insecure_storage: + if no_keychain: storage = _FileStorage.kind elif _keychain_available(): storage = _KeychainStorage.kind @@ -387,7 +387,7 @@ def forget(self, url: str) -> None: class _FileStorage: - """Tokens in the login file itself, for --insecure-storage.""" + """Tokens in the login file itself, for --no-keychain.""" kind = "file" diff --git a/python/lightning_sdk/lightning_cloud/credentials.py b/python/lightning_sdk/lightning_cloud/credentials.py index 5042d89d3..8366a6a7f 100644 --- a/python/lightning_sdk/lightning_cloud/credentials.py +++ b/python/lightning_sdk/lightning_cloud/credentials.py @@ -174,7 +174,7 @@ def _noting_refusal(self, read: Callable[[], Any]) -> Any: def start_login( no_browser: bool = False, - insecure_storage: bool = False, + no_keychain: bool = False, sign_in_params: Optional[Dict[str, str]] = None, ) -> Optional[cli_login.Login]: """Log in with the login ``LIGHTNING_CLI_OIDC_LOGIN`` selects when this is called. @@ -186,13 +186,13 @@ def start_login( if cli_login.oidc_login_enabled(): return cli_login.login( no_browser=no_browser, - insecure_storage=insecure_storage, + no_keychain=no_keychain, ) unavailable = [ flag for flag, given in ( ("--no-browser", no_browser), - ("--insecure-storage", insecure_storage), + ("--no-keychain", no_keychain), ) if given ] diff --git a/python/tests/lightning_cloud/test_cli_login.py b/python/tests/lightning_cloud/test_cli_login.py index 341f97685..b8221097d 100644 --- a/python/tests/lightning_cloud/test_cli_login.py +++ b/python/tests/lightning_cloud/test_cli_login.py @@ -211,10 +211,10 @@ def test_a_deployment_without_a_browser_says_to_include_sdk_credentials(platform cli_login.login() -def test_without_a_keychain_login_refuses_and_names_insecure_storage(platform, browser): +def test_without_a_keychain_login_refuses_and_names_no_keychain(platform, browser): keyring.set_keyring(fail.Keyring()) - with pytest.raises(LoginError, match="--insecure-storage"): + with pytest.raises(LoginError, match="--no-keychain"): cli_login.login() assert browser.visited == [] @@ -231,7 +231,7 @@ def set_password(self, service, username, password): keyring.set_keyring(Refusing()) - with pytest.raises(LoginError, match="--insecure-storage"): + with pytest.raises(LoginError, match="--no-keychain"): cli_login.login() [revoked] = platform.calls_to("/oidc/revoke") @@ -242,10 +242,10 @@ def set_password(self, service, username, password): assert platform.logouts() == [] -def test_insecure_storage_keeps_the_login_in_a_file_only_the_user_can_read(platform, browser, clock): +def test_no_keychain_keeps_the_login_in_a_file_only_the_user_can_read(platform, browser, clock): keyring.set_keyring(fail.Keyring()) - made = cli_login.login(insecure_storage=True) + made = cli_login.login(no_keychain=True) assert platform.calls_to("/oidc/token")[0]["storage"] == "file" assert platform.family_storage(made.refresh_token) == "file" diff --git a/python/tests/lightning_cloud/test_interim_login.py b/python/tests/lightning_cloud/test_interim_login.py index a638d690a..71f519b91 100644 --- a/python/tests/lightning_cloud/test_interim_login.py +++ b/python/tests/lightning_cloud/test_interim_login.py @@ -183,7 +183,7 @@ def test_the_switch_is_read_when_each_login_starts(platform, monkeypatch): ("options", "named"), [ ({"no_browser": True}, "--no-browser"), - ({"insecure_storage": True}, "--insecure-storage"), + ({"no_keychain": True}, "--no-keychain"), ], ) def test_without_the_switch_an_option_only_the_new_login_has_is_refused(platform, monkeypatch, options, named): From 6779e686933b191d670e5c5f45a500e09a2bcaf1 Mon Sep 17 00:00:00 2001 From: fnwoke-lightning Date: Thu, 5 Nov 2026 14:38:00 -0500 Subject: [PATCH 24/27] refactor(cli): rename lightning login --insecure-storage to --no-keychain Co-authored-by: Cursor --- python/lightning_sdk/cli/entrypoint.py | 6 +++--- python/lightning_sdk/cli/legacy/entrypoint.py | 6 +++--- python/lightning_sdk/cli/utils/auth.py | 6 +++--- python/tests/cli/test_login_commands.py | 6 +++--- 4 files changed, 12 insertions(+), 12 deletions(-) diff --git a/python/lightning_sdk/cli/entrypoint.py b/python/lightning_sdk/cli/entrypoint.py index c5f53f722..588c6b372 100644 --- a/python/lightning_sdk/cli/entrypoint.py +++ b/python/lightning_sdk/cli/entrypoint.py @@ -99,13 +99,13 @@ def main_cli(ctx: click.Context) -> None: help="Sign in without a browser on this machine: approve on another device, then paste the code it shows.", ) @click.option( - "--insecure-storage", + "--no-keychain", is_flag=True, help="Keep the login in a file instead of the OS keychain.", ) -def login(no_browser: bool, insecure_storage: bool) -> None: +def login(no_browser: bool, no_keychain: bool) -> None: """Sign in to Lightning AI.""" - run_login(no_browser=no_browser, insecure_storage=insecure_storage) + run_login(no_browser=no_browser, no_keychain=no_keychain) @main_cli.command diff --git a/python/lightning_sdk/cli/legacy/entrypoint.py b/python/lightning_sdk/cli/legacy/entrypoint.py index 7021e3dfd..daf5c6e88 100644 --- a/python/lightning_sdk/cli/legacy/entrypoint.py +++ b/python/lightning_sdk/cli/legacy/entrypoint.py @@ -75,13 +75,13 @@ def main_cli(ctx: click.Context) -> None: help="Sign in without a browser on this machine: approve on another device, then paste the code it shows.", ) @click.option( - "--insecure-storage", + "--no-keychain", is_flag=True, help="Keep the login in a file instead of the OS keychain.", ) -def login(no_browser: bool, insecure_storage: bool) -> None: +def login(no_browser: bool, no_keychain: bool) -> None: """Login to Lightning AI Studios.""" - run_login(no_browser=no_browser, insecure_storage=insecure_storage) + run_login(no_browser=no_browser, no_keychain=no_keychain) @main_cli.command diff --git a/python/lightning_sdk/cli/utils/auth.py b/python/lightning_sdk/cli/utils/auth.py index 6f27a3fc2..979450d19 100644 --- a/python/lightning_sdk/cli/utils/auth.py +++ b/python/lightning_sdk/cli/utils/auth.py @@ -35,7 +35,7 @@ def _env_set(name: str) -> bool: def run_login( no_browser: bool = False, - insecure_storage: bool = False, + no_keychain: bool = False, ) -> None: """Run the login the switch selects, without dropping the saved login first.""" if _env_set(env.LIGHTNING_AUTH_TOKEN_FILE): @@ -56,12 +56,12 @@ def run_login( if in_studio(): raise RuntimeError("Unable to login within a Studio. Did you change your shell setup?") from None - if insecure_storage and cli_login.oidc_login_enabled(): + if no_keychain and cli_login.oidc_login_enabled(): click.echo(f"This login is kept in a file at {env.LIGHTNING_LOGIN_PATH}.", err=True) try: start_login( no_browser=no_browser, - insecure_storage=insecure_storage, + no_keychain=no_keychain, ) except LoginUnavailable as err: raise click.UsageError(str(err)) from err diff --git a/python/tests/cli/test_login_commands.py b/python/tests/cli/test_login_commands.py index 52b4b7508..fca6a61e4 100644 --- a/python/tests/cli/test_login_commands.py +++ b/python/tests/cli/test_login_commands.py @@ -105,7 +105,7 @@ def test_login_uses_the_interim_flow_when_the_switch_is_unset(platform, browser, @pytest.mark.parametrize("cli", [main_cli, legacy_cli]) @pytest.mark.parametrize( "args", - [["--no-browser"], ["--insecure-storage"]], + [["--no-browser"], ["--no-keychain"]], ) def test_interim_login_refuses_flags_the_new_login_needs(isolated, cli, args) -> None: result = _invoke(cli, ["login", *args]) @@ -187,10 +187,10 @@ def test_login_with_the_switch_uses_the_new_login(platform, browser, monkeypatch assert not platform.sign_ins -def test_insecure_storage_warns_and_names_the_login_file(platform, browser, monkeypatch) -> None: +def test_no_keychain_warns_and_names_the_login_file(platform, browser, monkeypatch) -> None: monkeypatch.setenv("LIGHTNING_CLI_OIDC_LOGIN", "1") - result = _invoke(main_cli, ["login", "--insecure-storage"]) + result = _invoke(main_cli, ["login", "--no-keychain"]) assert result.exit_code == 0, _text(result) assert env.LIGHTNING_LOGIN_PATH in _text(result) From 28bf3cee02ee41f4337d798f0a37244faccd708b Mon Sep 17 00:00:00 2001 From: fnwoke-lightning Date: Thu, 5 Nov 2026 16:04:49 -0500 Subject: [PATCH 25/27] fix(cli): refuse login in a Studio before the credential env guards A Studio sets LIGHTNING_AUTH_TOKEN_FILE, so lightning login there reported the token-file refusal instead of the Studio-specific error. Check in_studio() first; the env guards still apply outside a Studio. --- python/lightning_sdk/cli/utils/auth.py | 5 +++-- python/tests/cli/test_login_commands.py | 13 +++++++++++++ 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/python/lightning_sdk/cli/utils/auth.py b/python/lightning_sdk/cli/utils/auth.py index 979450d19..f6caaaf65 100644 --- a/python/lightning_sdk/cli/utils/auth.py +++ b/python/lightning_sdk/cli/utils/auth.py @@ -38,6 +38,9 @@ def run_login( no_keychain: bool = False, ) -> None: """Run the login the switch selects, without dropping the saved login first.""" + if in_studio(): + raise RuntimeError("Unable to login within a Studio. Did you change your shell setup?") from None + if _env_set(env.LIGHTNING_AUTH_TOKEN_FILE): raise click.UsageError( "LIGHTNING_AUTH_TOKEN_FILE is set, so a new login would not be what requests send. " @@ -53,8 +56,6 @@ def run_login( "LIGHTNING_API_KEY is ignored until service accounts. This login continues without it.", err=True, ) - if in_studio(): - raise RuntimeError("Unable to login within a Studio. Did you change your shell setup?") from None if no_keychain and cli_login.oidc_login_enabled(): click.echo(f"This login is kept in a file at {env.LIGHTNING_LOGIN_PATH}.", err=True) diff --git a/python/tests/cli/test_login_commands.py b/python/tests/cli/test_login_commands.py index fca6a61e4..50bfa4f8c 100644 --- a/python/tests/cli/test_login_commands.py +++ b/python/tests/cli/test_login_commands.py @@ -211,6 +211,19 @@ def test_a_studio_without_credentials_still_refuses_to_log_in(monkeypatch) -> No assert "Unable to login within a Studio" in str(result.exception) +def test_a_studio_with_a_workload_token_still_refuses_to_log_in(monkeypatch) -> None: + monkeypatch.delenv("LIGHTNING_AUTH_TOKEN", raising=False) + monkeypatch.delenv("LIGHTNING_API_KEY", raising=False) + monkeypatch.setenv("LIGHTNING_AUTH_TOKEN_FILE", "/workload/token") + monkeypatch.setenv("LIGHTNING_INTERACTIVE", "true") + monkeypatch.setenv("LIGHTNING_CLOUD_SPACE_ID", "studio-1") + + result = _invoke(main_cli, ["login"]) + + assert result.exit_code != 0 + assert "Unable to login within a Studio" in str(result.exception) + + @pytest.mark.parametrize("cli", [main_cli, legacy_cli]) def test_logout_revokes_both_kinds_and_names_the_one_the_server_refused(platform, browser, monkeypatch, cli) -> None: monkeypatch.setenv("LIGHTNING_CLI_OIDC_LOGIN", "1") From 41f1847334ef5da2c076392aa122373027cc2c5f Mon Sep 17 00:00:00 2001 From: fnwoke-lightning Date: Thu, 5 Nov 2026 16:04:53 -0500 Subject: [PATCH 26/27] fix(api): fail log follow when the credential cannot be read Auth().auth_header raises LoginError, a ConnectionError and therefore an OSError, which the reconnect handler caught. With reconnect on and no idle timeout, following logs retried forever on an unreadable token file or refused login. Read the header before the try so the error reaches the caller, and pin it with an unreadable-token-file test. --- python/lightning_sdk/api/logs_api.py | 2 +- python/tests/api/test_logs_api.py | 18 ++++++++++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/python/lightning_sdk/api/logs_api.py b/python/lightning_sdk/api/logs_api.py index 9581d6941..e80489b1f 100644 --- a/python/lightning_sdk/api/logs_api.py +++ b/python/lightning_sdk/api/logs_api.py @@ -522,8 +522,8 @@ def follow( refreshed = False while True: ws = None + auth_header = Auth().auth_header try: - auth_header = Auth().auth_header # The handshake gets a longer budget than the deliberately short read poll. ws = websocket.create_connection( url, header=[f"Authorization: {auth_header}"], timeout=_CONNECT_TIMEOUT diff --git a/python/tests/api/test_logs_api.py b/python/tests/api/test_logs_api.py index f3328bece..2c4b2c034 100644 --- a/python/tests/api/test_logs_api.py +++ b/python/tests/api/test_logs_api.py @@ -13,6 +13,7 @@ _websocket_url, parse_log_entries, ) +from lightning_sdk.lightning_cloud.cli_login import LoginError from lightning_sdk.lightning_cloud.login import Auth from lightning_sdk.lightning_cloud.openapi import V1GetLogsResponse, V1JobLogEntry @@ -414,6 +415,23 @@ def test_follow_says_to_log_in_when_a_refused_token_cannot_refresh(fake_websocke sleep.assert_not_called() +def test_follow_fails_at_once_when_the_token_file_cannot_be_read(fake_websocket, monkeypatch, tmp_path) -> None: + token_path = tmp_path / "token" + token_path.mkdir() + monkeypatch.setattr("lightning_sdk.api.logs_api.Auth", Auth) + monkeypatch.setenv("LIGHTNING_AUTH_TOKEN_FILE", str(token_path)) + stop = mock.MagicMock(return_value=True) + + with mock.patch("lightning_sdk.api.logs_api.time.sleep") as sleep, pytest.raises( + LoginError, match="which can't be read" + ): + list(LogsApi(client=mock.MagicMock()).follow("wss://host/logs", stop=stop)) + + fake_websocket.create_connection.assert_not_called() + sleep.assert_not_called() + stop.assert_not_called() + + def test_follow_reconnects_once_with_a_refreshed_token(fake_websocket, token_file) -> None: def handshake(url, header, timeout): if header == ["Authorization: Bearer old-token"]: From 82ff06fe77a98a79044b4eb87fdf3f52382c7b1f Mon Sep 17 00:00:00 2001 From: fnwoke-lightning Date: Fri, 6 Nov 2026 13:38:02 -0500 Subject: [PATCH 27/27] fix(api): don't refresh or resend the credential on a 401 after a cross-host redirect requests follows redirects with a copy of the request, response hooks included, and drops Authorization when the host changes. The 401 hook then ran on the redirect target's response, refreshed, and resent the fresh token to that host, then told the user to log in. Retry only a 401 to a request that carried the credential. Co-authored-by: Cursor --- python/lightning_sdk/api/utils.py | 6 ++- python/tests/api/test_credential_requests.py | 47 ++++++++++++++++++++ 2 files changed, 52 insertions(+), 1 deletion(-) diff --git a/python/lightning_sdk/api/utils.py b/python/lightning_sdk/api/utils.py index f634ef3a4..4a68a11ec 100644 --- a/python/lightning_sdk/api/utils.py +++ b/python/lightning_sdk/api/utils.py @@ -1329,7 +1329,8 @@ class CredentialAuth(AuthBase): """Sets ``Authorization`` when a request is prepared, and retries one 401 after a refresh. A 401 that no refresh fixes raises ``ConnectionError`` saying to run ``lightning login``, - as the generated client does. + as the generated client does. A 401 to a request that didn't carry the credential, such + as a redirect to another host, is returned as it is. """ def __init__(self, source: Optional[CredentialSource] = None) -> None: @@ -1349,6 +1350,9 @@ def handle(response: requests.Response, **kwargs: Any) -> requests.Response: nonlocal retried if retried or response.status_code != 401: return response + # Redirects reuse this hook, and requests drops the header on a cross-host redirect. + if response.request.headers.get("Authorization") != header: + return response # The connection stays checked out until the body is read, so the retry can reuse it. _ = response.content response.close() diff --git a/python/tests/api/test_credential_requests.py b/python/tests/api/test_credential_requests.py index bf1b28943..7ecb215f4 100644 --- a/python/tests/api/test_credential_requests.py +++ b/python/tests/api/test_credential_requests.py @@ -184,6 +184,53 @@ def test_blob_upload_retries_a_401(refreshed_token, tmp_path): assert refreshed_token.seen == ["old-token", "new-token"] +def test_a_401_from_another_host_after_a_redirect_is_returned_without_sending_the_token(monkeypatch, tmp_path): + threading.Thread = _REAL_THREAD + for name in ("LIGHTNING_AUTH_TOKEN", "LIGHTNING_API_KEY"): + monkeypatch.delenv(name, raising=False) + token_path = tmp_path / "token" + token_path.write_text("old-token") + monkeypatch.setenv("LIGHTNING_AUTH_TOKEN_FILE", str(token_path)) + seen = {"api": [], "storage": []} + + def handler(name, status, location=None): + class Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_GET(self): # noqa: N802 + seen[name].append(self.headers.get("Authorization")) + if name == "storage": + token_path.write_text("new-token") + self.send_response(status) + if location: + self.send_header("Location", location) + self.send_header("Content-Length", "0") + self.end_headers() + + def log_message(self, fmt, *args): + return + + return Handler + + storage = ThreadingHTTPServer(("127.0.0.1", 0), handler("storage", 401)) + api = ThreadingHTTPServer( + ("127.0.0.1", 0), handler("api", 302, f"http://localhost:{storage.server_address[1]}/blob") + ) + threads = [_REAL_THREAD(target=server.serve_forever, daemon=True) for server in (api, storage)] + for thread in threads: + thread.start() + try: + response = requests.get(f"http://127.0.0.1:{api.server_address[1]}/download", auth=CredentialAuth(), timeout=5) + finally: + for server, thread in zip((api, storage), threads): + server.shutdown() + thread.join(timeout=5) + server.server_close() + + assert response.status_code == 401 + assert seen == {"api": ["Bearer old-token"], "storage": [None]} + + def test_folder_download_refreshes_when_the_access_token_expires_midway(monkeypatch, tmp_path, inline_executor_cls): threading.Thread = _REAL_THREAD for name in ("LIGHTNING_AUTH_TOKEN", "LIGHTNING_API_KEY", "LIGHTNING_AUTH_TOKEN_FILE"):