Merge pull request #1540 from kabilar/embargoed-zarr

Remove `NotImplementedError` to allow for uploading Zarr assets to embargoed Dandisets

dandi/files/zarr.py

@@ -22,7 +22,6 @@
     ZARR_DELETE_BATCH_SIZE,
     ZARR_MIME_TYPE,
     ZARR_UPLOAD_BATCH_SIZE,
-    EmbargoStatus,
 )
 from dandi.dandiapi import (
     RemoteAsset,
@@ -297,12 +296,6 @@ def iter_upload(
             ``"done"`` and an ``"asset"`` key containing the resulting
             `RemoteAsset`.
         """
-        # So that older clients don't get away with doing the wrong thing once
-        # Zarr upload to embargoed Dandisets is implemented in the API:
-        if dandiset.embargo_status is EmbargoStatus.EMBARGOED:
-            raise NotImplementedError(
-                "Uploading Zarr assets to embargoed Dandisets is currently not implemented"
-            )
         asset_path = metadata.setdefault("path", self.path)
         client = dandiset.client
         lgr.debug("%s: Producing asset", asset_path)

dandi/tests/fixtures.py

@@ -610,12 +610,23 @@ def sample_dandiset_factory(
     return SampleDandisetFactory(local_dandi_api, tmp_path_factory)
 
 
+# @sweep_embargo can be used along with `new_dandiset` fixture to
+# run the test against both embargoed and non-embargoed Dandisets.
+# "embargo: bool" argument should be added to the test function.
+sweep_embargo = pytest.mark.parametrize("embargo", [True, False])
+
+
 @pytest.fixture()
 def new_dandiset(
     request: pytest.FixtureRequest, sample_dandiset_factory: SampleDandisetFactory
 ) -> SampleDandiset:
+    kws = {}
+    if "embargo" in request.node.fixturenames:
+        kws["embargo"] = request.node.callspec.params["embargo"]
     return sample_dandiset_factory.mkdandiset(
-        f"Sample Dandiset for {request.node.name}"
+        f"Sample {'Embargoed' if kws.get('embargo') else 'Public'} "
+        f"Dandiset for {request.node.name}",
+        **kws,
     )
 
 

dandi/tests/test_upload.py

@@ -11,9 +11,9 @@
 from pytest_mock import MockerFixture
 import zarr
 
-from .fixtures import SampleDandiset
+from .fixtures import SampleDandiset, sweep_embargo
 from .test_helpers import assert_dirtrees_eq
-from ..consts import ZARR_MIME_TYPE, dandiset_metadata_file
+from ..consts import ZARR_MIME_TYPE, EmbargoStatus, dandiset_metadata_file
 from ..dandiapi import AssetType, RemoteBlobAsset, RemoteZarrAsset
 from ..dandiset import Dandiset
 from ..download import download
@@ -151,9 +151,10 @@ def test_upload_download_small_file(
     assert (tmp_path / dandiset_id / "file.txt").read_bytes() == contents
 
 
+@sweep_embargo
 @pytest.mark.parametrize("confirm", [True, False])
 def test_upload_sync(
-    confirm: bool, mocker: MockerFixture, text_dandiset: SampleDandiset
+    confirm: bool, mocker: MockerFixture, text_dandiset: SampleDandiset, embargo: bool
 ) -> None:
     (text_dandiset.dspath / "file.txt").unlink()
     confirm_mock = mocker.patch("click.confirm", return_value=confirm)
@@ -250,7 +251,13 @@ def test_upload_bids_non_nwb_file(bids_dandiset: SampleDandiset) -> None:
     assert [asset.path for asset in bids_dandiset.dandiset.get_assets()] == ["README"]
 
 
-def test_upload_sync_zarr(mocker, zarr_dandiset):
+@sweep_embargo
+def test_upload_sync_zarr(
+    mocker: MockerFixture, zarr_dandiset: SampleDandiset, embargo: bool
+) -> None:
+    assert zarr_dandiset.dandiset.embargo_status == (
+        EmbargoStatus.EMBARGOED if embargo else EmbargoStatus.OPEN
+    )
     rmtree(zarr_dandiset.dspath / "sample.zarr")
     zarr.save(zarr_dandiset.dspath / "identity.zarr", np.eye(5))
     confirm_mock = mocker.patch("click.confirm", return_value=True)