clean up

beaker/client.py

@@ -38,12 +38,11 @@ def user(self) -> str:
         return self.whoami()["name"]
 
     @classmethod
-    def from_env(cls, **kwargs) -> "Beaker":
+    def from_env(cls, **overrides) -> "Beaker":
         """
-        Initialize client from environment variables. Expects the beaker auth token
-        to be set as the ``BEAKER_TOKEN`` environment variable.
+        Initialize client from a config file and/or environment variables.
         """
-        return cls(Config.from_env())
+        return cls(Config.from_env(**overrides))
 
     @contextmanager
     def _session_with_backoff(self) -> requests.Session:
@@ -108,6 +107,13 @@ def get_workspace(self, workspace: Optional[str] = None) -> Dict[str, Any]:
     def ensure_workspace(self, workspace: str):
         """
         Ensure the given workspace exists.
+
+        Raises
+        ------
+        HTTPError
+        ValueError
+            If the workspace name is invalid.
+
         """
         try:
             self.get_workspace(workspace)

beaker/config.py

@@ -37,7 +37,7 @@ class Config:
     TOKEN_KEY: ClassVar[str] = "BEAKER_TOKEN"
 
     @classmethod
-    def from_env(cls) -> "Config":
+    def from_env(cls, **overrides) -> "Config":
         """
         Initialize a config from environment variables or a local config file if one
         can be found.
@@ -66,6 +66,13 @@ def from_env(cls) -> "Config":
         if cls.ADDRESS_KEY in os.environ:
             config.agent_address = os.environ[cls.ADDRESS_KEY]
 
+        # Override with any arguments passed to this method.
+        for name, value in overrides.items():
+            if hasattr(config, name):
+                setattr(config, name, value)
+            else:
+                raise AttributeError(f"'{cls.__name__}' has to attribute '{name}'")
+
         return config
 
     @classmethod

tests/client_test.py

@@ -5,14 +5,11 @@
 import pytest
 
 from beaker.client import Beaker
-from beaker.config import Config
 
 
 @pytest.fixture(scope="module")
 def client():
-    config = Config.from_env()
-    config.default_workspace = "ai2/beaker-py"
-    return Beaker(config)
+    return Beaker.from_env(default_workspace="ai2/beaker-py")
 
 
 @pytest.fixture(scope="module")