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 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 9cc195f62..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 @@ -31,7 +32,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 +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=self._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) @@ -135,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. @@ -157,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: @@ -164,7 +166,7 @@ def download() -> None: try: r = requests.get( url, - headers=self._auth_headers, + auth=request, stream=True, allow_redirects=True, timeout=(_DOWNLOAD_CONNECT_TIMEOUT_SECONDS, _DOWNLOAD_READ_TIMEOUT_SECONDS), @@ -253,6 +255,7 @@ def download_folder( mininterval=1, ) + auth = credential_auth() try: with ThreadPoolExecutor(max_workers=num_workers) as executor: futures = [ @@ -263,6 +266,7 @@ def download_folder( teamspace_id, pbar, expected_size=entry.get("size"), + auth=auth, ) for entry in files ] @@ -288,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=self._auth_headers, + auth=credential_auth(), ) if r.status_code not in (200, 204): raise _RemoteApiError( @@ -311,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=self._auth_headers, + auth=credential_auth(), ) 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/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 fa6b99043..e80489b1f 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, @@ -506,7 +508,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 @@ -518,8 +519,10 @@ def follow( last_line_at = time.monotonic() backoff = _MIN_RECONNECT_BACKOFF + refreshed = False while True: ws = None + auth_header = Auth().auth_header try: # The handshake gets a longer budget than the deliberately short read poll. ws = websocket.create_connection( @@ -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/lightning_sdk/api/sandbox_api.py b/python/lightning_sdk/api/sandbox_api.py index 81a9cd264..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.login import Auth +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, @@ -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,14 @@ 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) + 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 @@ -243,11 +247,14 @@ 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 "") + 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 30a791f4f..4a68a11ec 100644 --- a/python/lightning_sdk/api/utils.py +++ b/python/lightning_sdk/api/utils.py @@ -30,10 +30,12 @@ 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 import env +from lightning_sdk.lightning_cloud.credentials import CredentialSource from lightning_sdk.lightning_cloud.openapi import ( ModelsStoreCompleteMultiPartUploadBody, ModelsStoreCreateMultiPartUploadBody, @@ -44,7 +46,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 @@ -255,7 +257,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.""" @@ -300,7 +301,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, 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 @@ -1324,6 +1325,61 @@ 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. + + A 401 that no refresh fixes raises ``ConnectionError`` saying to run ``lightning login``, + 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: + self._source = source or CredentialSource() + + 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 + + retried = False + + 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() + 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 + + 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/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..4667d6bd7 --- /dev/null +++ b/python/lightning_sdk/cli/auth/status.py @@ -0,0 +1,131 @@ +"""Show the credential requests send, and what a saved login reaches.""" + +import os +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, env +from lightning_sdk.lightning_cloud.credentials import Credential, CredentialKind, CredentialSource + +_NOT_LOGGED_IN = "You are not logged in. Run `lightning login`." +_UNKNOWN_REACH = "not known locally" + + +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: + scopes = grant.get("orgs") + if scopes == "all": + return "all", "all" + if not isinstance(scopes, list) or not scopes: + return "none", "none" + + shown_orgs = [] + shown_teamspaces = [] + for scope in scopes: + if not isinstance(scope, dict): + continue + org_id = str(scope.get("org", "")) + 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" + + +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 the credential requests send, who it belongs to, and what a saved login reaches.""" + 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) + + client = cached_lightning_client() + user = _user(client) + table = Table(show_header=False, box=None) + table.add_column("Field", style="bold") + table.add_column("Value") + _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: + 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 new file mode 100644 index 000000000..3aa6583d9 --- /dev/null +++ b/python/lightning_sdk/cli/auth/token.py @@ -0,0 +1,41 @@ +"""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, 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 the credential requests send. A saved login has at least 30 minutes left.""" + 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: + try: + click.echo(cli_login.token()) + except cli_login.LoginError as err: + raise click.ClickException(str(err)) from err + return + 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 4c4a5cb97..588c6b372 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,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( + "--no-keychain", + is_flag=True, + help="Keep the login in a file instead of the OS keychain.", +) +def login(no_browser: bool, no_keychain: 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.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, no_keychain=no_keychain) @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/deploy/_auth.py b/python/lightning_sdk/cli/legacy/deploy/_auth.py index fcbe61e11..fd75e9eba 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, 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/lightning_sdk/cli/legacy/entrypoint.py b/python/lightning_sdk/cli/legacy/entrypoint.py index f36108b72..daf5c6e88 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,25 @@ 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( + "--no-keychain", + is_flag=True, + help="Keep the login in a file instead of the OS keychain.", +) +def login(no_browser: bool, no_keychain: bool) -> 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, no_keychain=no_keychain) @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/resource_completion.py b/python/lightning_sdk/cli/resource_completion.py index 08cbf2d10..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.auth_token: - return True - if not auth.load(): + try: + return bool(Auth().auth_header) + except Exception: return False - return bool(auth.auth_token) def accessible_teamspaces() -> dict[str, dict[str, str]]: diff --git a/python/lightning_sdk/cli/utils/auth.py b/python/lightning_sdk/cli/utils/auth.py index d8026ac66..f6caaaf65 100644 --- a/python/lightning_sdk/cli/utils/auth.py +++ b/python/lightning_sdk/cli/utils/auth.py @@ -1,384 +1,84 @@ -"""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 -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 ( + REVOKE_FAILED, + 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 - -# 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) +def require_auth_header() -> str: + """The credential requests send, or a usage error that tells the user to log in.""" 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 + 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 - 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) +def _env_set(name: str) -> bool: + return bool(os.environ.get(name, "").strip()) - 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 run_login( + no_browser: bool = False, + 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. " + "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_AUTH_TOKEN' 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, ) - 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: + 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: - 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) - + start_login( + no_browser=no_browser, + no_keychain=no_keychain, + ) + 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() 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[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/lightning_sdk/cli/utils/lockdown_status.py b/python/lightning_sdk/cli/utils/lockdown_status.py index 9e4d4e846..0589876f1 100644 --- a/python/lightning_sdk/cli/utils/lockdown_status.py +++ b/python/lightning_sdk/cli/utils/lockdown_status.py @@ -81,16 +81,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.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/lightning_sdk/cli/utils/logging.py b/python/lightning_sdk/cli/utils/logging.py index b2c07fb7b..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.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/lightning_sdk/lightning_cloud/login.py b/python/lightning_sdk/lightning_cloud/login.py index b318037e0..5b960377e 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_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..7ecb215f4 --- /dev/null +++ b/python/tests/api/test_credential_requests.py @@ -0,0 +1,293 @@ +"""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 + +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.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 +_LOGIN_FAILED = "Authentication failed. Please run `lightning login`." + + +class _RetryServer: + """Answers 401 for ``old-token`` after rewriting the token file, then 200. + + 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 = [] + + 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) + user_agents.append(self.headers.get("User-Agent")) + if token == "old-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" + 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.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) + self._thread.start() + + def stop(self): + self._httpd.shutdown() + self._thread.join(timeout=5) + self._httpd.server_close() + + +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, 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] + + 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"] + + +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): + _blob_request(refreshed_token.url, 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"): + 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 09e7222b2..1b82b22ba 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"}, @@ -318,7 +266,7 @@ def test_list_files_follows_cursor_pages(_authenticate_mock, _mock_auth, mock_re assert all(p["recursive"] == "true" and p["atMost"] == "1" for p in params) -@mock.patch("lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") @mock.patch("lightning_sdk.api.utils.LightningClient") @mock.patch("requests.get", autospec=True) def test_download_folder_preserves_each_file_mtime(mock_get, _mock_client, _mock_auth, tmp_path, inline_executor_cls): 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_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_logs_api.py b/python/tests/api/test_logs_api.py index 8316d97f1..2c4b2c034 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,12 @@ _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 +_LOGIN_FAILED = "Authentication failed. Please run `lightning login`." + def _api(*pages): """A LogsApi whose client returns ``pages`` in order.""" @@ -275,6 +280,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 +293,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 +389,90 @@ 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_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"]: + 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") diff --git a/python/tests/api/test_studio_api.py b/python/tests/api/test_studio_api.py index 64a489bb1..786d9af8c 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", "x-amz-meta-mtime": "1700000000.123456789"} @@ -1288,15 +1263,11 @@ def test_download_file(mock_authenticate, mock_requests_get, tmpdir): studio_api.download_file("file1", filepath, "st-abc", "ts-abc", "cluster-abc") assert os.stat(filepath).st_mtime_ns == 1700000000123456789 mock_requests_get.assert_called_once() - 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 = {} @@ -1310,17 +1281,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( @@ -1395,13 +1362,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=[ @@ -2067,15 +2029,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, @@ -2105,9 +2062,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) @@ -2149,15 +2105,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, @@ -2187,20 +2138,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"}, @@ -2222,7 +2168,7 @@ def test_list_files_follows_cursor_pages(_mock_authenticate, mock_requests_get): assert params[1]["cursor"] == "cursor-1" -@mock.patch("lightning_sdk.api.studio_api._authenticate_and_get_auth_headers") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") @mock.patch("lightning_sdk.api.utils.LightningClient") @mock.patch("requests.get", autospec=True) def test_download_folder_preserves_mtime(mock_get, _mock_client, _mock_auth, tmp_path, inline_executor_cls): diff --git a/python/tests/api/test_teamspace_api.py b/python/tests/api/test_teamspace_api.py index 850e7fbd9..044015e9d 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) mock_requests_get.return_value.headers["x-amz-meta-mtime"] = "1700000000.123456789" @@ -897,15 +862,11 @@ def test_download_file(mock_authenticate, mock_requests_get, tmpdir): assert Path(filepath).read_bytes() == b"data" * 256 assert Path(filepath).stat().st_mtime_ns == 1700000000123456789 mock_requests_get.assert_called_once() - 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() @@ -933,13 +894,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}] @@ -961,11 +917,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( @@ -1200,12 +1153,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"}, @@ -1349,7 +1298,7 @@ def test_list_workload_tags_tolerates_a_teamspace_with_no_tags(): assert teamspace_api.list_workload_tags("ts-abc") == [] -@mock.patch("lightning_sdk.api.teamspace_api._authenticate_and_get_auth_headers") +@mock.patch("lightning_sdk.lightning_cloud.credentials.Auth") @mock.patch("lightning_sdk.api.utils.LightningClient") @mock.patch("requests.get", autospec=True) def test_download_folder_preserves_mtime(mock_get, _mock_client, _mock_auth, tmp_path, inline_executor_cls): diff --git a/python/tests/api/test_utils.py b/python/tests/api/test_utils.py index 12f144272..6919bf466 100644 --- a/python/tests/api/test_utils.py +++ b/python/tests/api/test_utils.py @@ -11,6 +11,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, @@ -39,9 +40,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, @@ -115,7 +113,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}], @@ -795,18 +793,14 @@ def test_resolve_path_mappings(): 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): @@ -826,7 +820,7 @@ def test_single_part_uploader_basic(mock_requests, tmp_path): 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 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/test_entrypoint.py b/python/tests/cli/test_entrypoint.py index f5c0ef4aa..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,69 +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" - - -def test_login_not_required_with_credentials_in_env(monkeypatch, tmp_path): - _clear_auth(monkeypatch, tmp_path) - monkeypatch.setenv("LIGHTNING_AUTH_TOKEN", "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() - - -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): - 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..50bfa4f8c --- /dev/null +++ b/python/tests/cli/test_login_commands.py @@ -0,0 +1,467 @@ +"""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 +from tests.lightning_cloud.utils.fake_platform import FakeBrowser, FakePlatform +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"], ["--no-keychain"]], +) +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() + + +@pytest.mark.parametrize("cli", [main_cli, legacy_cli]) +@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 "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_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", "--no-browser"]) + + assert result.exit_code == 2, _text(result) + assert "LIGHTNING_CLI_OIDC_LOGIN" in _text(result) + 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_no_keychain_warns_and_names_the_login_file(platform, browser, monkeypatch) -> None: + monkeypatch.setenv("LIGHTNING_CLI_OIDC_LOGIN", "1") + + result = _invoke(main_cli, ["login", "--no-keychain"]) + + 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" + + +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) + + +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") + 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 + + +@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 + + 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 + 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 + 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) + + 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"]) + 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: + 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_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) + + 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" + 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 b8dbfad84..0084420d0 100644 --- a/python/tests/cli/utils/test_auth.py +++ b/python/tests/cli/utils/test_auth.py @@ -1,438 +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(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) - - 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.auth_header == "Bearer sess_user-1.secret" - assert json.loads(secrets_file.read_text()) == { - "user_id": "user-1", - "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 - - -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" 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/filesystem/test_filesystem_copy.py b/python/tests/core/filesystem/test_filesystem_copy.py index 4f19ccb16..1c22da824 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,18 +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", @@ -50,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): @@ -81,13 +70,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, @@ -101,7 +88,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 @@ -128,11 +114,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", @@ -151,11 +136,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): 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/lightning_cloud/test_credentials.py b/python/tests/lightning_cloud/test_credentials.py index 78a59a4e7..beb02cda8 100644 --- a/python/tests/lightning_cloud/test_credentials.py +++ b/python/tests/lightning_cloud/test_credentials.py @@ -627,11 +627,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" 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 "" 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()