Add weka option to DataSource (#280)

* Add `weka` option to `DataSource`

* fix

* fix tests

CHANGELOG.md

@@ -9,6 +9,10 @@ use patch releases for compatibility fixes instead.
 
 ## Unreleased
 
+### Added
+
+- Added `weka` option to `DataSource` model.
+
 ### Fixed
 
 - Only check for upgrades once every 12 hours by default.

beaker/data_model/experiment_spec.py

@@ -101,6 +101,11 @@ class DataSource(BaseModel, frozen=False):
     - ``/var/beaker/share`` as a shared local scratch space.
     """
 
+    weka: Optional[str] = None
+    """
+    The name of a weka bucket.
+    """
+
     result: Optional[str] = None
     """
     Name of a previous task whose result will be mounted.
@@ -162,6 +167,7 @@ def new(
         sub_path: Optional[str] = None,
         beaker: Optional[str] = None,
         host_path: Optional[str] = None,
+        weka: Optional[str] = None,
         result: Optional[str] = None,
         secret: Optional[str] = None,
     ) -> "DataMount":
@@ -172,6 +178,7 @@ def new(
         :param sub_path: The :data:`sub_path`.
         :param beaker: The :data:`beaker <DataSource.beaker>` argument to :class:`DataSource`.
         :param host_path: The :data:`host_path <DataSource.host_path>` argument to :class:`DataSource`.
+        :param weka: The :data:`weka <DataSource.weka>` argument to :class:`DataSource`.
         :param result: The :data:`result <DataSource.result>` argument to :class:`DataSource`.
         :param url: The :data:`url <DataSource.url>` argument to :class:`DataSource`.
         :param secret: The :data:`secret <DataSource.secret>` argument to :class:`DataSource`.
@@ -182,6 +189,7 @@ def new(
             source=DataSource(
                 beaker=beaker,
                 host_path=host_path,
+                weka=weka,
                 result=result,
                 secret=secret,
             ),

tests/client_test.py

@@ -2,13 +2,16 @@
 from flaky import flaky
 
 from beaker import Beaker
+from beaker.config import InternalConfig
 
 
 @flaky  # this can fail if the request to GitHub fails
 def test_warn_for_newer_version(monkeypatch):
     import beaker.client
     import beaker.version
 
+    InternalConfig().save()
+
     monkeypatch.setattr(Beaker, "CLIENT_VERSION", "0.1.0")
     monkeypatch.setattr(beaker.client, "_LATEST_VERSION_CHECKED", False)
 

tests/data_model_test.py

@@ -93,10 +93,16 @@ def test_experiment_spec_validation():
 
 def test_snake_case_vs_lower_camel_case():
     for x in (DataSource(host_path="/tmp/foo"), DataSource(hostPath="/tmp/foo")):  # type: ignore
-        assert str(x) == "DataSource(beaker=None, host_path='/tmp/foo', result=None, secret=None)"
+        assert (
+            str(x)
+            == "DataSource(beaker=None, host_path='/tmp/foo', weka=None, result=None, secret=None)"
+        )
         assert x.host_path == "/tmp/foo"
         x.host_path = "/tmp/bar"
-        assert str(x) == "DataSource(beaker=None, host_path='/tmp/bar', result=None, secret=None)"
+        assert (
+            str(x)
+            == "DataSource(beaker=None, host_path='/tmp/bar', weka=None, result=None, secret=None)"
+        )
         assert x.to_json() == {"hostPath": "/tmp/bar"}