From a35e8aa14be0743227f447babe21642ee319f280 Mon Sep 17 00:00:00 2001 From: epwalsh Date: Wed, 20 Mar 2024 16:23:12 -0700 Subject: [PATCH] fix dataset tests --- beaker/client.py | 14 -------------- beaker/conftest.py | 2 ++ beaker/services/dataset.py | 6 +++--- conftest.py | 16 ++++++++++------ tests/dataset_test.py | 22 ++++------------------ 5 files changed, 19 insertions(+), 41 deletions(-) diff --git a/beaker/client.py b/beaker/client.py index f0dffeb..2d96ff7 100644 --- a/beaker/client.py +++ b/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 `_ 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 `_ for a walk-through of the main methods, or check out the `Dataset API Docs <#dataset>`_ diff --git a/beaker/conftest.py b/beaker/conftest.py index 0bf1aad..5948e9f 100644 --- a/beaker/conftest.py +++ b/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 diff --git a/beaker/services/dataset.py b/beaker/services/dataset.py index 12534fd..0c2a824 100644 --- a/beaker/services/dataset.py +++ b/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)) diff --git a/conftest.py b/conftest.py index 57110b4..acfc5f4 100644 --- a/conftest.py +++ b/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() diff --git a/tests/dataset_test.py b/tests/dataset_test.py index 2d4dde3..e470b7f 100644 --- a/tests/dataset_test.py +++ b/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")