diff --git a/tests/commands/project/test_envvars.py b/tests/commands/project/test_envvars.py new file mode 100644 index 00000000..6571ba57 --- /dev/null +++ b/tests/commands/project/test_envvars.py @@ -0,0 +1,69 @@ +from contextlib import contextmanager + +import pytest +import requests_mock + +from tests.fixture_data import PROJECT_DATA +from valohai_cli.commands.project.environment_variables.create import create +from valohai_cli.commands.project.environment_variables.delete import delete +from valohai_cli.commands.project.environment_variables.list import list + + +@contextmanager +def env_var_api_mock(): + with requests_mock.mock() as m: + project_url = f"https://app.valohai.com/api/v0/projects/{PROJECT_DATA['id']}/" + m.get(project_url, json=PROJECT_DATA) + m.post(f"{project_url}environment_variable/", json={}) + m.delete(f"{project_url}environment_variable/", json={}) + yield m + + +def test_list_envvars(runner, logged_in_and_linked): + with env_var_api_mock(): + result = runner.invoke(list) + for key, ev in PROJECT_DATA["environment_variables"].items(): + assert key in result.output + assert (ev["value"] if not ev["secret"] else "****") in result.output + + +def test_create_envvar_secret(runner, logged_in_and_linked): + with env_var_api_mock() as m: + result = runner.invoke(create, ["--secret", "hey=ho=lets=go", "eeuu=oouuuhhh"]) + assert result.exit_code == 0 + assert [r.json() for r in m.request_history] == [ + {"name": "hey", "secret": True, "value": "ho=lets=go"}, + {"name": "eeuu", "secret": True, "value": "oouuuhhh"}, + ] + + +def test_create_envvar_nonsecret(runner, logged_in_and_linked): + with env_var_api_mock() as m: + result = runner.invoke(create, ["BANANA_HAMMOCK_ENABLED=true"]) + print(result.output) + assert result.exit_code == 0 + assert [r.json() for r in m.request_history] == [ + {"name": "BANANA_HAMMOCK_ENABLED", "secret": False, "value": "true"}, + ] + + +@pytest.mark.parametrize( + "args", + [ + "foo", + "^38=fbar", + "3aa=kdkd", + ], +) +def test_create_envvar_bad(runner, logged_in_and_linked, args): + result = runner.invoke(create, args) + assert result.exit_code != 0 + + +def test_delete_envvar(runner, logged_in_and_linked): + with env_var_api_mock() as m: + result = runner.invoke(delete, ["hey"]) + assert result.exit_code == 0 + assert [r.json() for r in m.request_history] == [ + {"name": "hey"}, + ] diff --git a/tests/fixture_data.py b/tests/fixture_data.py index d5d16e67..371f0f4f 100644 --- a/tests/fixture_data.py +++ b/tests/fixture_data.py @@ -19,6 +19,16 @@ "urls": { "display": "https://app.valohai.com/p/nyan/nyan/", }, + "environment_variables": { + "test1": { + "value": "val1", + "secret": False, + }, + "ssshhh": { + "value": None, + "secret": True, + }, + }, } execution_id = str(uuid.uuid4()) diff --git a/valohai_cli/commands/project/environment_variables/__init__.py b/valohai_cli/commands/project/environment_variables/__init__.py new file mode 100644 index 00000000..a147bc0c --- /dev/null +++ b/valohai_cli/commands/project/environment_variables/__init__.py @@ -0,0 +1,10 @@ +import click + +from valohai_cli.plugin_cli import PluginCLI + + +@click.command(cls=PluginCLI, commands_module="valohai_cli.commands.project.environment_variables") +def environment_variables() -> None: + """ + Project environment variable related commands. + """ diff --git a/valohai_cli/commands/project/environment_variables/create.py b/valohai_cli/commands/project/environment_variables/create.py new file mode 100644 index 00000000..b45cb485 --- /dev/null +++ b/valohai_cli/commands/project/environment_variables/create.py @@ -0,0 +1,69 @@ +from __future__ import annotations + +import re +from typing import Any + +import click + +from valohai_cli.api import request +from valohai_cli.ctx import get_project +from valohai_cli.messages import success, warn + +VALID_ENVVAR_NAME_RE = re.compile(r"^[a-z0-9_-]+$", re.I) + + +def parse_env_variable(s: str) -> tuple[str, str]: + """ + Parse an environment variable string of the form NAME=VALUE. + """ + if "=" not in s: + raise click.BadParameter("Environment variable must be in the form NAME=VALUE") + key, _, value = s.partition("=") + key = key.strip() + if not VALID_ENVVAR_NAME_RE.match(key): + raise click.BadParameter( + f"Environment variable name must be alphanumeric-and-dashes-and-underscores, got {key}", + ) + return key, value.strip() + + +@click.command() +@click.argument( + "environment_variables", + metavar="NAME=VALUE", + nargs=-1, + required=True, + type=parse_env_variable, +) +@click.option("--secret/--public", help="Set environment variable as a secret. Default value is --public") +def create( + *, + environment_variables: list[tuple[str, str]], + secret: bool, +) -> Any: + """ + Add an environment variable to a project + """ + project = get_project(require=True) + + if not environment_variables: # pragma: no cover + warn("Nothing to do.") + return + + for key, value in environment_variables: + payload = { + "name": key, + "value": value, + "secret": secret, + } + + request( + method="post", + url=f"/api/v0/projects/{project.id}/environment_variable/", + json=payload, + ) + + if len(environment_variables) == 1: + success(f"Added environment variable {environment_variables[0][0]} to project {project.name}.") + else: + success(f"Added {len(environment_variables)} environment variables to project {project.name}.") diff --git a/valohai_cli/commands/project/environment_variables/delete.py b/valohai_cli/commands/project/environment_variables/delete.py new file mode 100644 index 00000000..0d9ddc27 --- /dev/null +++ b/valohai_cli/commands/project/environment_variables/delete.py @@ -0,0 +1,35 @@ +from __future__ import annotations + +from typing import Any + +import click +from click.exceptions import Exit + +from valohai_cli.api import request +from valohai_cli.ctx import get_project +from valohai_cli.exceptions import APINotFoundError +from valohai_cli.messages import error, success + + +@click.command() +@click.argument("names", metavar="NAME", nargs=-1, required=True) +def delete(*, names: list[str]) -> Any: + """ + Delete one or more environment variables. + """ + + project = get_project(require=True) + fail = False + for name in names: + payload = {"name": name} + + try: + request(method="delete", url=f"/api/v0/projects/{project.id}/environment_variable/", json=payload) + except APINotFoundError: # pragma: no cover + error(f"Environment variable ({name}) not found in {project.name}") + fail = True + else: + success(f"Successfully deleted environment variable {name} from {project.name}") + + if fail: + raise Exit(1) diff --git a/valohai_cli/commands/project/environment_variables/list.py b/valohai_cli/commands/project/environment_variables/list.py new file mode 100644 index 00000000..7b78d2f4 --- /dev/null +++ b/valohai_cli/commands/project/environment_variables/list.py @@ -0,0 +1,43 @@ +from __future__ import annotations + +from typing import Any + +import click + +from valohai_cli.api import request +from valohai_cli.ctx import get_project +from valohai_cli.messages import info +from valohai_cli.settings import settings +from valohai_cli.table import print_json, print_table + + +@click.command() +def list() -> Any: + """ + List environment variables of a project + """ + + project = get_project(require=True) + data = request(method="get", url=f"/api/v0/projects/{project.id}/").json()["environment_variables"] + + if settings.output_format == "json": + return print_json(data) + + if not data: # pragma: no cover + info(f"{project}: No environment variables.") + return + + formatted_envvars = [ + { + "name": name, + "value": "****" if info["secret"] else info["value"], + "secret": info["secret"], + } + for name, info in sorted(data.items()) + ] + + print_table( + formatted_envvars, + columns=["name", "value", "secret"], + headers=["Name", "Value", "Secret"], + ) diff --git a/valohai_cli/plugin_cli.py b/valohai_cli/plugin_cli.py index 8384be7a..2f69ad22 100644 --- a/valohai_cli/plugin_cli.py +++ b/valohai_cli/plugin_cli.py @@ -16,6 +16,7 @@ class PluginCLI(click.MultiCommand): aliases = { + "add": "create", "new": "create", "start": "run", } @@ -138,7 +139,7 @@ def get_help(self, ctx: Context) -> str: # (see https://github.com/pallets/click/pull/1623). import json - return json.dumps(ctx.to_info_dict()) + return json.dumps(ctx.to_info_dict(), default=str) return super().get_help(ctx)