From 06524864a748a6ada26528bdcca6ae2f202b9575 Mon Sep 17 00:00:00 2001 From: Justus Perillieux Date: Thu, 30 Jul 2026 16:18:49 +0200 Subject: [PATCH] feat(sdk): add environment and secret management --- go/secrets_test.go | 96 +++++++++++++++++ go/shared.go | 7 +- go/studio.go | 25 +++++ go/studio_lifecycle_test.go | 46 ++++++++ go/teamspace.go | 33 +++++- go/user.go | 31 +++++- python/lightning_sdk/api/deployment_api.py | 2 +- python/lightning_sdk/api/teamspace_api.py | 36 ++++++- python/lightning_sdk/api/user_api.py | 25 ++++- .../lightning_sdk/cli/deployment/__init__.py | 2 + python/lightning_sdk/cli/deployment/env.py | 98 +++++++++++++++++ python/lightning_sdk/cli/entrypoint.py | 6 +- python/lightning_sdk/cli/groups.py | 14 +++ python/lightning_sdk/cli/studio/__init__.py | 2 + python/lightning_sdk/cli/studio/env.py | 93 ++++++++++++++++ .../lightning_sdk/cli/teamspace/__init__.py | 10 ++ python/lightning_sdk/cli/teamspace/secret.py | 81 ++++++++++++++ python/lightning_sdk/cli/user/__init__.py | 10 ++ python/lightning_sdk/cli/user/secret.py | 60 +++++++++++ .../lightning_sdk/cli/utils/configuration.py | 59 ++++++++++ .../cli/utils/resource_resolution.py | 13 +++ python/lightning_sdk/deployment.py | 53 +++++++++ python/lightning_sdk/studio.py | 29 +++++ python/lightning_sdk/teamspace.py | 17 +++ python/lightning_sdk/user.py | 17 +++ python/tests/api/test_teamspace_api.py | 51 ++++++++- python/tests/api/test_user_api.py | 51 ++++++++- python/tests/cli/deployment/test_env.py | 62 +++++++++++ python/tests/cli/studio/test_env.py | 57 ++++++++++ python/tests/cli/teamspace/__init__.py | 1 + python/tests/cli/teamspace/test_secret.py | 51 +++++++++ python/tests/cli/user/__init__.py | 1 + python/tests/cli/user/test_secret.py | 60 +++++++++++ python/tests/cli/utils/test_configuration.py | 73 +++++++++++++ .../cli/utils/test_resource_resolution.py | 30 ++++++ .../core/studio/test_studio_environment.py | 57 ++++++++++ python/tests/core/test_deployment.py | 101 +++++++++++++++++- python/tests/core/test_teamspace.py | 32 +++++- python/tests/core/test_user.py | 23 ++++ 39 files changed, 1500 insertions(+), 15 deletions(-) create mode 100644 python/lightning_sdk/cli/deployment/env.py create mode 100644 python/lightning_sdk/cli/studio/env.py create mode 100644 python/lightning_sdk/cli/teamspace/__init__.py create mode 100644 python/lightning_sdk/cli/teamspace/secret.py create mode 100644 python/lightning_sdk/cli/user/__init__.py create mode 100644 python/lightning_sdk/cli/user/secret.py create mode 100644 python/lightning_sdk/cli/utils/configuration.py create mode 100644 python/tests/cli/deployment/test_env.py create mode 100644 python/tests/cli/studio/test_env.py create mode 100644 python/tests/cli/teamspace/__init__.py create mode 100644 python/tests/cli/teamspace/test_secret.py create mode 100644 python/tests/cli/user/__init__.py create mode 100644 python/tests/cli/user/test_secret.py create mode 100644 python/tests/cli/utils/test_configuration.py diff --git a/go/secrets_test.go b/go/secrets_test.go index 298de72d..e8b5f42b 100644 --- a/go/secrets_test.go +++ b/go/secrets_test.go @@ -127,3 +127,99 @@ func TestTeamspaceSecretsAndSetSecretUseGeneratedRoutes(t *testing.T) { } } + +func TestUserSecretsDeleteAndTypedCollision(t *testing.T) { + var seen []string + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + seen = append(seen, r.Method+" "+r.URL.RequestURI()) + w.Header().Set("Content-Type", "application/json") + switch r.Method + " " + r.URL.Path { + case "GET /v1/secrets": + _ = json.NewEncoder(w).Encode(map[string]any{ + "secrets": []map[string]any{ + {"id": "typed-secret", "name": "TOKEN", "type": "SECRET_TYPE_DOCKER_REGISTRY"}, + {"id": "generic-secret", "name": "OLD_TOKEN", "type": "SECRET_TYPE_UNSPECIFIED"}, + }, + }) + case "POST /v1/secrets": + var body struct { + Name string `json:"name"` + Value string `json:"value"` + } + require.NoError(t, json.NewDecoder(r.Body).Decode(&body)) + assert.Equal(t, "TOKEN", body.Name) + assert.Equal(t, "created", body.Value) + _ = json.NewEncoder(w).Encode(map[string]any{}) + case "DELETE /v1/secrets/generic-secret": + _ = json.NewEncoder(w).Encode(map[string]any{}) + default: + assert.Fail(t, fmt.Sprintf("unexpected request: %s %s", r.Method, r.URL.RequestURI())) + } + })) + defer server.Close() + t.Setenv("LIGHTNING_CLOUD_URL", server.URL) + + u := mustUser(t, "user-1", "alice") + require.NoError(t, u.SetSecret("TOKEN", "created")) + require.NoError(t, u.DeleteSecret("OLD_TOKEN")) + require.ErrorContains(t, u.DeleteSecret("MISSING"), `generic user secret "MISSING" was not found`) + require.Error(t, u.DeleteSecret("invalid-name")) + + assert.Equal(t, []string{ + "GET /v1/secrets", + "POST /v1/secrets", + "GET /v1/secrets", + "DELETE /v1/secrets/generic-secret", + "GET /v1/secrets", + }, seen) +} + +func TestTeamspaceSecretsDeleteAndTypedCollision(t *testing.T) { + var seen []string + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + seen = append(seen, r.Method+" "+r.URL.RequestURI()) + w.Header().Set("Content-Type", "application/json") + switch r.Method + " " + r.URL.Path { + case "GET /v1/projects/project-1/secrets": + _ = json.NewEncoder(w).Encode(map[string]any{ + "secrets": []map[string]any{ + {"id": "typed-secret", "name": "TOKEN", "type": "SECRET_TYPE_HF_TOKEN"}, + {"id": "generic-secret", "name": "OLD_TOKEN", "type": "SECRET_TYPE_UNSPECIFIED"}, + }, + }) + case "POST /v1/projects/project-1/secrets": + var body struct { + Name string `json:"name"` + Type string `json:"type"` + Value string `json:"value"` + } + require.NoError(t, json.NewDecoder(r.Body).Decode(&body)) + assert.Equal(t, "TOKEN", body.Name) + assert.Equal(t, "SECRET_TYPE_UNSPECIFIED", body.Type) + assert.Equal(t, "created", body.Value) + _ = json.NewEncoder(w).Encode(map[string]any{}) + case "DELETE /v1/projects/project-1/secrets/generic-secret": + _ = json.NewEncoder(w).Encode(map[string]any{}) + default: + assert.Fail(t, fmt.Sprintf("unexpected request: %s %s", r.Method, r.URL.RequestURI())) + } + })) + defer server.Close() + t.Setenv("LIGHTNING_CLOUD_URL", server.URL) + + ts := mustTeamspace(t, "project-1", "default", "alice") + require.NoError(t, ts.SetSecret("TOKEN", "created")) + require.NoError(t, ts.DeleteSecret("OLD_TOKEN")) + require.ErrorContains(t, ts.DeleteSecret("MISSING"), `generic teamspace secret "MISSING" was not found`) + require.Error(t, ts.DeleteSecret("invalid-name")) + + assert.Equal(t, []string{ + "GET /v1/projects/project-1/secrets", + "POST /v1/projects/project-1/secrets", + "GET /v1/projects/project-1/secrets", + "DELETE /v1/projects/project-1/secrets/generic-secret", + "GET /v1/projects/project-1/secrets", + }, seen) +} diff --git a/go/shared.go b/go/shared.go index ab0f2cdf..db6cfc42 100644 --- a/go/shared.go +++ b/go/shared.go @@ -115,10 +115,15 @@ func validSecretName(name string) bool { return secretNamePattern.MatchString(name) } +func isGenericSecret(secret *models.V1Secret) bool { + return secret != nil && + (secret.Type == nil || *secret.Type == models.V1SecretTypeSECRETTYPEUNSPECIFIED) +} + func redactedSecrets(secrets []*models.V1Secret) map[string]string { result := map[string]string{} for _, secret := range secrets { - if secret == nil || secret.Type == nil || *secret.Type != models.V1SecretTypeSECRETTYPEUNSPECIFIED { + if !isGenericSecret(secret) { continue } result[secret.Name] = "***REDACTED***" diff --git a/go/studio.go b/go/studio.go index 44fae0b8..3a542a92 100644 --- a/go/studio.go +++ b/go/studio.go @@ -879,6 +879,11 @@ func (s *Studio) SetEnv(newEnv map[string]string, partial bool) error { if s == nil || s.teamspaceID == "" || s.id == "" { return errors.New("studio set env requires teamspace ID and studio ID") } + for key := range newEnv { + if !validSecretName(key) { + return errors.New("environment variable names must start with a letter or underscore and contain only letters, numbers, and underscores") + } + } updatedEnv := map[string]string{} if partial { for key, value := range s.env { @@ -919,6 +924,26 @@ func (s *Studio) SetEnv(newEnv map[string]string, partial bool) error { return nil } +// DeleteEnv deletes one directly configured Studio environment variable. +func (s *Studio) DeleteEnv(key string) error { + if s == nil || s.teamspaceID == "" || s.id == "" { + return errors.New("studio delete env requires teamspace ID and studio ID") + } + if !validSecretName(key) { + return errors.New("environment variable names must start with a letter or underscore and contain only letters, numbers, and underscores") + } + if _, ok := s.env[key]; !ok { + return fmt.Errorf("studio environment variable %q was not found", key) + } + updated := make(map[string]string, len(s.env)-1) + for name, value := range s.env { + if name != key { + updated[name] = value + } + } + return s.SetEnv(updated, false) +} + // AvailablePlugins lists plugins available for installation. func (s *Studio) AvailablePlugins() (map[string]string, error) { if s == nil || s.teamspaceID == "" || s.id == "" { diff --git a/go/studio_lifecycle_test.go b/go/studio_lifecycle_test.go index 3f0ed66a..a32c6f4b 100644 --- a/go/studio_lifecycle_test.go +++ b/go/studio_lifecycle_test.go @@ -710,6 +710,52 @@ func TestStudioSetEnvMergesByDefault(t *testing.T) { } +func TestStudioDeleteEnvValidatesAndPreservesUnrelatedVariables(t *testing.T) { + var seen []string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + seen = append(seen, r.Method+" "+r.URL.RequestURI()) + w.Header().Set("Content-Type", "application/json") + assert.Equal(t, http.MethodPut, r.Method) + assert.Equal(t, "/v1/projects/project-1/cloudspaces/studio-1", r.URL.Path) + + var body struct { + Env []struct { + Name string `json:"name"` + Value string `json:"value"` + } `json:"env"` + } + require.NoError(t, json.NewDecoder(r.Body).Decode(&body)) + require.Len(t, body.Env, 1) + assert.Equal(t, "KEEP", body.Env[0].Name) + assert.Equal(t, "yes", body.Env[0].Value) + + _ = json.NewEncoder(w).Encode(map[string]any{ + "id": "studio-1", + "name": "dev", + "projectId": "project-1", + "env": []map[string]any{{"name": "KEEP", "value": "yes"}}, + }) + })) + defer server.Close() + t.Setenv("LIGHTNING_CLOUD_URL", server.URL) + + s, err := lit.GetStudio("dev", lit.StudioOptions{ + ID: "studio-1", + Teamspace: mustTeamspace(t, "project-1", ""), + Env: map[string]string{"KEEP": "yes", "OLD": "remove"}, + }) + require.NoError(t, err) + + require.Error(t, s.SetEnv(map[string]string{"INVALID-NAME": "value"}, true)) + require.Error(t, s.DeleteEnv("INVALID-NAME")) + require.ErrorContains(t, s.DeleteEnv("MISSING"), `studio environment variable "MISSING" was not found`) + require.Empty(t, seen) + + require.NoError(t, s.DeleteEnv("OLD")) + assert.Equal(t, map[string]string{"KEEP": "yes"}, s.Env()) + assert.Equal(t, []string{"PUT /v1/projects/project-1/cloudspaces/studio-1"}, seen) +} + func TestStudioListsPluginsThroughGeneratedRoutes(t *testing.T) { var seen []string diff --git a/go/teamspace.go b/go/teamspace.go index 4663d976..7d2d4cca 100644 --- a/go/teamspace.go +++ b/go/teamspace.go @@ -542,7 +542,7 @@ func (t *Teamspace) SetSecret(key, value string) error { return err } for _, secret := range secrets { - if secret != nil && secret.Name == key { + if isGenericSecret(secret) && secret.Name == key { _, err := api.SecretService.SecretServiceUpdateSecret( secret_service.NewSecretServiceUpdateSecretParamsWithContext(context.Background()). WithProjectID(id). @@ -561,6 +561,37 @@ func (t *Teamspace) SetSecret(key, value string) error { return err } +// DeleteSecret deletes a generic teamspace secret by name. +func (t *Teamspace) DeleteSecret(key string) error { + if !validSecretName(key) { + return errors.New("secret keys must only contain alphanumeric characters and underscores and not begin with a number") + } + id, err := t.requireID("delete secret") + if err != nil { + return err + } + secrets, err := t.listSecrets() + if err != nil { + return err + } + for _, secret := range secrets { + if !isGenericSecret(secret) || secret.Name != key { + continue + } + api, err := sdkclient.New() + if err != nil { + return err + } + _, err = api.SecretService.SecretServiceDeleteSecret( + secret_service.NewSecretServiceDeleteSecretParamsWithContext(context.Background()). + WithProjectID(id). + WithID(secret.ID), + ) + return err + } + return fmt.Errorf("generic teamspace secret %q was not found", key) +} + // NewFolder creates a managed folder in the teamspace. func (t *Teamspace) NewFolder(name string, opts ...FolderOptions) error { id, err := t.requireID("new folder") diff --git a/go/user.go b/go/user.go index d9369092..0f7ea7ff 100644 --- a/go/user.go +++ b/go/user.go @@ -184,7 +184,7 @@ func (u *User) SetSecret(key, value string) error { return err } for _, secret := range secrets { - if secret != nil && secret.Name == key { + if isGenericSecret(secret) && secret.Name == key { _, err := api.SecretService.SecretServiceUpdateUserSecret( secret_service.NewSecretServiceUpdateUserSecretParamsWithContext(context.Background()). WithID(secret.ID). @@ -200,6 +200,35 @@ func (u *User) SetSecret(key, value string) error { return err } +// DeleteSecret deletes a generic user secret by name. +func (u *User) DeleteSecret(key string) error { + if u == nil { + return errors.New("user delete secret requires user") + } + if !validSecretName(key) { + return errors.New("secret keys must only contain alphanumeric characters and underscores and not begin with a number") + } + secrets, err := u.listSecrets() + if err != nil { + return err + } + for _, secret := range secrets { + if !isGenericSecret(secret) || secret.Name != key { + continue + } + api, err := sdkclient.New() + if err != nil { + return err + } + _, err = api.SecretService.SecretServiceDeleteUserSecret( + secret_service.NewSecretServiceDeleteUserSecretParamsWithContext(context.Background()). + WithID(secret.ID), + ) + return err + } + return fmt.Errorf("generic user secret %q was not found", key) +} + func (u *User) listSecrets() ([]*models.V1Secret, error) { api, err := sdkclient.New() if err != nil { diff --git a/python/lightning_sdk/api/deployment_api.py b/python/lightning_sdk/api/deployment_api.py index 6cc44ba4..944197b8 100644 --- a/python/lightning_sdk/api/deployment_api.py +++ b/python/lightning_sdk/api/deployment_api.py @@ -1404,7 +1404,7 @@ def restore_env(env: List[V1EnvVar]) -> List[Union[Secret, Env]]: def to_env(env: Union[List[Union[Secret, Env]], Dict[str, str], None] = None) -> Optional[List[V1EnvVar]]: - if not env: + if env is None: return None env_list = [] diff --git a/python/lightning_sdk/api/teamspace_api.py b/python/lightning_sdk/api/teamspace_api.py index 1bea2402..14b20758 100644 --- a/python/lightning_sdk/api/teamspace_api.py +++ b/python/lightning_sdk/api/teamspace_api.py @@ -96,6 +96,12 @@ def _resolve_secret_type(secret_type: Union[str, SecretType]) -> V1SecretType: return _SECRET_TYPE_TO_API[resolved] +def _secret_matches_type(secret: V1Secret, secret_type: V1SecretType) -> bool: + if secret_type == V1SecretType.UNSPECIFIED: + return secret.type in (None, V1SecretType.UNSPECIFIED) + return secret.type == secret_type + + class TeamspaceApi: """Internal API client for Teamspace requests (mainly http requests).""" @@ -941,7 +947,11 @@ def get_secrets(self, teamspace_id: str) -> Dict[str, str]: # this returns encrypted values for security. It doesn't make sense to show them, # so we just return a placeholder # not a security issue to replace in the client as we get the encrypted values from the server. - return {secret.name: "***REDACTED***" for secret in secrets if secret.type == V1SecretType.UNSPECIFIED} + return { + secret.name: "***REDACTED***" + for secret in secrets + if _secret_matches_type(secret, V1SecretType.UNSPECIFIED) + } def set_secret( self, @@ -965,10 +975,32 @@ def set_secret( resolved_type = _resolve_secret_type(secret_type) secrets = self._get_secrets(teamspace_id) for secret in secrets: - if secret.name == key: + if secret.name == key and _secret_matches_type(secret, resolved_type): return self._update_secret(teamspace_id, secret.id, value) return self._create_secret(teamspace_id, key, value, secret_type=resolved_type) + def delete_secret(self, teamspace_id: str, key: str) -> None: + """Delete a generic secret from a teamspace. + + Args: + teamspace_id: ID of the teamspace that owns the secret. + key: The generic secret name. + + Raises: + ValueError: If no generic secret with the given name exists. + """ + secret = next( + ( + secret + for secret in self._get_secrets(teamspace_id) + if secret.name == key and _secret_matches_type(secret, V1SecretType.UNSPECIFIED) + ), + None, + ) + if secret is None: + raise ValueError(f"Generic teamspace secret {key!r} was not found.") + self._client.secret_service_delete_secret(project_id=teamspace_id, id=secret.id) + def _get_secrets(self, teamspace_id: str) -> List[V1Secret]: """Fetch all raw secret objects for the teamspace. diff --git a/python/lightning_sdk/api/user_api.py b/python/lightning_sdk/api/user_api.py index e6622901..a3faecc6 100644 --- a/python/lightning_sdk/api/user_api.py +++ b/python/lightning_sdk/api/user_api.py @@ -19,6 +19,10 @@ from lightning_sdk.lightning_cloud.rest_client import LightningClient +def _is_generic_secret(secret: V1Secret) -> bool: + return secret.type in (None, V1SecretType.UNSPECIFIED) + + class UserApi: """Internal API Client for user requests (mainly http requests).""" @@ -144,7 +148,7 @@ def get_secrets(self) -> Dict[str, str]: # this returns encrypted values for security. It doesn't make sense to show them, # so we just return a placeholder # not a security issue to replace in the client as we get the encrypted values from the server. - return {secret.name: "***REDACTED***" for secret in secrets if secret.type == V1SecretType.UNSPECIFIED} + return {secret.name: "***REDACTED***" for secret in secrets if _is_generic_secret(secret)} def set_secret(self, key: str, value: str) -> None: """Create or replace an encrypted secret for the authenticated user. @@ -155,10 +159,27 @@ def set_secret(self, key: str, value: str) -> None: """ secrets = self._get_secrets() for secret in secrets: - if secret.name == key: + if secret.name == key and _is_generic_secret(secret): return self._update_secret(secret.id, value) return self._create_secret(key, value) + def delete_secret(self, key: str) -> None: + """Delete a generic secret for the authenticated user. + + Args: + key: The generic secret name. + + Raises: + ValueError: If no generic secret with the given name exists. + """ + secret = next( + (secret for secret in self._get_secrets() if secret.name == key and _is_generic_secret(secret)), + None, + ) + if secret is None: + raise ValueError(f"Generic user secret {key!r} was not found.") + self._client.secret_service_delete_user_secret(id=secret.id) + def _get_secrets(self) -> List[V1Secret]: """Fetch all raw secret objects for the authenticated user. diff --git a/python/lightning_sdk/cli/deployment/__init__.py b/python/lightning_sdk/cli/deployment/__init__.py index 75e75fb0..ff089300 100644 --- a/python/lightning_sdk/cli/deployment/__init__.py +++ b/python/lightning_sdk/cli/deployment/__init__.py @@ -7,6 +7,7 @@ def register_commands(group: click.Group) -> None: """Register deployment commands with the given group.""" from lightning_sdk.cli.deployment.create import create_deployment from lightning_sdk.cli.deployment.delete import delete_deployment + from lightning_sdk.cli.deployment.env import env from lightning_sdk.cli.deployment.inspect import inspect_deployment from lightning_sdk.cli.deployment.list import list_deployments from lightning_sdk.cli.deployment.logs import deployment_logs @@ -18,5 +19,6 @@ def register_commands(group: click.Group) -> None: group.add_command(inspect_deployment, name="inspect") group.add_command(update_deployment, name="update") group.add_command(delete_deployment, name="delete") + group.add_command(env) group.add_command(deployment_logs, name="logs") group.add_command(reload_weights, name="reload-weights") diff --git a/python/lightning_sdk/cli/deployment/env.py b/python/lightning_sdk/cli/deployment/env.py new file mode 100644 index 00000000..1984ddb5 --- /dev/null +++ b/python/lightning_sdk/cli/deployment/env.py @@ -0,0 +1,98 @@ +"""Deployment environment-variable commands.""" + +from typing import Optional + +import rich_click as click +from rich.table import Table + +from lightning_sdk.cli.utils.configuration import ( + deployment_environment_records, + parse_assignment, + validate_name, +) +from lightning_sdk.cli.utils.json_output import echo_json +from lightning_sdk.cli.utils.logging import LightningCommand, LightningGroup +from lightning_sdk.cli.utils.resource_resolution import resolve_deployment +from lightning_sdk.cli.utils.richt_print import rich_to_str +from lightning_sdk.cli.utils.teamspace_option import resolve_teamspace, teamspace_option + + +@click.group("env", cls=LightningGroup) +def env() -> None: + """Manage Deployment environment variables.""" + + +@env.command("list", cls=LightningCommand) +@click.option("--name", required=True, help="Deployment name.") +@teamspace_option +@click.option("--json", "as_json", is_flag=True, default=False, help="Output as JSON.") +def list_env( + name: str, + teamspace: Optional[str], + org: Optional[str], + user: Optional[str], + as_json: bool, +) -> None: + """List directly configured Deployment environment variables.""" + resolved_teamspace = resolve_teamspace(teamspace=teamspace, org=org, user=user) + deployment = resolve_deployment(name, resolved_teamspace) + records = deployment_environment_records(deployment.env or []) + if as_json: + echo_json({"environment_variables": records}) + return + + table = Table() + table.add_column("Name") + table.add_column("Value") + table.add_column("From secret") + for record in records: + table.add_row(record["name"], record.get("value", ""), record.get("from_secret", "")) + click.echo(rich_to_str(table), color=True) + + +@env.command("set", cls=LightningCommand) +@click.argument("assignment") +@click.option("--name", required=True, help="Deployment name.") +@teamspace_option +@click.option("--json", "as_json", is_flag=True, default=False, help="Output as JSON.") +def set_env( + assignment: str, + name: str, + teamspace: Optional[str], + org: Optional[str], + user: Optional[str], + as_json: bool, +) -> None: + """Set one Deployment environment variable from KEY=VALUE.""" + key, value = parse_assignment(assignment) + resolved_teamspace = resolve_teamspace(teamspace=teamspace, org=org, user=user) + deployment = resolve_deployment(name, resolved_teamspace) + deployment.set_env({key: value}) + if as_json: + echo_json({"name": key, "status": "set"}) + return + click.echo(f"Environment variable {key} set.") + + +@env.command("delete", cls=LightningCommand) +@click.argument("key") +@click.option("--name", required=True, help="Deployment name.") +@teamspace_option +@click.option("--json", "as_json", is_flag=True, default=False, help="Output as JSON.") +def delete_env( + key: str, + name: str, + teamspace: Optional[str], + org: Optional[str], + user: Optional[str], + as_json: bool, +) -> None: + """Delete one directly configured Deployment environment variable.""" + validate_name(key) + resolved_teamspace = resolve_teamspace(teamspace=teamspace, org=org, user=user) + deployment = resolve_deployment(name, resolved_teamspace) + deployment.delete_env(key) + if as_json: + echo_json({"name": key, "status": "deleted"}) + return + click.echo(f"Environment variable {key} deleted.") diff --git a/python/lightning_sdk/cli/entrypoint.py b/python/lightning_sdk/cli/entrypoint.py index bf7ab6e3..606b1f36 100644 --- a/python/lightning_sdk/cli/entrypoint.py +++ b/python/lightning_sdk/cli/entrypoint.py @@ -29,6 +29,8 @@ sandbox, ssh, studio, + teamspace, + user, ) from lightning_sdk.cli.legacy_redirects import ( build_hidden_alias_group, @@ -44,7 +46,7 @@ {"name": "GET STARTED", "commands": ["login", "logout", "config"]}, {"name": "COMPUTE", "commands": ["studio", "base-studio", "machine", "container", "sandbox"]}, {"name": "TRAIN & DEPLOY", "commands": ["job", "mmt", "model", "deployment"]}, - {"name": "ACCESS", "commands": ["auth", "api-key", "ssh", "license"]}, + {"name": "ACCESS", "commands": ["user", "teamspace", "auth", "api-key", "ssh", "license"]}, {"name": "DATA & FILES", "commands": ["cp"]}, ] } @@ -131,6 +133,8 @@ def logout() -> None: main_cli.add_command(folder) main_cli.add_command(ssh) main_cli.add_command(studio) +main_cli.add_command(user) +main_cli.add_command(teamspace) main_cli.add_command(sandbox) main_cli.add_command(base_studio) main_cli.add_command(dataset) diff --git a/python/lightning_sdk/cli/groups.py b/python/lightning_sdk/cli/groups.py index d88fc076..f7b982f1 100644 --- a/python/lightning_sdk/cli/groups.py +++ b/python/lightning_sdk/cli/groups.py @@ -23,6 +23,8 @@ from lightning_sdk.cli.sandbox import register_commands as register_sandbox_commands from lightning_sdk.cli.ssh import register_commands as register_ssh_commands from lightning_sdk.cli.studio import register_commands as register_studio_commands +from lightning_sdk.cli.teamspace import register_commands as register_teamspace_commands +from lightning_sdk.cli.user import register_commands as register_user_commands from lightning_sdk.cli.utils.logging import LightningCommand, LightningGroup @@ -31,6 +33,16 @@ def studio() -> None: """Persistent GPU dev workspaces.""" +@click.group(name="user", cls=LightningGroup) +def user() -> None: + """Manage the authenticated user.""" + + +@click.group(name="teamspace", cls=LightningGroup) +def teamspace() -> None: + """Manage Lightning teamspaces.""" + + @click.group(name="job", cls=LightningGroup) def job() -> None: """Run batch jobs and sweeps.""" @@ -192,6 +204,8 @@ def cp() -> None: register_mmt_commands(mmt) register_machine_commands(machine) register_studio_commands(studio) +register_user_commands(user) +register_teamspace_commands(teamspace) register_config_commands(config) register_api_commands(api) register_deployment_commands(deployment) diff --git a/python/lightning_sdk/cli/studio/__init__.py b/python/lightning_sdk/cli/studio/__init__.py index 74beab8a..87720075 100644 --- a/python/lightning_sdk/cli/studio/__init__.py +++ b/python/lightning_sdk/cli/studio/__init__.py @@ -9,6 +9,7 @@ def register_commands(group: click.Group) -> None: from lightning_sdk.cli.studio.cp import cp_studio_file from lightning_sdk.cli.studio.create import create_studio from lightning_sdk.cli.studio.delete import delete_studio + from lightning_sdk.cli.studio.env import env from lightning_sdk.cli.studio.list import list_studios from lightning_sdk.cli.studio.ls import ls_studio from lightning_sdk.cli.studio.open import open_studio @@ -19,6 +20,7 @@ def register_commands(group: click.Group) -> None: from lightning_sdk.cli.studio.switch import switch_studio group.add_command(delete_studio) + group.add_command(env) group.add_command(create_studio) group.add_command(list_studios) group.add_command(ssh_studio) diff --git a/python/lightning_sdk/cli/studio/env.py b/python/lightning_sdk/cli/studio/env.py new file mode 100644 index 00000000..f8ff0747 --- /dev/null +++ b/python/lightning_sdk/cli/studio/env.py @@ -0,0 +1,93 @@ +"""Studio environment-variable commands.""" + +from typing import Optional + +import rich_click as click +from rich.table import Table + +from lightning_sdk.cli.utils.configuration import environment_records, parse_assignment, validate_name +from lightning_sdk.cli.utils.json_output import echo_json +from lightning_sdk.cli.utils.logging import LightningCommand, LightningGroup +from lightning_sdk.cli.utils.resource_resolution import resolve_studio +from lightning_sdk.cli.utils.richt_print import rich_to_str +from lightning_sdk.cli.utils.teamspace_option import resolve_teamspace, teamspace_option + + +@click.group("env", cls=LightningGroup) +def env() -> None: + """Manage Studio environment variables.""" + + +@env.command("list", cls=LightningCommand) +@click.option("--name", help="Studio to use. Falls back to the current Studio or configured default.") +@teamspace_option +@click.option("--json", "as_json", is_flag=True, default=False, help="Output as JSON.") +def list_env( + name: Optional[str], + teamspace: Optional[str], + org: Optional[str], + user: Optional[str], + as_json: bool, +) -> None: + """List directly configured Studio environment variables.""" + resolved_teamspace = resolve_teamspace(teamspace=teamspace, org=org, user=user) + studio = resolve_studio(name, resolved_teamspace) + records = environment_records(studio.env) + if as_json: + echo_json({"environment_variables": records}) + return + + table = Table() + table.add_column("Name") + table.add_column("Value") + for record in records: + table.add_row(record["name"], record["value"]) + click.echo(rich_to_str(table), color=True) + + +@env.command("set", cls=LightningCommand) +@click.argument("assignment") +@click.option("--name", help="Studio to use. Falls back to the current Studio or configured default.") +@teamspace_option +@click.option("--json", "as_json", is_flag=True, default=False, help="Output as JSON.") +def set_env( + assignment: str, + name: Optional[str], + teamspace: Optional[str], + org: Optional[str], + user: Optional[str], + as_json: bool, +) -> None: + """Set one Studio environment variable from KEY=VALUE.""" + key, value = parse_assignment(assignment) + resolved_teamspace = resolve_teamspace(teamspace=teamspace, org=org, user=user) + studio = resolve_studio(name, resolved_teamspace) + studio.set_env({key: value}) + if as_json: + echo_json({"name": key, "status": "set"}) + return + click.echo(f"Environment variable {key} set.") + + +@env.command("delete", cls=LightningCommand) +@click.argument("key") +@click.option("--name", help="Studio to use. Falls back to the current Studio or configured default.") +@teamspace_option +@click.option("--json", "as_json", is_flag=True, default=False, help="Output as JSON.") +def delete_env( + key: str, + name: Optional[str], + teamspace: Optional[str], + org: Optional[str], + user: Optional[str], + as_json: bool, +) -> None: + """Delete one directly configured Studio environment variable.""" + validate_name(key) + resolved_teamspace = resolve_teamspace(teamspace=teamspace, org=org, user=user) + studio = resolve_studio(name, resolved_teamspace) + studio.delete_env(key) + if as_json: + echo_json({"name": key, "status": "deleted"}) + return + click.echo(f"Environment variable {key} deleted.") diff --git a/python/lightning_sdk/cli/teamspace/__init__.py b/python/lightning_sdk/cli/teamspace/__init__.py new file mode 100644 index 00000000..314d8caf --- /dev/null +++ b/python/lightning_sdk/cli/teamspace/__init__.py @@ -0,0 +1,10 @@ +"""Teamspace CLI commands.""" + +import rich_click as click + + +def register_commands(group: click.Group) -> None: + """Register Teamspace commands.""" + from lightning_sdk.cli.teamspace.secret import secret + + group.add_command(secret) diff --git a/python/lightning_sdk/cli/teamspace/secret.py b/python/lightning_sdk/cli/teamspace/secret.py new file mode 100644 index 00000000..18754a5e --- /dev/null +++ b/python/lightning_sdk/cli/teamspace/secret.py @@ -0,0 +1,81 @@ +"""Generic Teamspace secret commands.""" + +from typing import Optional + +import rich_click as click +from rich.table import Table + +from lightning_sdk.cli.utils.configuration import read_secret_value, secret_records, validate_name +from lightning_sdk.cli.utils.json_output import echo_json +from lightning_sdk.cli.utils.logging import LightningCommand, LightningGroup +from lightning_sdk.cli.utils.richt_print import rich_to_str +from lightning_sdk.cli.utils.teamspace_option import resolve_teamspace, teamspace_option + + +@click.group("secret", cls=LightningGroup) +def secret() -> None: + """Manage generic Teamspace secrets.""" + + +@secret.command("list", cls=LightningCommand) +@teamspace_option +@click.option("--json", "as_json", is_flag=True, default=False, help="Output as JSON.") +def list_secrets(teamspace: Optional[str], org: Optional[str], user: Optional[str], as_json: bool) -> None: + """List generic Teamspace secrets.""" + resolved = resolve_teamspace(teamspace=teamspace, org=org, user=user) + records = secret_records(resolved.secrets) + if as_json: + echo_json({"secrets": records}) + return + + table = Table() + table.add_column("Name") + table.add_column("Value") + for record in records: + table.add_row(record["name"], record["value"]) + click.echo(rich_to_str(table), color=True) + + +@secret.command("set", cls=LightningCommand) +@click.argument("key") +@teamspace_option +@click.option("--value-stdin", is_flag=True, default=False, help="Read the secret value from stdin.") +@click.option("--json", "as_json", is_flag=True, default=False, help="Output as JSON.") +def set_secret( + key: str, + teamspace: Optional[str], + org: Optional[str], + user: Optional[str], + value_stdin: bool, + as_json: bool, +) -> None: + """Create or update a generic Teamspace secret.""" + validate_name(key) + value = read_secret_value(value_stdin) + resolved = resolve_teamspace(teamspace=teamspace, org=org, user=user) + resolved.set_secret(key, value) + if as_json: + echo_json({"name": key, "status": "set"}) + return + click.echo(f"Secret {key} set.") + + +@secret.command("delete", cls=LightningCommand) +@click.argument("key") +@teamspace_option +@click.option("--json", "as_json", is_flag=True, default=False, help="Output as JSON.") +def delete_secret( + key: str, + teamspace: Optional[str], + org: Optional[str], + user: Optional[str], + as_json: bool, +) -> None: + """Delete a generic Teamspace secret.""" + validate_name(key) + resolved = resolve_teamspace(teamspace=teamspace, org=org, user=user) + resolved.delete_secret(key) + if as_json: + echo_json({"name": key, "status": "deleted"}) + return + click.echo(f"Secret {key} deleted.") diff --git a/python/lightning_sdk/cli/user/__init__.py b/python/lightning_sdk/cli/user/__init__.py new file mode 100644 index 00000000..fce46e82 --- /dev/null +++ b/python/lightning_sdk/cli/user/__init__.py @@ -0,0 +1,10 @@ +"""Authenticated-user CLI commands.""" + +import rich_click as click + + +def register_commands(group: click.Group) -> None: + """Register authenticated-user commands.""" + from lightning_sdk.cli.user.secret import secret + + group.add_command(secret) diff --git a/python/lightning_sdk/cli/user/secret.py b/python/lightning_sdk/cli/user/secret.py new file mode 100644 index 00000000..5f01d6ad --- /dev/null +++ b/python/lightning_sdk/cli/user/secret.py @@ -0,0 +1,60 @@ +"""Generic secret commands for the authenticated user.""" + +import rich_click as click +from rich.table import Table + +from lightning_sdk.cli.utils.configuration import read_secret_value, secret_records, validate_name +from lightning_sdk.cli.utils.json_output import echo_json +from lightning_sdk.cli.utils.logging import LightningCommand, LightningGroup +from lightning_sdk.cli.utils.richt_print import rich_to_str +from lightning_sdk.utils.resolve import _get_authed_user + + +@click.group("secret", cls=LightningGroup) +def secret() -> None: + """Manage generic secrets for the authenticated user.""" + + +@secret.command("list", cls=LightningCommand) +@click.option("--json", "as_json", is_flag=True, default=False, help="Output as JSON.") +def list_secrets(as_json: bool) -> None: + """List generic user secrets.""" + records = secret_records(_get_authed_user().secrets) + if as_json: + echo_json({"secrets": records}) + return + + table = Table() + table.add_column("Name") + table.add_column("Value") + for record in records: + table.add_row(record["name"], record["value"]) + click.echo(rich_to_str(table), color=True) + + +@secret.command("set", cls=LightningCommand) +@click.argument("key") +@click.option("--value-stdin", is_flag=True, default=False, help="Read the secret value from stdin.") +@click.option("--json", "as_json", is_flag=True, default=False, help="Output as JSON.") +def set_secret(key: str, value_stdin: bool, as_json: bool) -> None: + """Create or update a generic user secret.""" + validate_name(key) + value = read_secret_value(value_stdin) + _get_authed_user().set_secret(key, value) + if as_json: + echo_json({"name": key, "status": "set"}) + return + click.echo(f"Secret {key} set.") + + +@secret.command("delete", cls=LightningCommand) +@click.argument("key") +@click.option("--json", "as_json", is_flag=True, default=False, help="Output as JSON.") +def delete_secret(key: str, as_json: bool) -> None: + """Delete a generic user secret.""" + validate_name(key) + _get_authed_user().delete_secret(key) + if as_json: + echo_json({"name": key, "status": "deleted"}) + return + click.echo(f"Secret {key} deleted.") diff --git a/python/lightning_sdk/cli/utils/configuration.py b/python/lightning_sdk/cli/utils/configuration.py new file mode 100644 index 00000000..400f9698 --- /dev/null +++ b/python/lightning_sdk/cli/utils/configuration.py @@ -0,0 +1,59 @@ +"""Shared parsing and rendering records for environment and secret commands.""" + +import re +import sys +from typing import Dict, List, Tuple, Union + +import rich_click as click + +from lightning_sdk.api.deployment_api import Env, Secret + +VARIABLE_NAME_PATTERN = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") + + +def validate_name(name: str) -> None: + """Validate an environment variable or generic secret name.""" + if VARIABLE_NAME_PATTERN.fullmatch(name) is None: + raise click.BadParameter( + "Names must start with a letter or underscore and contain only letters, numbers, and underscores." + ) + + +def parse_assignment(assignment: str) -> Tuple[str, str]: + """Parse a ``KEY=VALUE`` assignment, preserving empty values.""" + if "=" not in assignment: + raise click.BadParameter("Expected KEY=VALUE.") + key, value = assignment.split("=", 1) + validate_name(key) + return key, value + + +def read_secret_value(value_stdin: bool) -> str: + """Read a secret from stdin or a hidden interactive prompt.""" + value = sys.stdin.read() if value_stdin else click.prompt("Secret value", hide_input=True, err=True) + if value_stdin: + value = value.removesuffix("\n").removesuffix("\r") + if value == "": + raise click.UsageError("Secret value cannot be empty.") + return value + + +def secret_records(secrets: Dict[str, str]) -> List[Dict[str, str]]: + """Build sorted, always-redacted generic secret records.""" + return [{"name": name, "value": "***REDACTED***"} for name in sorted(secrets)] + + +def environment_records(env: Dict[str, str]) -> List[Dict[str, str]]: + """Build sorted literal environment-variable records.""" + return [{"name": name, "value": env[name]} for name in sorted(env)] + + +def deployment_environment_records(entries: List[Union[Env, Secret]]) -> List[Dict[str, str]]: + """Build sorted records for Deployment literals and secret references.""" + records = [ + {"name": entry.name, "value": entry.value} + if isinstance(entry, Env) + else {"name": entry.name, "from_secret": entry.name} + for entry in entries + ] + return sorted(records, key=lambda record: record["name"]) diff --git a/python/lightning_sdk/cli/utils/resource_resolution.py b/python/lightning_sdk/cli/utils/resource_resolution.py index 6ec19081..b0331fdf 100644 --- a/python/lightning_sdk/cli/utils/resource_resolution.py +++ b/python/lightning_sdk/cli/utils/resource_resolution.py @@ -3,6 +3,7 @@ import rich_click as click from lightning_sdk.api.cloud_account_api import CloudAccountApi +from lightning_sdk.deployment import Deployment from lightning_sdk.job import Job from lightning_sdk.lightning_cloud.openapi import V1ClusterType from lightning_sdk.mmt import MMT @@ -56,6 +57,18 @@ def resolve_studio(name: Optional[str], teamspace: Teamspace) -> Studio: raise click.UsageError(f"Could not resolve studio{detail}. Pass --name STUDIO.") from ex +def resolve_deployment(name: Optional[str], teamspace: Teamspace) -> Deployment: + if not name: + raise click.UsageError("Missing deployment name. Pass --name DEPLOYMENT.") + try: + deployment = Deployment(name=name, teamspace=teamspace) + except ValueError as ex: + raise click.UsageError(f"Could not resolve deployment {name!r} in teamspace {teamspace.name!r}.") from ex + if not deployment._is_created: + raise click.UsageError(f"Could not resolve deployment {name!r} in teamspace {teamspace.name!r}.") + return deployment + + def resolve_job(name: Optional[str], teamspace: Teamspace) -> Job: if not name: raise click.UsageError("Missing job name. Pass JOB.") diff --git a/python/lightning_sdk/deployment.py b/python/lightning_sdk/deployment.py index 0b8401d5..71a026b8 100644 --- a/python/lightning_sdk/deployment.py +++ b/python/lightning_sdk/deployment.py @@ -1,4 +1,5 @@ import os +import re import sys from datetime import datetime from typing import Any, Dict, List, Optional, Union @@ -68,6 +69,17 @@ "TokenAuth", ] +_ENV_NAME_PATTERN = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") +_ENV_NAME_ERROR = ( + "Environment variable names must start with a letter or underscore " + "and contain only letters, numbers, and underscores." +) + + +def _validate_env_names(env: Dict[str, str]) -> None: + if any(_ENV_NAME_PATTERN.fullmatch(name) is None for name in env): + raise ValueError(_ENV_NAME_ERROR) + class Deployment(metaclass=TrackCallsMeta): """The Lightning AI Deployment. @@ -632,6 +644,47 @@ def env(self) -> Optional[List[Union[Secret, Env]]]: return restore_env(self._deployment.spec.env) return None + def set_env(self, new_env: Dict[str, str], partial: bool = True) -> None: + """Set directly configured literal Deployment environment variables. + + Args: + new_env: Literal environment variable names and values. + partial: Preserve unrelated literal variables when ``True``. Secret references + are preserved unless shadowed by a new literal in either mode. + + Raises: + ValueError: If a name is invalid or the Deployment does not exist. + """ + _validate_env_names(new_env) + current = self.env + if current is None: + raise ValueError("Deployment must exist before its environment can be changed.") + + names = set(new_env) + if partial: + kept = [entry for entry in current if entry.name not in names] + else: + kept = [entry for entry in current if isinstance(entry, Secret) and entry.name not in names] + self.update(env=[*kept, *(Env(name=name, value=value) for name, value in new_env.items())]) + + def delete_env(self, key: str) -> None: + """Delete one directly configured literal Deployment environment variable. + + Args: + key: Name of the literal environment variable to delete. + + Raises: + ValueError: If the name is invalid, the Deployment does not exist, or the + literal variable is not directly configured. + """ + _validate_env_names({key: ""}) + current = self.env + if current is None: + raise ValueError("Deployment must exist before its environment can be changed.") + if not any(isinstance(entry, Env) and entry.name == key for entry in current): + raise ValueError(f"Deployment environment variable {key!r} was not found.") + self.update(env=[entry for entry in current if not (isinstance(entry, Env) and entry.name == key)]) + @property def urls(self) -> Optional[List[str]]: """The urls to reach the deployment. diff --git a/python/lightning_sdk/studio.py b/python/lightning_sdk/studio.py index b4b848c2..013af565 100644 --- a/python/lightning_sdk/studio.py +++ b/python/lightning_sdk/studio.py @@ -1,5 +1,6 @@ import glob import os +import re import threading import warnings from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union @@ -34,6 +35,11 @@ from lightning_sdk.mmt import MMT _logger = _setup_logger(__name__) +_ENV_NAME_PATTERN = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") +_ENV_NAME_ERROR = ( + "Environment variable names must start with a letter or underscore " + "and contain only letters, numbers, and underscores." +) class Studio(metaclass=TrackCallsMeta): @@ -973,8 +979,31 @@ def set_env(self, new_env: Dict[str, str], partial: bool = True) -> None: If False, existing environment variables that are not in new_env will be removed. If True, existing environment variables that are not in new_env will be kept. """ + if any(_ENV_NAME_PATTERN.fullmatch(key) is None for key in new_env): + raise ValueError(_ENV_NAME_ERROR) self._studio_api.set_env(self._studio, self._teamspace.id, new_env, partial=partial) + def delete_env(self, key: str) -> None: + """Delete one directly configured Studio environment variable. + + Args: + key: Name of the environment variable to delete. + + Raises: + ValueError: If the name is invalid or is not directly configured. + """ + if _ENV_NAME_PATTERN.fullmatch(key) is None: + raise ValueError(_ENV_NAME_ERROR) + current = self.env + if key not in current: + raise ValueError(f"Studio environment variable {key!r} was not found.") + self._studio_api.set_env( + self._studio, + self._teamspace.id, + {name: value for name, value in current.items() if name != key}, + partial=False, + ) + def __eq__(self, other: "Studio") -> bool: """Checks for equality with other Studios. diff --git a/python/lightning_sdk/teamspace.py b/python/lightning_sdk/teamspace.py index a531e750..bc0d3ba8 100644 --- a/python/lightning_sdk/teamspace.py +++ b/python/lightning_sdk/teamspace.py @@ -358,6 +358,23 @@ def set_secret(self, key: str, value: str, secret_type: Union[str, SecretType] = self._teamspace_api.set_secret(self.id, key, value, secret_type=secret_type) + def delete_secret(self, key: str) -> None: + """Delete a generic encrypted secret from the teamspace. + + Args: + key: Secret name. Must start with a letter or underscore and contain only + alphanumeric characters and underscores. + + Raises: + ValueError: If the key is invalid or no generic secret with that name exists. + """ + if not self._teamspace_api.verify_secret_name(key): + raise ValueError( + "Secret keys must only contain alphanumeric characters and underscores and not begin with a number." + ) + + self._teamspace_api.delete_secret(self.id, key) + def list_machines(self, cloud_account: Optional[str] = None, machine: Optional[str] = None) -> List[Machine]: """List available machines across cloud accounts. diff --git a/python/lightning_sdk/user.py b/python/lightning_sdk/user.py index e2cb1932..de305d36 100644 --- a/python/lightning_sdk/user.py +++ b/python/lightning_sdk/user.py @@ -68,6 +68,23 @@ def set_secret(self, key: str, value: str) -> None: self._user_api.set_secret(key, value) + def delete_secret(self, key: str) -> None: + """Delete an encrypted secret for the user. + + Args: + key: Secret name. Must contain only alphanumeric characters and underscores, + and must not start with a digit. + + Raises: + ValueError: If the key contains invalid characters or does not exist. + """ + if not self._user_api.verify_secret_name(key): + raise ValueError( + "Secret keys must only contain alphanumeric characters and underscores and not begin with a number." + ) + + self._user_api.delete_secret(key) + def create_teamspace(self, name: str) -> "Teamspace": """Create a new teamspace owned by this user. diff --git a/python/tests/api/test_teamspace_api.py b/python/tests/api/test_teamspace_api.py index 9df0d7eb..4b96ef28 100644 --- a/python/tests/api/test_teamspace_api.py +++ b/python/tests/api/test_teamspace_api.py @@ -975,8 +975,8 @@ def test_set_secret_update_existing(): teamspace_api = TeamspaceApi() existing_secrets = [ - V1Secret(id="secret-1", name="API_KEY"), - V1Secret(id="secret-2", name="DATABASE_URL"), + V1Secret(id="secret-1", name="API_KEY", type=V1SecretType.UNSPECIFIED), + V1Secret(id="secret-2", name="DATABASE_URL", type=V1SecretType.UNSPECIFIED), ] with mock.patch.object(teamspace_api, "_get_secrets", return_value=existing_secrets), mock.patch.object( @@ -987,6 +987,53 @@ def test_set_secret_update_existing(): mock_update.assert_called_once_with("ts-abc", "secret-1", "new_secret_value") +@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +def test_set_generic_secret_does_not_update_same_name_typed_secret(): + teamspace_api = TeamspaceApi() + existing_secrets = [ + V1Secret(id="hf-secret", name="API_KEY", type=V1SecretType.HF_TOKEN), + ] + + with mock.patch.object(teamspace_api, "_get_secrets", return_value=existing_secrets), mock.patch.object( + teamspace_api, "_update_secret" + ) as mock_update, mock.patch.object(teamspace_api, "_create_secret") as mock_create: + teamspace_api.set_secret("ts-abc", "API_KEY", "new_secret_value") + + mock_update.assert_not_called() + mock_create.assert_called_once_with("ts-abc", "API_KEY", "new_secret_value", secret_type=V1SecretType.UNSPECIFIED) + + +@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +def test_delete_secret_resolves_generic_secret_id(): + teamspace_api = TeamspaceApi() + existing_secrets = [ + V1Secret(id="hf-secret", name="API_KEY", type=V1SecretType.HF_TOKEN), + V1Secret(id="generic-secret", name="API_KEY", type=V1SecretType.UNSPECIFIED), + ] + + with mock.patch.object(teamspace_api, "_get_secrets", return_value=existing_secrets), mock.patch.object( + teamspace_api._client, "secret_service_delete_secret" + ) as mock_delete: + teamspace_api.delete_secret("ts-abc", "API_KEY") + + mock_delete.assert_called_once_with(project_id="ts-abc", id="generic-secret") + + +@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +def test_delete_secret_rejects_missing_generic_secret(): + teamspace_api = TeamspaceApi() + existing_secrets = [ + V1Secret(id="hf-secret", name="API_KEY", type=V1SecretType.HF_TOKEN), + ] + + with mock.patch.object(teamspace_api, "_get_secrets", return_value=existing_secrets), mock.patch.object( + teamspace_api._client, "secret_service_delete_secret" + ) as mock_delete, pytest.raises(ValueError, match="Generic teamspace secret 'API_KEY' was not found"): + teamspace_api.delete_secret("ts-abc", "API_KEY") + + mock_delete.assert_not_called() + + @mock.patch("lightning_sdk.api.teamspace_api.LightningClient") @mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) def test_get_secrets_api_call(mock_client): diff --git a/python/tests/api/test_user_api.py b/python/tests/api/test_user_api.py index e8ee2d48..f21a8b06 100644 --- a/python/tests/api/test_user_api.py +++ b/python/tests/api/test_user_api.py @@ -82,8 +82,8 @@ def test_set_secret_update_existing(): user_api = UserApi() existing_secrets = [ - V1Secret(id="secret-1", name="API_KEY"), - V1Secret(id="secret-2", name="DATABASE_URL"), + V1Secret(id="secret-1", name="API_KEY", type=V1SecretType.UNSPECIFIED), + V1Secret(id="secret-2", name="DATABASE_URL", type=V1SecretType.UNSPECIFIED), ] with mock.patch.object(user_api, "_get_secrets", return_value=existing_secrets), mock.patch.object( @@ -94,6 +94,53 @@ def test_set_secret_update_existing(): mock_update.assert_called_once_with("secret-1", "new_secret_value") +@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +def test_set_secret_does_not_update_same_name_typed_secret(): + user_api = UserApi() + existing_secrets = [ + V1Secret(id="docker-secret", name="API_KEY", type=V1SecretType.DOCKER_REGISTRY), + ] + + with mock.patch.object(user_api, "_get_secrets", return_value=existing_secrets), mock.patch.object( + user_api, "_update_secret" + ) as mock_update, mock.patch.object(user_api, "_create_secret") as mock_create: + user_api.set_secret("API_KEY", "new_secret_value") + + mock_update.assert_not_called() + mock_create.assert_called_once_with("API_KEY", "new_secret_value") + + +@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +def test_delete_secret_resolves_generic_secret_id(): + user_api = UserApi() + existing_secrets = [ + V1Secret(id="docker-secret", name="API_KEY", type=V1SecretType.DOCKER_REGISTRY), + V1Secret(id="generic-secret", name="API_KEY", type=V1SecretType.UNSPECIFIED), + ] + + with mock.patch.object(user_api, "_get_secrets", return_value=existing_secrets), mock.patch.object( + user_api._client, "secret_service_delete_user_secret" + ) as mock_delete: + user_api.delete_secret("API_KEY") + + mock_delete.assert_called_once_with(id="generic-secret") + + +@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +def test_delete_secret_rejects_missing_generic_secret(): + user_api = UserApi() + existing_secrets = [ + V1Secret(id="docker-secret", name="API_KEY", type=V1SecretType.DOCKER_REGISTRY), + ] + + with mock.patch.object(user_api, "_get_secrets", return_value=existing_secrets), mock.patch.object( + user_api._client, "secret_service_delete_user_secret" + ) as mock_delete, pytest.raises(ValueError, match="Generic user secret 'API_KEY' was not found"): + user_api.delete_secret("API_KEY") + + mock_delete.assert_not_called() + + @mock.patch("lightning_sdk.api.user_api.LightningClient") @mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) def test_get_secrets_api_call(mock_client): diff --git a/python/tests/cli/deployment/test_env.py b/python/tests/cli/deployment/test_env.py new file mode 100644 index 00000000..42887896 --- /dev/null +++ b/python/tests/cli/deployment/test_env.py @@ -0,0 +1,62 @@ +import json +from unittest.mock import MagicMock, call, patch + +from click.testing import CliRunner + +from lightning_sdk.api.deployment_api import Env, Secret +from lightning_sdk.cli.entrypoint import main_cli + + +def test_deployment_env_list_renders_literals_and_secret_references(): + teamspace = MagicMock() + deployment = MagicMock() + deployment.env = [Secret("TOKEN"), Env("DEBUG", "true")] + + with patch("lightning_sdk.cli.deployment.env.resolve_teamspace", return_value=teamspace), patch( + "lightning_sdk.cli.deployment.env.resolve_deployment", return_value=deployment + ) as resolve_deployment: + result = CliRunner().invoke( + main_cli, + ["deployment", "env", "list", "--name", "serve", "--teamspace", "acme/research", "--json"], + ) + + assert result.exit_code == 0, result.output + assert json.loads(result.output) == { + "environment_variables": [ + {"name": "DEBUG", "value": "true"}, + {"from_secret": "TOKEN", "name": "TOKEN"}, + ] + } + resolve_deployment.assert_called_once_with("serve", teamspace) + + +def test_deployment_env_set_and_delete_delegate_to_resource(): + teamspace = MagicMock() + deployment = MagicMock() + + with patch("lightning_sdk.cli.deployment.env.resolve_teamspace", return_value=teamspace), patch( + "lightning_sdk.cli.deployment.env.resolve_deployment", return_value=deployment + ) as resolve_deployment: + set_result = CliRunner().invoke( + main_cli, + ["deployment", "env", "set", "DEBUG=", "--name", "serve", "--teamspace", "acme/research", "--json"], + ) + delete_result = CliRunner().invoke( + main_cli, + ["deployment", "env", "delete", "DEBUG", "--name", "serve", "--teamspace", "acme/research", "--json"], + ) + + assert set_result.exit_code == 0, set_result.output + assert delete_result.exit_code == 0, delete_result.output + deployment.set_env.assert_called_once_with({"DEBUG": ""}) + deployment.delete_env.assert_called_once_with("DEBUG") + assert resolve_deployment.call_args_list == [call("serve", teamspace), call("serve", teamspace)] + + +def test_deployment_env_requires_name_and_all_lists_offer_json(): + missing = CliRunner().invoke(main_cli, ["deployment", "env", "list"]) + + assert missing.exit_code != 0 + assert "--name" in missing.output + assert "--json" in CliRunner().invoke(main_cli, ["deployment", "env", "list", "--help"]).output + assert "--json" in CliRunner().invoke(main_cli, ["studio", "env", "list", "--help"]).output diff --git a/python/tests/cli/studio/test_env.py b/python/tests/cli/studio/test_env.py new file mode 100644 index 00000000..ec3afc61 --- /dev/null +++ b/python/tests/cli/studio/test_env.py @@ -0,0 +1,57 @@ +import json +from unittest.mock import MagicMock, call, patch + +from click.testing import CliRunner + +from lightning_sdk.cli.entrypoint import main_cli + + +def test_studio_env_list_set_and_delete_resolve_resource(): + teamspace = MagicMock() + studio = MagicMock() + studio.env = {"Z_VALUE": "last", "DEBUG": "true"} + + with patch("lightning_sdk.cli.studio.env.resolve_teamspace", return_value=teamspace) as resolve_teamspace, patch( + "lightning_sdk.cli.studio.env.resolve_studio", return_value=studio + ) as resolve_studio: + listed = CliRunner().invoke( + main_cli, ["studio", "env", "list", "--name", "dev", "--teamspace", "acme/research", "--json"] + ) + set_result = CliRunner().invoke( + main_cli, ["studio", "env", "set", "DEBUG=", "--name", "dev", "--teamspace", "acme/research", "--json"] + ) + deleted = CliRunner().invoke( + main_cli, + ["studio", "env", "delete", "DEBUG", "--name", "dev", "--teamspace", "acme/research", "--json"], + ) + + assert listed.exit_code == 0, listed.output + assert json.loads(listed.output) == { + "environment_variables": [ + {"name": "DEBUG", "value": "true"}, + {"name": "Z_VALUE", "value": "last"}, + ] + } + assert json.loads(set_result.output) == {"name": "DEBUG", "status": "set"} + assert json.loads(deleted.output) == {"name": "DEBUG", "status": "deleted"} + studio.set_env.assert_called_once_with({"DEBUG": ""}) + studio.delete_env.assert_called_once_with("DEBUG") + assert resolve_teamspace.call_count == 3 + assert resolve_studio.call_args_list == [call("dev", teamspace), call("dev", teamspace), call("dev", teamspace)] + + +def test_studio_env_omitted_name_uses_context_and_invalid_assignment_fails_first(): + teamspace = MagicMock() + studio = MagicMock() + studio.env = {} + + with patch("lightning_sdk.cli.studio.env.resolve_teamspace", return_value=teamspace), patch( + "lightning_sdk.cli.studio.env.resolve_studio", return_value=studio + ) as resolve_studio: + listed = CliRunner().invoke(main_cli, ["studio", "env", "list", "--json"]) + invalid = CliRunner().invoke(main_cli, ["studio", "env", "set", "INVALID"]) + + assert listed.exit_code == 0, listed.output + resolve_studio.assert_called_once_with(None, teamspace) + assert invalid.exit_code != 0 + studio.set_env.assert_not_called() diff --git a/python/tests/cli/teamspace/__init__.py b/python/tests/cli/teamspace/__init__.py new file mode 100644 index 00000000..4653f632 --- /dev/null +++ b/python/tests/cli/teamspace/__init__.py @@ -0,0 +1 @@ +"""Teamspace CLI tests.""" diff --git a/python/tests/cli/teamspace/test_secret.py b/python/tests/cli/teamspace/test_secret.py new file mode 100644 index 00000000..4bd73b0f --- /dev/null +++ b/python/tests/cli/teamspace/test_secret.py @@ -0,0 +1,51 @@ +import json +from unittest.mock import MagicMock, patch + +from click.testing import CliRunner + +from lightning_sdk.cli.entrypoint import main_cli + + +def test_teamspace_secret_commands_resolve_slug_and_redact_lists(): + teamspace = MagicMock() + teamspace.secrets = {"TOKEN": "sdk-secret-canary-7f3a"} + + with patch("lightning_sdk.cli.teamspace.secret.resolve_teamspace", return_value=teamspace) as resolve: + listed = CliRunner().invoke(main_cli, ["teamspace", "secret", "list", "--teamspace", "acme/research", "--json"]) + deleted = CliRunner().invoke( + main_cli, ["teamspace", "secret", "delete", "TOKEN", "--teamspace", "acme/research", "--json"] + ) + + assert listed.exit_code == 0, listed.output + assert json.loads(listed.output) == {"secrets": [{"name": "TOKEN", "value": "***REDACTED***"}]} + assert "sdk-secret-canary-7f3a" not in listed.output + assert json.loads(deleted.output) == {"name": "TOKEN", "status": "deleted"} + teamspace.delete_secret.assert_called_once_with("TOKEN") + assert resolve.call_count == 2 + resolve.assert_called_with(teamspace="acme/research", org=None, user=None) + + +def test_teamspace_secret_set_reads_stdin_and_never_exposes_value(): + canary = "sdk-secret-canary-7f3a" + teamspace = MagicMock() + + with patch("lightning_sdk.cli.teamspace.secret.resolve_teamspace", return_value=teamspace): + result = CliRunner().invoke( + main_cli, + ["teamspace", "secret", "set", "TOKEN", "--teamspace", "acme/research", "--value-stdin", "--json"], + input=canary + "\n", + ) + + assert result.exit_code == 0, result.output + assert canary not in result.output + teamspace.set_secret.assert_called_once_with("TOKEN", canary) + assert json.loads(result.output) == {"name": "TOKEN", "status": "set"} + + +def test_secret_resource_groups_are_registered(): + assert "user" in main_cli.commands + assert "teamspace" in main_cli.commands + assert "secret" in main_cli.commands["user"].commands + assert "secret" in main_cli.commands["teamspace"].commands + for resource in ("user", "teamspace"): + assert set(main_cli.commands[resource].commands["secret"].commands) == {"list", "set", "delete"} diff --git a/python/tests/cli/user/__init__.py b/python/tests/cli/user/__init__.py new file mode 100644 index 00000000..e5532be0 --- /dev/null +++ b/python/tests/cli/user/__init__.py @@ -0,0 +1 @@ +"""User CLI tests.""" diff --git a/python/tests/cli/user/test_secret.py b/python/tests/cli/user/test_secret.py new file mode 100644 index 00000000..1bdc0832 --- /dev/null +++ b/python/tests/cli/user/test_secret.py @@ -0,0 +1,60 @@ +import json +from unittest.mock import MagicMock, call, patch + +from click.testing import CliRunner + +from lightning_sdk.cli.entrypoint import main_cli + + +def test_user_secret_list_json_is_sorted_and_redacted(): + user = MagicMock() + user.secrets = {"Z_TOKEN": "sdk-secret-canary-7f3a", "A_TOKEN": "encrypted"} + + with patch("lightning_sdk.cli.user.secret._get_authed_user", return_value=user): + result = CliRunner().invoke(main_cli, ["user", "secret", "list", "--json"]) + + assert result.exit_code == 0, result.output + assert json.loads(result.output) == { + "secrets": [ + {"name": "A_TOKEN", "value": "***REDACTED***"}, + {"name": "Z_TOKEN", "value": "***REDACTED***"}, + ] + } + assert "sdk-secret-canary-7f3a" not in result.output + + +def test_user_secret_set_supports_hidden_prompt_and_stdin_without_leaking(): + canary = "sdk-secret-canary-7f3a" + user = MagicMock() + + with patch("lightning_sdk.cli.user.secret._get_authed_user", return_value=user): + prompt_result = CliRunner().invoke(main_cli, ["user", "secret", "set", "TOKEN", "--json"], input=canary + "\n") + stdin_result = CliRunner().invoke( + main_cli, + ["user", "secret", "set", "OTHER_TOKEN", "--value-stdin", "--json"], + input=canary + "\n", + ) + + assert prompt_result.exit_code == 0, prompt_result.output + assert stdin_result.exit_code == 0, stdin_result.output + assert canary not in prompt_result.output + assert canary not in stdin_result.output + assert user.set_secret.call_args_list == [call("TOKEN", canary), call("OTHER_TOKEN", canary)] + assert '"status": "set"' in prompt_result.output + + +def test_user_secret_delete_and_invalid_inputs(): + user = MagicMock() + with patch("lightning_sdk.cli.user.secret._get_authed_user", return_value=user): + deleted = CliRunner().invoke(main_cli, ["user", "secret", "delete", "TOKEN", "--json"]) + invalid_name = CliRunner().invoke(main_cli, ["user", "secret", "delete", "BAD-NAME"]) + value_option = CliRunner().invoke( + main_cli, ["user", "secret", "set", "TOKEN", "--value", "sdk-secret-canary-7f3a"] + ) + + assert deleted.exit_code == 0, deleted.output + assert json.loads(deleted.output) == {"name": "TOKEN", "status": "deleted"} + user.delete_secret.assert_called_once_with("TOKEN") + assert invalid_name.exit_code != 0 + assert value_option.exit_code != 0 + assert "sdk-secret-canary-7f3a" not in value_option.output diff --git a/python/tests/cli/utils/test_configuration.py b/python/tests/cli/utils/test_configuration.py new file mode 100644 index 00000000..dbf83fc3 --- /dev/null +++ b/python/tests/cli/utils/test_configuration.py @@ -0,0 +1,73 @@ +import io + +import pytest +import rich_click as click + +from lightning_sdk.api.deployment_api import Env, Secret +from lightning_sdk.cli.utils.configuration import ( + deployment_environment_records, + environment_records, + parse_assignment, + read_secret_value, + secret_records, +) + + +def test_parse_assignment_splits_once_and_preserves_empty_value(): + assert parse_assignment("TOKEN=part=two") == ("TOKEN", "part=two") + assert parse_assignment("_EMPTY=") == ("_EMPTY", "") + + +@pytest.mark.parametrize("assignment", ["TOKEN", "=value", "2_TOKEN=value", "BAD-NAME=value"]) +def test_parse_assignment_rejects_missing_separator_and_invalid_names(assignment): + with pytest.raises(click.BadParameter): + parse_assignment(assignment) + + +def test_read_secret_value_uses_hidden_prompt(monkeypatch, capsys): + canary = "sdk-secret-canary-7f3a" + prompt = monkeypatch.setattr(click, "prompt", lambda *args, **kwargs: canary) + + assert read_secret_value(False) == canary + captured = capsys.readouterr() + assert canary not in captured.out + assert canary not in captured.err + assert prompt is None + + +def test_read_secret_value_reads_stdin_without_echo(monkeypatch, capsys): + canary = "sdk-secret-canary-7f3a" + monkeypatch.setattr("sys.stdin", io.StringIO(canary + "\r\n")) + + assert read_secret_value(True) == canary + captured = capsys.readouterr() + assert canary not in captured.out + assert canary not in captured.err + + +@pytest.mark.parametrize("value_stdin", [False, True]) +def test_read_secret_value_rejects_empty_input(monkeypatch, value_stdin): + if value_stdin: + monkeypatch.setattr("sys.stdin", io.StringIO("")) + else: + monkeypatch.setattr(click, "prompt", lambda *args, **kwargs: "") + + with pytest.raises(click.UsageError, match="Secret value cannot be empty"): + read_secret_value(value_stdin) + + +def test_configuration_records_are_sorted_and_never_expose_secrets(): + canary = "sdk-secret-canary-7f3a" + + assert secret_records({"Z_TOKEN": canary, "A_TOKEN": "anything"}) == [ + {"name": "A_TOKEN", "value": "***REDACTED***"}, + {"name": "Z_TOKEN", "value": "***REDACTED***"}, + ] + assert environment_records({"Z": "last", "A": ""}) == [ + {"name": "A", "value": ""}, + {"name": "Z", "value": "last"}, + ] + assert deployment_environment_records([Secret("TOKEN"), Env("DEBUG", "true")]) == [ + {"name": "DEBUG", "value": "true"}, + {"from_secret": "TOKEN", "name": "TOKEN"}, + ] diff --git a/python/tests/cli/utils/test_resource_resolution.py b/python/tests/cli/utils/test_resource_resolution.py index a1db9b97..029faa86 100644 --- a/python/tests/cli/utils/test_resource_resolution.py +++ b/python/tests/cli/utils/test_resource_resolution.py @@ -7,6 +7,7 @@ from lightning_sdk.cli.utils.resource_resolution import ( join_teamspace_slug, resolve_cluster, + resolve_deployment, resolve_job, resolve_mmt, resolve_studio, @@ -130,6 +131,35 @@ def test_resolve_studio_converts_not_found_to_usage_error() -> None: resolve_studio("missing", MagicMock()) +def test_resolve_deployment_requires_name() -> None: + with pytest.raises(click.UsageError, match="--name"): + resolve_deployment(None, MagicMock()) + + +def test_resolve_deployment_fetches_existing_deployment() -> None: + teamspace = MagicMock() + resolved = MagicMock() + resolved._is_created = True + with patch( + "lightning_sdk.cli.utils.resource_resolution.Deployment", + return_value=resolved, + ) as deployment: + assert resolve_deployment("api", teamspace) is resolved + deployment.assert_called_once_with(name="api", teamspace=teamspace) + + +def test_resolve_deployment_rejects_uncreated_deployment() -> None: + teamspace = MagicMock(name="teamspace") + teamspace.name = "research" + resolved = MagicMock() + resolved._is_created = False + with patch( + "lightning_sdk.cli.utils.resource_resolution.Deployment", + return_value=resolved, + ), pytest.raises(click.UsageError, match="api"): + resolve_deployment("api", teamspace) + + def test_resolve_job_requires_name() -> None: with pytest.raises(click.UsageError, match="JOB"): resolve_job(None, MagicMock()) diff --git a/python/tests/core/studio/test_studio_environment.py b/python/tests/core/studio/test_studio_environment.py index a44893fb..59d2dfc9 100644 --- a/python/tests/core/studio/test_studio_environment.py +++ b/python/tests/core/studio/test_studio_environment.py @@ -1,5 +1,7 @@ from unittest import mock +import pytest + from lightning_sdk.lightning_cloud.openapi import V1CloudSpace, V1ListCloudSpacesResponse from lightning_sdk.studio import Studio @@ -76,3 +78,58 @@ def test_studio_set_env_partial_false(): studio.set_env(new_env, partial=False) studio._studio_api.set_env.assert_called_once_with(mock_studio, "ts-abc", new_env, partial=False) + + +@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +def test_studio_set_env_validates_names_and_accepts_empty_values(): + studio = Studio.__new__(Studio) + studio._studio = V1CloudSpace(id="st-abc", name="st-abc", cluster_id="c-abc") + studio._studio_api = mock.MagicMock() + studio._teamspace = mock.MagicMock(id="ts-abc") + + new_env = {"_NAME": "", "NAME_2": "value"} + studio.set_env(new_env) + + studio._studio_api.set_env.assert_called_once_with(studio._studio, "ts-abc", new_env, partial=True) + + for invalid_name in ("", "2_NAME", "INVALID-NAME"): + with pytest.raises(ValueError, match="Environment variable names must start"): + studio.set_env({invalid_name: "value"}) + + assert studio._studio_api.set_env.call_count == 1 + + +@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +def test_studio_delete_env_preserves_unrelated_variables(): + studio = Studio.__new__(Studio) + studio._studio = V1CloudSpace(id="st-abc", name="st-abc", cluster_id="c-abc") + studio._studio_api = mock.MagicMock() + studio._teamspace = mock.MagicMock(id="ts-abc") + + with mock.patch.object( + Studio, "env", new_callable=mock.PropertyMock, return_value={"KEEP": "yes", "OLD": "remove"} + ): + studio.delete_env("OLD") + + studio._studio_api.set_env.assert_called_once_with(studio._studio, "ts-abc", {"KEEP": "yes"}, partial=False) + + +@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +def test_studio_delete_env_handles_final_missing_and_invalid_names(): + studio = Studio.__new__(Studio) + studio._studio = V1CloudSpace(id="st-abc", name="st-abc", cluster_id="c-abc") + studio._studio_api = mock.MagicMock() + studio._teamspace = mock.MagicMock(id="ts-abc") + + with mock.patch.object(Studio, "env", new_callable=mock.PropertyMock, return_value={"ONLY": "value"}): + studio.delete_env("ONLY") + studio._studio_api.set_env.assert_called_once_with(studio._studio, "ts-abc", {}, partial=False) + + studio._studio_api.reset_mock() + with mock.patch.object(Studio, "env", new_callable=mock.PropertyMock, return_value={}), pytest.raises( + ValueError, match="Studio environment variable 'MISSING' was not found" + ): + studio.delete_env("MISSING") + with pytest.raises(ValueError, match="Environment variable names must start"): + studio.delete_env("INVALID-NAME") + studio._studio_api.set_env.assert_not_called() diff --git a/python/tests/core/test_deployment.py b/python/tests/core/test_deployment.py index 4677bf4c..602025af 100644 --- a/python/tests/core/test_deployment.py +++ b/python/tests/core/test_deployment.py @@ -1,6 +1,6 @@ import re from types import SimpleNamespace -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock, PropertyMock, patch import pytest @@ -64,6 +64,100 @@ def _mock_deployment_dependencies(monkeypatch, auth_instance): return teamspace_mock +def _assert_env_entries(entries, expected): + actual = [ + ("secret", entry.name, None) if isinstance(entry, Secret) else ("literal", entry.name, entry.value) + for entry in entries + ] + assert actual == expected + + +def test_deployment_set_env_partial_preserves_unrelated_entries(): + deployment = deployment_module.Deployment.__new__(deployment_module.Deployment) + deployment.update = MagicMock() + current = [Env("KEEP", "yes"), Env("CHANGE", "old"), Secret("TOKEN")] + + with patch.object(deployment_module.Deployment, "env", new_callable=PropertyMock, return_value=current): + deployment.set_env({"CHANGE": "new", "EMPTY": ""}) + + entries = deployment.update.call_args.kwargs["env"] + _assert_env_entries( + entries, + [ + ("literal", "KEEP", "yes"), + ("secret", "TOKEN", None), + ("literal", "CHANGE", "new"), + ("literal", "EMPTY", ""), + ], + ) + + +def test_deployment_set_env_replacement_preserves_only_unshadowed_secret_references(): + deployment = deployment_module.Deployment.__new__(deployment_module.Deployment) + deployment.update = MagicMock() + current = [Env("OLD", "value"), Secret("KEEP_SECRET"), Secret("SHADOWED")] + + with patch.object(deployment_module.Deployment, "env", new_callable=PropertyMock, return_value=current): + deployment.set_env({"SHADOWED": "literal"}, partial=False) + + entries = deployment.update.call_args.kwargs["env"] + _assert_env_entries( + entries, + [ + ("secret", "KEEP_SECRET", None), + ("literal", "SHADOWED", "literal"), + ], + ) + + +def test_deployment_delete_env_removes_only_literal_and_can_clear_final_entry(): + deployment = deployment_module.Deployment.__new__(deployment_module.Deployment) + deployment.update = MagicMock() + + current = [Env("TOKEN", "literal"), Secret("TOKEN"), Secret("OTHER")] + with patch.object(deployment_module.Deployment, "env", new_callable=PropertyMock, return_value=current): + deployment.delete_env("TOKEN") + entries = deployment.update.call_args.kwargs["env"] + _assert_env_entries(entries, [("secret", "TOKEN", None), ("secret", "OTHER", None)]) + + deployment.update.reset_mock() + with patch.object( + deployment_module.Deployment, "env", new_callable=PropertyMock, return_value=[Env("ONLY", "value")] + ): + deployment.delete_env("ONLY") + deployment.update.assert_called_once_with(env=[]) + + +def test_deployment_env_mutations_reject_missing_invalid_and_uncreated(): + deployment = deployment_module.Deployment.__new__(deployment_module.Deployment) + deployment.update = MagicMock() + + with patch.object( + deployment_module.Deployment, "env", new_callable=PropertyMock, return_value=[Secret("TOKEN")] + ), pytest.raises(ValueError, match="Deployment environment variable 'TOKEN' was not found"): + deployment.delete_env("TOKEN") + + with patch.object( + deployment_module.Deployment, + "env", + new_callable=PropertyMock, + side_effect=AssertionError("environment should not be read"), + ): + for invalid_name in ("", "2_NAME", "INVALID-NAME"): + with pytest.raises(ValueError, match="Environment variable names must start"): + deployment.set_env({invalid_name: "value"}) + with pytest.raises(ValueError, match="Environment variable names must start"): + deployment.delete_env(invalid_name) + + with patch.object(deployment_module.Deployment, "env", new_callable=PropertyMock, return_value=None): + with pytest.raises(ValueError, match="Deployment must exist before its environment can be changed"): + deployment.set_env({"TOKEN": "value"}) + with pytest.raises(ValueError, match="Deployment must exist before its environment can be changed"): + deployment.delete_env("TOKEN") + + deployment.update.assert_not_called() + + @patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) @patch("lightning_sdk.deployment.login.Auth", new=MagicMock()) def test_unstarted_deployment_id_is_none(monkeypatch): @@ -365,6 +459,11 @@ def test_to_env(): assert env == [V1EnvVar(name="key", value="value"), V1EnvVar(from_secret="secret")] +def test_to_env_distinguishes_no_change_from_clear(): + assert deployment_api_module.to_env(None) is None + assert deployment_api_module.to_env([]) == [] + + @patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=MagicMock()) @patch("lightning_sdk.deployment.login.Auth", new=MagicMock()) def test_to_endpoint(): diff --git a/python/tests/core/test_teamspace.py b/python/tests/core/test_teamspace.py index da06fa5f..1744ae3b 100644 --- a/python/tests/core/test_teamspace.py +++ b/python/tests/core/test_teamspace.py @@ -48,6 +48,7 @@ def __init__(self, id: str) -> None: @pytest.mark.parametrize("user", ["user-abc", None, -1]) @pytest.mark.parametrize("org", ["org-abc", None, -1]) @mock.patch.dict(os.environ, clear=True) +@mock.patch("lightning_sdk.utils.config._DEFAULT_CONFIG_FILE_PATH", "/nonexistent/config.yaml") @mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) def test_teamspace_init( internal_teamspace_api_list_mocker, internal_user_api_mocker, internal_get_org_api_mocker, user, org @@ -77,6 +78,7 @@ def test_teamspace_init( @pytest.mark.parametrize("user", ["user-abc", None, -1]) @pytest.mark.parametrize("org", ["org-abc", None, -1]) +@mock.patch("lightning_sdk.utils.config._DEFAULT_CONFIG_FILE_PATH", "/nonexistent/config.yaml") @mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) def test_teamspace_init_env( internal_teamspace_api_list_mocker, internal_user_api_mocker, internal_get_org_api_mocker, user, org @@ -107,7 +109,6 @@ def test_teamspace_init_env( org = None with context, mock.patch.dict(os.environ, new_dict, clear=True): - print(user, org) Teamspace("ts-abc", user=user, org=org) @@ -858,6 +859,35 @@ def test_teamspace_set_secret_invalid_name( ts.set_secret("123_INVALID", "secret_value") +@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +def test_teamspace_delete_secret( + internal_teamspace_api_list_mocker, + internal_user_api_mocker, +): + ts = Teamspace("ts-abc", user="user-abc") + + with mock.patch.object(ts._teamspace_api, "delete_secret") as mock_delete: + ts.delete_secret("OLD_SECRET") + + mock_delete.assert_called_once_with("ts-abc002", "OLD_SECRET") + + +@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +def test_teamspace_delete_secret_invalid_name( + internal_teamspace_api_list_mocker, + internal_user_api_mocker, +): + ts = Teamspace("ts-abc", user="user-abc") + + with mock.patch.object(ts._teamspace_api, "delete_secret") as mock_delete, pytest.raises( + ValueError, + match="Secret keys must only contain alphanumeric characters and underscores and not begin with a number.", + ): + ts.delete_secret("123_INVALID") + + mock_delete.assert_not_called() + + @mock.patch("lightning_sdk.api.teamspace_api.LightningClient") @mock.patch("lightning_sdk.api.cloud_account_api.LightningClient") @mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) diff --git a/python/tests/core/test_user.py b/python/tests/core/test_user.py index d59eb0d5..71609a99 100644 --- a/python/tests/core/test_user.py +++ b/python/tests/core/test_user.py @@ -97,6 +97,29 @@ def test_user_set_secret_invalid_name(internal_user_api_mocker): user.set_secret("123_INVALID", "secret_value") +@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +def test_user_delete_secret(internal_user_api_mocker): + user = User("my-user-name") + + with mock.patch.object(user._user_api, "delete_secret") as mock_delete: + user.delete_secret("OLD_SECRET") + + mock_delete.assert_called_once_with("OLD_SECRET") + + +@mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock()) +def test_user_delete_secret_invalid_name(internal_user_api_mocker): + user = User("my-user-name") + + with mock.patch.object(user._user_api, "delete_secret") as mock_delete, pytest.raises( + ValueError, + match="Secret keys must only contain alphanumeric characters and underscores and not begin with a number.", + ): + user.delete_secret("123_INVALID") + + mock_delete.assert_not_called() + + @mock.patch("lightning_sdk.teamspace.Teamspace.__init__", return_value=None) @mock.patch("lightning_sdk.user._get_authed_user") @mock.patch("lightning_sdk.lightning_cloud.rest_client.Auth", new=mock.MagicMock())