fix dataset tests

beaker/client.py

@@ -308,15 +308,6 @@ def workspace(self) -> WorkspaceClient:
         """
         Manage workspaces.
 
-        :examples:
-
-        >>> beaker.workspace.datasets(
-        ...     match="squad",
-        ...     uncommitted=False,
-        ...     results=False,
-        ... )[0].full_name
-        'petew/squad-train'
-
         .. tip::
             See the `Workspaces Overview <overview.html#workspaces>`_ for a walk-through of the
             main methods, or check out the `Workspace API Docs <#workspace>`_
@@ -363,11 +354,6 @@ def dataset(self) -> DatasetClient:
         """
         Manage datasets.
 
-        :examples:
-
-        >>> [file_info.path for file_info in beaker.dataset.ls("petew/squad-train")]
-        ['squad-train.arrow']
-
         .. tip::
             See the `Datasets Overview <overview.html#datasets>`_ for a walk-through of the
             main methods, or check out the `Dataset API Docs <#dataset>`_

beaker/conftest.py

@@ -19,6 +19,7 @@ def doctest_fixtures(
     group_name,
     hello_world_experiment_name,
     squad_dataset_name,
+    squad_dataset_file_name,
     tmp_path,
 ):
     doctest_namespace["beaker"] = client
@@ -36,4 +37,5 @@ def doctest_fixtures(
     doctest_namespace["group_name"] = group_name
     doctest_namespace["hello_world_experiment_name"] = hello_world_experiment_name
     doctest_namespace["squad_dataset_name"] = squad_dataset_name
+    doctest_namespace["squad_dataset_file_name"] = squad_dataset_file_name
     doctest_namespace["tmp_path"] = tmp_path

beaker/services/dataset.py

@@ -341,8 +341,8 @@ def stream_file(
         :examples:
 
         >>> total_bytes = 0
-        >>> with open(tmp_path / "squad-train.arrow", "wb") as f:
-        ...     for chunk in beaker.dataset.stream_file(squad_dataset_name, "squad-train.arrow", quiet=True):
+        >>> with open(tmp_path / squad_dataset_file_name, "wb") as f:
+        ...     for chunk in beaker.dataset.stream_file(squad_dataset_name, squad_dataset_file_name, quiet=True):
         ...         total_bytes += f.write(chunk)
         """
         dataset = self.resolve_dataset(dataset, ensure_storage=True)
@@ -400,7 +400,7 @@ def get_file(
 
         :examples:
 
-        >>> contents = beaker.dataset.get_file(squad_dataset_name, "squad-train.arrow", quiet=True)
+        >>> contents = beaker.dataset.get_file(squad_dataset_name, squad_dataset_file_name, quiet=True)
         """
 
         @retriable(recoverable_errors=(RequestException, ChecksumFailedError))

conftest.py

@@ -17,8 +17,8 @@ def unique_name() -> str:
     return petname.generate() + "-" + str(uuid.uuid4())[:8]  # type: ignore
 
 
-def beaker_object_fixture(client: Beaker, service: str):
-    name = unique_name()
+def beaker_object_fixture(client: Beaker, service: str, prefix: str = ""):
+    name = prefix + unique_name()
     service_client = getattr(client, service)
     not_found_exception = getattr(exceptions, f"{service.title()}NotFound")
     yield name
@@ -168,13 +168,17 @@ def archived_workspace(client: Beaker, archived_workspace_name: str) -> Workspac
 
 
 @pytest.fixture()
-def squad_dataset_name() -> str:
-    return "petew/squad-train"
+def squad_dataset_file_name() -> str:
+    return "squad-train.arrow"
 
 
 @pytest.fixture()
-def squad_dataset_file_name() -> str:
-    return "squad-train.arrow"
+def squad_dataset_name(client: Beaker, squad_dataset_file_name) -> Generator[str, None, None]:
+    for dataset_name in beaker_object_fixture(client, "dataset", prefix="squad"):
+        dataset = client.dataset.create(dataset_name, commit=False)
+        client.dataset.upload(dataset, b"blahblahblah", squad_dataset_file_name)
+        client.dataset.commit(dataset)
+        yield dataset_name
 
 
 @pytest.fixture()

tests/dataset_test.py

@@ -1,23 +1,9 @@
 from beaker.client import Beaker
 
 
-def test_dataset_get(client: Beaker, squad_dataset_name: str):
-    dataset = client.dataset.get(squad_dataset_name)
-    assert dataset.name is not None
-    # Try with ID.
-    client.dataset.get(dataset.id)
-    # Try with just name (without account prefix).
-    client.dataset.get(dataset.name)
-
-
-def test_dataset_ls(client: Beaker, squad_dataset_name: str):
-    client.dataset.ls(squad_dataset_name)
-
-
-def test_file_info(client: Beaker, squad_dataset_name: str, squad_dataset_file_name: str):
-    client.dataset.file_info(squad_dataset_name, squad_dataset_file_name)
-
-
-def test_upload(client: Beaker, dataset_name: str):
+def test_create_upload_commit(client: Beaker, dataset_name: str):
     ds = client.dataset.create(dataset_name, commit=False)
     client.dataset.upload(ds, b"foo-bar", "foo-bar")
+    client.dataset.commit(ds)
+    client.dataset.ls(ds)
+    client.dataset.file_info(ds, "foo-bar")