Add ability to pass in ENV to workstation

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "cloud-workstation"
-version = "0.1.3"
+version = "0.1.4"
 requires-python = ">=3.8"
 readme = "README.md"
 description = ""

src/workstation/cli/crud.py

@@ -180,6 +180,15 @@ def common_options(func):  # noqa: D103
     type=str,
     metavar="<str>",
 )
+@click.option(
+    "--envs",
+    "-e",
+    "envs",
+    type=(str, str),
+    multiple=True,
+    help="Environment variables to set at runtime.",
+    metavar="<key value>",
+)
 @click.pass_context
 def create(
     context: click.Context,
@@ -190,6 +199,7 @@ def create(
     project: Optional[str],
     proxy: Optional[str],
     no_proxy: Optional[str],
+    envs: Tuple[Tuple[str, str]],
     **kwargs,
 ):
     """Create a workstation."""
@@ -227,6 +237,7 @@ def create(
         location=location,
         proxy=proxy,
         no_proxy=no_proxy,
+        envs=envs,
     )
 
     config_manager.write_ssh_config(

src/workstation/core.py

@@ -1,5 +1,5 @@
 import sys
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Tuple
 
 from google.api_core.exceptions import AlreadyExists
 from google.api_core.operation import Operation
@@ -107,6 +107,7 @@ def create_workstation(
     user: str,
     proxy: Optional[str] = None,
     no_proxy: Optional[str] = None,
+    envs: Optional[Tuple[Tuple[str, str]]] = None,
 ) -> Workstation:
     """
     Create a new workstation with the specified configuration.
@@ -151,6 +152,18 @@ def create_workstation(
         env["no_proxy"] = no_proxy
         env["NO_PROXY"] = no_proxy
 
+    if envs:
+        user_envs = dict(envs)
+        # ensure that no duplicate keys are added to env
+        for key, value in user_envs.items():
+            if key not in env:
+                env[key] = value
+            else:
+                logger.warning(
+                    f"Environment variable {key} already exists in the environment, skipping"
+                )
+
+
     request = workstations_v1beta.CreateWorkstationRequest(
         parent=f"projects/{project}/locations/{location}/workstationClusters/{cluster}/workstationConfigs/{config}",
         workstation_id=name,

tests/test_ConfigManager.py

@@ -2,7 +2,6 @@
 from pathlib import Path
 
 import pytest
-
 from workstation.config import ConfigManager
 
 

tests/test_core.py

@@ -0,0 +1,35 @@
+import getpass
+from unittest.mock import patch
+
+from workstation.core import create_workstation
+
+
+@patch("workstation.core.workstations_v1beta.WorkstationsClient")
+@patch("workstation.core.config_manager.write_configuration")
+def test_create_workstation_env_dict(mock_write_configuration, mock_workstations_client):
+    mock_client_instance = mock_workstations_client.return_value
+    mock_operation = mock_client_instance.create_workstation.return_value
+    mock_operation.result.return_value = {}
+
+    envs = (("KEY1", "VALUE1"), ("KEY2", "VALUE2"))
+    create_workstation(
+        project="test-project",
+        location="us-central1",
+        cluster="default-cluster",
+        config="default-config",
+        name="test-workstation",
+        account="test-account",
+        user="test-user",
+        envs=envs,
+    )
+
+    expected_env = {
+        "LDAP": "test-user",
+        "ACCOUNT": "test-account",
+        "KEY1": "VALUE1",
+        "KEY2": "VALUE2",
+    }
+
+    assert mock_client_instance.create_workstation.call_count == 1
+    _, kwargs = mock_client_instance.create_workstation.call_args
+    assert kwargs["request"].workstation.env == expected_env

tests/test_crud.py

@@ -3,7 +3,6 @@
 
 import pytest
 from click.testing import CliRunner
-
 from workstation.cli import crud
 
 
@@ -119,3 +118,6 @@ def test_list(mock_get_gcloud_config, mock_check_gcloud_auth, mock_list_workstat
         "]\n"
     )
     assert result.output == expected_json_output
+
+
+

tests/test_utils.py

@@ -1,6 +1,5 @@
 import pytest
 from pytest_mock import MockerFixture
-
 from workstation.utils import get_instance_assignment, process_entry