RSDK-10208 - create app client from env vars (#874)

Author: stuqdog

src/viam/app/viam_client.py

@@ -1,3 +1,4 @@
+import os
 from typing import Mapping, Optional
 
 from grpclib.client import Channel
@@ -10,7 +11,7 @@
 from viam.app.ml_training_client import MLTrainingClient
 from viam.app.provisioning_client import ProvisioningClient
 from viam.robot.client import RobotClient
-from viam.rpc.dial import DialOptions, _dial_app, _get_access_token
+from viam.rpc.dial import Credentials, DialOptions, _dial_app, _get_access_token
 
 LOGGER = logging.getLogger(__name__)
 
@@ -24,14 +25,46 @@ class ViamClient:
         ViamClient.create_from_dial_options(...)
     """
 
+    @classmethod
+    async def create_from_env_vars(cls, dial_options: Optional[DialOptions] = None, app_url: Optional[str] = None) -> Self:
+        """Create `ViamClient` using credentials set in the environment as `VIAM_API_KEY` and `VIAM_API_KEY_ID`.
+
+        ::
+
+            client = await ViamClient.create_from_env_vars()
+
+        Args:
+            dial_options (Optional[viam.rpc.dial.DialOptions]): Options for authorization and connection to app.
+                If not provided, default options will be selected. Note that `creds` and `auth_entity`
+                fields will be overwritten by the values set by a module.
+            app_url: (Optional[str]): URL of app. Uses app.viam.com if not specified.
+
+        Raises:
+            ValueError: If there are no env vars set by the module, or if they are set improperly
+
+        """
+        dial_options = dial_options if dial_options else DialOptions()
+        api_key = os.environ.get("VIAM_API_KEY")
+        if api_key is None:
+            raise ValueError("api key cannot be None")
+        api_key_id = os.environ.get("VIAM_API_KEY_ID")
+        if api_key_id is None:
+            raise ValueError("api key ID cannot be None")
+        credentials = Credentials(type="api-key", payload=api_key)
+        dial_options.credentials = credentials
+        dial_options.auth_entity = api_key_id
+
+        return await cls.create_from_dial_options(dial_options, app_url)
+
+
     @classmethod
     async def create_from_dial_options(cls, dial_options: DialOptions, app_url: Optional[str] = None) -> Self:
         """Create `ViamClient` that establishes a connection to the Viam app.
 
         ::
 
             dial_options = DialOptions.with_api_key("<API-KEY>", "<API-KEY-ID>")
-            ViamClient.create_from_dial_options(dial_options)
+            client = await ViamClient.create_from_dial_options(dial_options)
 
         Args:
             dial_options (viam.rpc.dial.DialOptions): Required information for authorization and connection to app.

tests/test_viam_client.py

@@ -1,4 +1,5 @@
 import random
+import os
 from unittest.mock import patch
 from uuid import uuid4
 
@@ -73,6 +74,35 @@ async def test_clients(self):
                     assert client.provisioning_client._channel == channel
                     assert client.provisioning_client._metadata == METADATA
 
+    async def test_client_from_env_vars(self):
+        async with ChannelFor([]) as channel:
+            with patch("viam.app.viam_client._dial_app") as patched_dial:
+                patched_dial.return_value = channel
+                with patch("viam.app.viam_client._get_access_token") as patched_auth:
+                    ACCESS_TOKEN = "MY_ACCESS_TOKEN"
+                    METADATA = {"authorization": f"Bearer {ACCESS_TOKEN}"}
+                    patched_auth.return_value = ACCESS_TOKEN
+
+                    os.environ["VIAM_API_KEY"] = "MY_API_KEY"
+                    os.environ["VIAM_API_KEY_ID"] = str(uuid4())
+
+                    client = await ViamClient.create_from_env_vars()
+
+                    assert client.data_client._channel == channel
+                    assert client.data_client._metadata == METADATA
+
+                    assert client.app_client._channel == channel
+                    assert client.app_client._metadata == METADATA
+
+                    assert client.ml_training_client._channel == channel
+                    assert client.ml_training_client._metadata == METADATA
+
+                    assert client.billing_client._channel == channel
+                    assert client.billing_client._metadata == METADATA
+
+                    assert client.provisioning_client._channel == channel
+                    assert client.provisioning_client._metadata == METADATA
+
     async def test_closes(self):
         async with ChannelFor([]) as channel:
             with patch.object(channel, "close"):