arkindex-base-worker 0.3.7rc6__py3-none-any.whl → 0.3.7rc7__py3-none-any.whl

tests/test_dataset_worker.py

@@ -1,13 +1,59 @@
 import logging
+import uuid
+from argparse import ArgumentTypeError
 
 import pytest
 from apistar.exceptions import ErrorResponse
 
-from arkindex_worker.worker import MissingDatasetArchive
+from arkindex_worker.models import Dataset, Set
+from arkindex_worker.worker import MissingDatasetArchive, check_dataset_set
 from arkindex_worker.worker.dataset import DatasetState
 from tests.conftest import FIXTURES_DIR, PROCESS_ID
 from tests.test_elements_worker import BASE_API_CALLS
 
+RANDOM_UUID = uuid.uuid4()
+
+
+@pytest.fixture()
+def tmp_archive(tmp_path):
+    archive = tmp_path / "test_archive.tar.zst"
+    archive.touch()
+
+    yield archive
+
+    archive.unlink(missing_ok=True)
+
+
+@pytest.mark.parametrize(
+    ("value", "error"),
+    [("train", ""), (f"{RANDOM_UUID}:train:val", ""), ("not_uuid:train", "")],
+)
+def test_check_dataset_set_errors(value, error):
+    with pytest.raises(ArgumentTypeError, match=error):
+        check_dataset_set(value)
+
+
+def test_check_dataset_set():
+    assert check_dataset_set(f"{RANDOM_UUID}:train") == (RANDOM_UUID, "train")
+
+
+def test_cleanup_downloaded_artifact_no_download(mock_dataset_worker):
+    assert not mock_dataset_worker.downloaded_artifact
+    # Do nothing
+    mock_dataset_worker.cleanup_downloaded_artifact()
+
+
+def test_cleanup_downloaded_artifact(mock_dataset_worker, tmp_archive):
+    mock_dataset_worker.downloaded_artifact = tmp_archive
+
+    assert mock_dataset_worker.downloaded_artifact.exists()
+    # Unlink the downloaded archive
+    mock_dataset_worker.cleanup_downloaded_artifact()
+    assert not mock_dataset_worker.downloaded_artifact.exists()
+
+    # Unlinking again does not raise an error even if the archive no longer exists
+    mock_dataset_worker.cleanup_downloaded_artifact()
+
 
 def test_download_dataset_artifact_list_api_error(
     responses, mock_dataset_worker, default_dataset
@@ -127,8 +173,15 @@ def test_download_dataset_artifact_no_archive(
     ]
 
 
+@pytest.mark.parametrize("downloaded_cache", [False, True])
 def test_download_dataset_artifact(
-    mocker, tmp_path, responses, mock_dataset_worker, default_dataset
+    mocker,
+    tmp_path,
+    responses,
+    mock_dataset_worker,
+    default_dataset,
+    downloaded_cache,
+    tmp_archive,
 ):
     task_id = default_dataset.task_id
     archive_path = (
@@ -176,10 +229,22 @@ def test_download_dataset_artifact(
         content_type="application/zstd",
     )
 
-    archive = mock_dataset_worker.download_dataset_artifact(default_dataset)
-    assert archive == tmp_path / "dataset_id.tar.zst"
-    assert archive.read_bytes() == archive_path.read_bytes()
-    archive.unlink()
+    if downloaded_cache:
+        mock_dataset_worker.downloaded_artifact = tmp_archive
+    previous_artifact = mock_dataset_worker.downloaded_artifact
+
+    mock_dataset_worker.download_dataset_artifact(default_dataset)
+
+    # We removed the artifact that was downloaded previously
+    if previous_artifact:
+        assert not previous_artifact.exists()
+
+    assert mock_dataset_worker.downloaded_artifact == tmp_path / "dataset_id.tar.zst"
+    assert (
+        mock_dataset_worker.downloaded_artifact.read_bytes()
+        == archive_path.read_bytes()
+    )
+    mock_dataset_worker.downloaded_artifact.unlink()
 
     assert len(responses.calls) == len(BASE_API_CALLS) + 2
     assert [
@@ -190,189 +255,107 @@ def test_download_dataset_artifact(
     ]
 
 
-def test_list_dataset_elements_per_split_api_error(
-    responses, mock_dataset_worker, default_dataset
+def test_download_dataset_artifact_already_exists(
+    mocker, tmp_path, responses, mock_dataset_worker, default_dataset
 ):
-    responses.add(
-        responses.GET,
-        f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?set=set_1&with_count=true",
-        status=500,
+    mocker.patch(
+        "arkindex_worker.worker.base.BaseWorker.find_extras_directory",
+        return_value=tmp_path,
     )
+    already_downloaded = tmp_path / "dataset_id.tar.zst"
+    already_downloaded.write_bytes(b"Some content")
+    mock_dataset_worker.downloaded_artifact = already_downloaded
 
-    with pytest.raises(
-        Exception, match="Stopping pagination as data will be incomplete"
-    ):
-        mock_dataset_worker.list_dataset_elements_per_split(default_dataset)
+    mock_dataset_worker.download_dataset_artifact(default_dataset)
 
-    assert len(responses.calls) == len(BASE_API_CALLS) + 5
-    assert [
-        (call.request.method, call.request.url) for call in responses.calls
-    ] == BASE_API_CALLS + [
-        # The API call is retried 5 times
-        (
-            "GET",
-            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?set=set_1&with_count=true",
-        ),
-        (
-            "GET",
-            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?set=set_1&with_count=true",
-        ),
-        (
-            "GET",
-            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?set=set_1&with_count=true",
-        ),
-        (
-            "GET",
-            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?set=set_1&with_count=true",
-        ),
-        (
-            "GET",
-            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?set=set_1&with_count=true",
-        ),
-    ]
+    assert mock_dataset_worker.downloaded_artifact == already_downloaded
+    already_downloaded.unlink()
 
-
-def test_list_dataset_elements_per_split(
-    responses, mock_dataset_worker, default_dataset
-):
-    expected_results = []
-    for selected_set in default_dataset.selected_sets:
-        index = selected_set[-1]
-        expected_results.append(
-            {
-                "set": selected_set,
-                "element": {
-                    "id": str(index) * 4,
-                    "type": "page",
-                    "name": f"Test {index}",
-                    "corpus": {},
-                    "thumbnail_url": None,
-                    "zone": {},
-                    "best_classes": None,
-                    "has_children": None,
-                    "worker_version_id": None,
-                    "worker_run_id": None,
-                },
-            }
-        )
-        responses.add(
-            responses.GET,
-            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?set={selected_set}&with_count=true",
-            status=200,
-            json={
-                "count": 1,
-                "next": None,
-                "results": [expected_results[-1]],
-            },
-        )
-
-    assert list(
-        mock_dataset_worker.list_dataset_elements_per_split(default_dataset)
-    ) == [
-        ("set_1", [expected_results[0]["element"]]),
-        ("set_2", [expected_results[1]["element"]]),
-        ("set_3", [expected_results[2]["element"]]),
-    ]
-
-    assert len(responses.calls) == len(BASE_API_CALLS) + 3
+    assert len(responses.calls) == len(BASE_API_CALLS)
     assert [
         (call.request.method, call.request.url) for call in responses.calls
-    ] == BASE_API_CALLS + [
-        (
-            "GET",
-            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?set=set_1&with_count=true",
-        ),
-        (
-            "GET",
-            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?set=set_2&with_count=true",
-        ),
-        (
-            "GET",
-            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?set=set_3&with_count=true",
-        ),
-    ]
+    ] == BASE_API_CALLS
 
 
-def test_list_datasets_read_only(mock_dev_dataset_worker):
-    assert list(mock_dev_dataset_worker.list_datasets()) == [
-        "11111111-1111-1111-1111-111111111111",
-        "22222222-2222-2222-2222-222222222222",
-    ]
-
-
-def test_list_datasets_api_error(responses, mock_dataset_worker):
+def test_list_sets_api_error(responses, mock_dataset_worker):
     responses.add(
         responses.GET,
-        f"http://testserver/api/v1/process/{PROCESS_ID}/datasets/",
+        f"http://testserver/api/v1/process/{PROCESS_ID}/sets/",
         status=500,
     )
 
     with pytest.raises(
         Exception, match="Stopping pagination as data will be incomplete"
     ):
-        next(mock_dataset_worker.list_datasets())
+        next(mock_dataset_worker.list_sets())
 
     assert len(responses.calls) == len(BASE_API_CALLS) + 5
     assert [
         (call.request.method, call.request.url) for call in responses.calls
     ] == BASE_API_CALLS + [
         # The API call is retried 5 times
-        ("GET", f"http://testserver/api/v1/process/{PROCESS_ID}/datasets/"),
-        ("GET", f"http://testserver/api/v1/process/{PROCESS_ID}/datasets/"),
-        ("GET", f"http://testserver/api/v1/process/{PROCESS_ID}/datasets/"),
-        ("GET", f"http://testserver/api/v1/process/{PROCESS_ID}/datasets/"),
-        ("GET", f"http://testserver/api/v1/process/{PROCESS_ID}/datasets/"),
+        ("GET", f"http://testserver/api/v1/process/{PROCESS_ID}/sets/"),
+        ("GET", f"http://testserver/api/v1/process/{PROCESS_ID}/sets/"),
+        ("GET", f"http://testserver/api/v1/process/{PROCESS_ID}/sets/"),
+        ("GET", f"http://testserver/api/v1/process/{PROCESS_ID}/sets/"),
+        ("GET", f"http://testserver/api/v1/process/{PROCESS_ID}/sets/"),
     ]
 
 
-def test_list_datasets(responses, mock_dataset_worker):
+def test_list_sets(responses, mock_dataset_worker):
     expected_results = [
         {
-            "id": "process_dataset_1",
+            "id": "set_1",
             "dataset": {
                 "id": "dataset_1",
                 "name": "Dataset 1",
                 "description": "My first great dataset",
-                "sets": ["train", "val", "test"],
+                "sets": [
+                    {"id": "set_1", "name": "train"},
+                    {"id": "set_2", "name": "val"},
+                ],
                 "state": "open",
                 "corpus_id": "corpus_id",
                 "creator": "test@teklia.com",
                 "task_id": "task_id_1",
             },
-            "sets": ["test"],
+            "set_name": "train",
         },
         {
-            "id": "process_dataset_2",
+            "id": "set_2",
             "dataset": {
-                "id": "dataset_2",
-                "name": "Dataset 2",
-                "description": "My second great dataset",
-                "sets": ["train", "val"],
-                "state": "complete",
+                "id": "dataset_1",
+                "name": "Dataset 1",
+                "description": "My first great dataset",
+                "sets": [
+                    {"id": "set_1", "name": "train"},
+                    {"id": "set_2", "name": "val"},
+                ],
+                "state": "open",
                 "corpus_id": "corpus_id",
                 "creator": "test@teklia.com",
-                "task_id": "task_id_2",
+                "task_id": "task_id_1",
             },
-            "sets": ["train", "val"],
+            "set_name": "val",
         },
         {
-            "id": "process_dataset_3",
+            "id": "set_3",
             "dataset": {
-                "id": "dataset_3",
-                "name": "Dataset 3 (TRASHME)",
-                "description": "My third dataset, in error",
-                "sets": ["nonsense", "random set"],
-                "state": "error",
+                "id": "dataset_2",
+                "name": "Dataset 2",
+                "description": "My second great dataset",
+                "sets": [{"id": "set_3", "name": "my_set"}],
+                "state": "complete",
                 "corpus_id": "corpus_id",
                 "creator": "test@teklia.com",
-                "task_id": "task_id_3",
+                "task_id": "task_id_2",
             },
-            "sets": ["random set"],
+            "set_name": "my_set",
         },
     ]
     responses.add(
         responses.GET,
-        f"http://testserver/api/v1/process/{PROCESS_ID}/datasets/",
+        f"http://testserver/api/v1/process/{PROCESS_ID}/sets/",
         status=200,
         json={
             "count": 3,
@@ -381,154 +364,109 @@ def test_list_datasets(responses, mock_dataset_worker):
         },
     )
 
-    for idx, dataset in enumerate(mock_dataset_worker.list_process_datasets()):
-        assert dataset == {
-            **expected_results[idx]["dataset"],
-            "selected_sets": expected_results[idx]["sets"],
-        }
+    for idx, dataset_set in enumerate(mock_dataset_worker.list_process_sets()):
+        assert isinstance(dataset_set, Set)
+        assert dataset_set.name == expected_results[idx]["set_name"]
+
+        assert isinstance(dataset_set.dataset, Dataset)
+        assert dataset_set.dataset == expected_results[idx]["dataset"]
 
     assert len(responses.calls) == len(BASE_API_CALLS) + 1
     assert [
         (call.request.method, call.request.url) for call in responses.calls
     ] == BASE_API_CALLS + [
-        ("GET", f"http://testserver/api/v1/process/{PROCESS_ID}/datasets/"),
+        ("GET", f"http://testserver/api/v1/process/{PROCESS_ID}/sets/"),
     ]
 
 
-@pytest.mark.parametrize("generator", [True, False])
-def test_run_no_datasets(mocker, caplog, mock_dataset_worker, generator):
-    mocker.patch("arkindex_worker.worker.DatasetWorker.list_datasets", return_value=[])
-    mock_dataset_worker.generator = generator
+def test_list_sets_retrieve_dataset_api_error(
+    responses, mock_dev_dataset_worker, default_dataset
+):
+    mock_dev_dataset_worker.args.set = [
+        (default_dataset.id, "train"),
+        (default_dataset.id, "val"),
+    ]
 
-    with pytest.raises(SystemExit):
-        mock_dataset_worker.run()
+    responses.add(
+        responses.GET,
+        f"http://testserver/api/v1/datasets/{default_dataset.id}/",
+        status=500,
+    )
 
-    assert [(level, message) for _, level, message in caplog.record_tuples] == [
-        (logging.INFO, "Loaded Worker Fake worker @ 123412 from API"),
-        (logging.WARNING, "No datasets to process, stopping."),
+    with pytest.raises(ErrorResponse):
+        next(mock_dev_dataset_worker.list_sets())
+
+    assert len(responses.calls) == 5
+    assert [(call.request.method, call.request.url) for call in responses.calls] == [
+        # The API call is retried 5 times
+        ("GET", f"http://testserver/api/v1/datasets/{default_dataset.id}/"),
+        ("GET", f"http://testserver/api/v1/datasets/{default_dataset.id}/"),
+        ("GET", f"http://testserver/api/v1/datasets/{default_dataset.id}/"),
+        ("GET", f"http://testserver/api/v1/datasets/{default_dataset.id}/"),
+        ("GET", f"http://testserver/api/v1/datasets/{default_dataset.id}/"),
     ]
 
 
-@pytest.mark.parametrize(
-    ("generator", "error"),
-    [
-        (True, "When generating a new dataset, its state should be Open or Error."),
-        (False, "When processing an existing dataset, its state should be Complete."),
-    ],
-)
-def test_run_initial_dataset_state_error(
-    mocker, responses, caplog, mock_dataset_worker, default_dataset, generator, error
-):
-    default_dataset.state = DatasetState.Building.value
-    mocker.patch(
-        "arkindex_worker.worker.DatasetWorker.list_datasets",
-        return_value=[default_dataset],
+def test_list_sets_read_only(responses, mock_dev_dataset_worker, default_dataset):
+    mock_dev_dataset_worker.args.set = [
+        (default_dataset.id, "train"),
+        (default_dataset.id, "val"),
+    ]
+
+    responses.add(
+        responses.GET,
+        f"http://testserver/api/v1/datasets/{default_dataset.id}/",
+        status=200,
+        json=default_dataset,
     )
-    mock_dataset_worker.generator = generator
-
-    extra_call = []
-    if generator:
-        responses.add(
-            responses.PATCH,
-            f"http://testserver/api/v1/datasets/{default_dataset.id}/",
-            status=200,
-            json={},
-        )
-        extra_call = [
-            ("PATCH", f"http://testserver/api/v1/datasets/{default_dataset.id}/"),
-        ]
+
+    assert list(mock_dev_dataset_worker.list_sets()) == [
+        Set(name="train", dataset=default_dataset),
+        Set(name="val", dataset=default_dataset),
+    ]
+
+    assert len(responses.calls) == 1
+    assert [(call.request.method, call.request.url) for call in responses.calls] == [
+        ("GET", f"http://testserver/api/v1/datasets/{default_dataset.id}/"),
+    ]
+
+
+def test_run_no_sets(mocker, caplog, mock_dataset_worker):
+    mocker.patch("arkindex_worker.worker.DatasetWorker.list_sets", return_value=[])
 
     with pytest.raises(SystemExit):
         mock_dataset_worker.run()
 
-    assert len(responses.calls) == len(BASE_API_CALLS) * 2 + len(extra_call)
-    assert [
-        (call.request.method, call.request.url) for call in responses.calls
-    ] == BASE_API_CALLS * 2 + extra_call
-
     assert [(level, message) for _, level, message in caplog.record_tuples] == [
         (logging.INFO, "Loaded Worker Fake worker @ 123412 from API"),
-        (
-            logging.WARNING,
-            f"Failed running worker on dataset dataset_id: AssertionError('{error}')",
-        ),
-    ] + (
-        [
-            (
-                logging.WARNING,
-                "This API helper `update_dataset_state` did not update the cache database",
-            )
-        ]
-        if generator
-        else []
-    ) + [
-        (logging.ERROR, "Ran on 1 dataset: 0 completed, 1 failed"),
+        (logging.WARNING, "No sets to process, stopping."),
     ]
 
 
-def test_run_update_dataset_state_api_error(
+def test_run_initial_dataset_state_error(
     mocker, responses, caplog, mock_dataset_worker, default_dataset
 ):
+    default_dataset.state = DatasetState.Building.value
     mocker.patch(
-        "arkindex_worker.worker.DatasetWorker.list_datasets",
-        return_value=[default_dataset],
-    )
-    mock_dataset_worker.generator = True
-
-    responses.add(
-        responses.PATCH,
-        f"http://testserver/api/v1/datasets/{default_dataset.id}/",
-        status=500,
+        "arkindex_worker.worker.DatasetWorker.list_sets",
+        return_value=[Set(name="train", dataset=default_dataset)],
     )
 
     with pytest.raises(SystemExit):
         mock_dataset_worker.run()
 
-    assert len(responses.calls) == len(BASE_API_CALLS) * 2 + 10
+    assert len(responses.calls) == len(BASE_API_CALLS) * 2
     assert [
         (call.request.method, call.request.url) for call in responses.calls
-    ] == BASE_API_CALLS * 2 + [
-        # We retry 5 times the API call to update the Dataset as Building
-        ("PATCH", f"http://testserver/api/v1/datasets/{default_dataset.id}/"),
-        ("PATCH", f"http://testserver/api/v1/datasets/{default_dataset.id}/"),
-        ("PATCH", f"http://testserver/api/v1/datasets/{default_dataset.id}/"),
-        ("PATCH", f"http://testserver/api/v1/datasets/{default_dataset.id}/"),
-        ("PATCH", f"http://testserver/api/v1/datasets/{default_dataset.id}/"),
-        # We retry 5 times the API call to update the Dataset as in Error
-        ("PATCH", f"http://testserver/api/v1/datasets/{default_dataset.id}/"),
-        ("PATCH", f"http://testserver/api/v1/datasets/{default_dataset.id}/"),
-        ("PATCH", f"http://testserver/api/v1/datasets/{default_dataset.id}/"),
-        ("PATCH", f"http://testserver/api/v1/datasets/{default_dataset.id}/"),
-        ("PATCH", f"http://testserver/api/v1/datasets/{default_dataset.id}/"),
-    ]
+    ] == BASE_API_CALLS * 2
 
-    retries = [3.0, 4.0, 8.0, 16.0]
     assert [(level, message) for _, level, message in caplog.record_tuples] == [
         (logging.INFO, "Loaded Worker Fake worker @ 123412 from API"),
-        (logging.INFO, "Processing Dataset (dataset_id) (1/1)"),
-        (logging.INFO, "Building Dataset (dataset_id) (1/1)"),
-        *[
-            (
-                logging.INFO,
-                f"Retrying arkindex_worker.worker.base.BaseWorker.request in {retry} seconds as it raised ErrorResponse: .",
-            )
-            for retry in retries
-        ],
         (
             logging.WARNING,
-            "An API error occurred while processing dataset dataset_id: 500 Internal Server Error - None",
-        ),
-        *[
-            (
-                logging.INFO,
-                f"Retrying arkindex_worker.worker.base.BaseWorker.request in {retry} seconds as it raised ErrorResponse: .",
-            )
-            for retry in retries
-        ],
-        (
-            logging.ERROR,
-            "Ran on 1 dataset: 0 completed, 1 failed",
+            "Failed running worker on Set (train) from Dataset (dataset_id): AssertionError('When processing a set, its dataset state should be Complete.')",
         ),
+        (logging.ERROR, "Ran on 1 set: 0 completed, 1 failed"),
     ]
 
 
@@ -541,10 +479,9 @@ def test_run_download_dataset_artifact_api_error(
     default_dataset,
 ):
     default_dataset.state = DatasetState.Complete.value
-
     mocker.patch(
-        "arkindex_worker.worker.DatasetWorker.list_datasets",
-        return_value=[default_dataset],
+        "arkindex_worker.worker.DatasetWorker.list_sets",
+        return_value=[Set(name="train", dataset=default_dataset)],
     )
     mocker.patch(
         "arkindex_worker.worker.base.BaseWorker.find_extras_directory",
@@ -574,8 +511,11 @@ def test_run_download_dataset_artifact_api_error(
 
     assert [(level, message) for _, level, message in caplog.record_tuples] == [
         (logging.INFO, "Loaded Worker Fake worker @ 123412 from API"),
-        (logging.INFO, "Processing Dataset (dataset_id) (1/1)"),
-        (logging.INFO, "Downloading data for Dataset (dataset_id) (1/1)"),
+        (
+            logging.INFO,
+            "Retrieving data for Set (train) from Dataset (dataset_id) (1/1)",
+        ),
+        (logging.INFO, "Downloading artifact for Dataset (dataset_id)"),
         *[
             (
                 logging.INFO,
@@ -585,11 +525,11 @@ def test_run_download_dataset_artifact_api_error(
         ],
         (
             logging.WARNING,
-            "An API error occurred while processing dataset dataset_id: 500 Internal Server Error - None",
+            "An API error occurred while processing Set (train) from Dataset (dataset_id): 500 Internal Server Error - None",
         ),
         (
             logging.ERROR,
-            "Ran on 1 dataset: 0 completed, 1 failed",
+            "Ran on 1 set: 0 completed, 1 failed",
         ),
     ]
 
@@ -603,10 +543,9 @@ def test_run_no_downloaded_artifact_error(
     default_dataset,
 ):
     default_dataset.state = DatasetState.Complete.value
-
     mocker.patch(
-        "arkindex_worker.worker.DatasetWorker.list_datasets",
-        return_value=[default_dataset],
+        "arkindex_worker.worker.DatasetWorker.list_sets",
+        return_value=[Set(name="train", dataset=default_dataset)],
     )
     mocker.patch(
         "arkindex_worker.worker.base.BaseWorker.find_extras_directory",
@@ -632,22 +571,22 @@ def test_run_no_downloaded_artifact_error(
 
     assert [(level, message) for _, level, message in caplog.record_tuples] == [
         (logging.INFO, "Loaded Worker Fake worker @ 123412 from API"),
-        (logging.INFO, "Processing Dataset (dataset_id) (1/1)"),
-        (logging.INFO, "Downloading data for Dataset (dataset_id) (1/1)"),
+        (
+            logging.INFO,
+            "Retrieving data for Set (train) from Dataset (dataset_id) (1/1)",
+        ),
+        (logging.INFO, "Downloading artifact for Dataset (dataset_id)"),
         (
             logging.WARNING,
-            "Failed running worker on dataset dataset_id: MissingDatasetArchive('The dataset compressed archive artifact was not found.')",
+            "Failed running worker on Set (train) from Dataset (dataset_id): MissingDatasetArchive('The dataset compressed archive artifact was not found.')",
         ),
         (
             logging.ERROR,
-            "Ran on 1 dataset: 0 completed, 1 failed",
+            "Ran on 1 set: 0 completed, 1 failed",
         ),
     ]
 
 
-@pytest.mark.parametrize(
-    ("generator", "state"), [(True, DatasetState.Open), (False, DatasetState.Complete)]
-)
 def test_run(
     mocker,
     tmp_path,
@@ -656,100 +595,68 @@ def test_run(
     mock_dataset_worker,
     default_dataset,
     default_artifact,
-    generator,
-    state,
 ):
-    mock_dataset_worker.generator = generator
-    default_dataset.state = state.value
-
+    default_dataset.state = DatasetState.Complete.value
     mocker.patch(
-        "arkindex_worker.worker.DatasetWorker.list_datasets",
-        return_value=[default_dataset],
+        "arkindex_worker.worker.DatasetWorker.list_sets",
+        return_value=[Set(name="train", dataset=default_dataset)],
     )
     mocker.patch(
         "arkindex_worker.worker.base.BaseWorker.find_extras_directory",
         return_value=tmp_path,
     )
-    mock_process = mocker.patch("arkindex_worker.worker.DatasetWorker.process_dataset")
-
-    extra_calls = []
-    extra_logs = []
-    if generator:
-        responses.add(
-            responses.PATCH,
-            f"http://testserver/api/v1/datasets/{default_dataset.id}/",
-            status=200,
-            json={},
-        )
-        extra_calls += [
-            ("PATCH", f"http://testserver/api/v1/datasets/{default_dataset.id}/"),
-        ] * 2
-        extra_logs += [
-            (logging.INFO, "Building Dataset (dataset_id) (1/1)"),
-            (
-                logging.WARNING,
-                "This API helper `update_dataset_state` did not update the cache database",
-            ),
-            (logging.INFO, "Completed Dataset (dataset_id) (1/1)"),
-            (
-                logging.WARNING,
-                "This API helper `update_dataset_state` did not update the cache database",
-            ),
-        ]
-    else:
-        archive_path = (
-            FIXTURES_DIR
-            / "extract_parent_archives"
-            / "first_parent"
-            / "arkindex_data.tar.zst"
-        )
-        responses.add(
-            responses.GET,
-            f"http://testserver/api/v1/task/{default_dataset.task_id}/artifacts/",
-            status=200,
-            json=[default_artifact],
-        )
-        responses.add(
-            responses.GET,
-            f"http://testserver/api/v1/task/{default_dataset.task_id}/artifact/dataset_id.tar.zst",
-            status=200,
-            body=archive_path.read_bytes(),
-            content_type="application/zstd",
-        )
-        extra_calls += [
-            (
-                "GET",
-                f"http://testserver/api/v1/task/{default_dataset.task_id}/artifacts/",
-            ),
-            (
-                "GET",
-                f"http://testserver/api/v1/task/{default_dataset.task_id}/artifact/dataset_id.tar.zst",
-            ),
-        ]
-        extra_logs += [
-            (logging.INFO, "Downloading data for Dataset (dataset_id) (1/1)"),
-        ]
+    mock_process = mocker.patch("arkindex_worker.worker.DatasetWorker.process_set")
+
+    archive_path = (
+        FIXTURES_DIR
+        / "extract_parent_archives"
+        / "first_parent"
+        / "arkindex_data.tar.zst"
+    )
+    responses.add(
+        responses.GET,
+        f"http://testserver/api/v1/task/{default_dataset.task_id}/artifacts/",
+        status=200,
+        json=[default_artifact],
+    )
+    responses.add(
+        responses.GET,
+        f"http://testserver/api/v1/task/{default_dataset.task_id}/artifact/dataset_id.tar.zst",
+        status=200,
+        body=archive_path.read_bytes(),
+        content_type="application/zstd",
+    )
 
     mock_dataset_worker.run()
 
     assert mock_process.call_count == 1
 
-    assert len(responses.calls) == len(BASE_API_CALLS) * 2 + len(extra_calls)
+    assert len(responses.calls) == len(BASE_API_CALLS) * 2 + 2
     assert [
         (call.request.method, call.request.url) for call in responses.calls
-    ] == BASE_API_CALLS * 2 + extra_calls
+    ] == BASE_API_CALLS * 2 + [
+        (
+            "GET",
+            f"http://testserver/api/v1/task/{default_dataset.task_id}/artifacts/",
+        ),
+        (
+            "GET",
+            f"http://testserver/api/v1/task/{default_dataset.task_id}/artifact/dataset_id.tar.zst",
+        ),
+    ]
 
     assert [(level, message) for _, level, message in caplog.record_tuples] == [
         (logging.INFO, "Loaded Worker Fake worker @ 123412 from API"),
-        (logging.INFO, "Processing Dataset (dataset_id) (1/1)"),
-        *extra_logs,
-        (logging.INFO, "Ran on 1 dataset: 1 completed, 0 failed"),
+        (
+            logging.INFO,
+            "Retrieving data for Set (train) from Dataset (dataset_id) (1/1)",
+        ),
+        (logging.INFO, "Downloading artifact for Dataset (dataset_id)"),
+        (logging.INFO, "Processing Set (train) from Dataset (dataset_id) (1/1)"),
+        (logging.INFO, "Ran on 1 set: 1 completed, 0 failed"),
     ]
 
 
-@pytest.mark.parametrize(
-    ("generator", "state"), [(True, DatasetState.Open), (False, DatasetState.Complete)]
-)
 def test_run_read_only(
     mocker,
     tmp_path,
@@ -758,90 +665,61 @@ def test_run_read_only(
     mock_dev_dataset_worker,
     default_dataset,
     default_artifact,
-    generator,
-    state,
 ):
-    mock_dev_dataset_worker.generator = generator
-    default_dataset.state = state.value
-
+    default_dataset.state = DatasetState.Complete.value
     mocker.patch(
-        "arkindex_worker.worker.DatasetWorker.list_datasets",
-        return_value=[default_dataset.id],
+        "arkindex_worker.worker.DatasetWorker.list_sets",
+        return_value=[Set(name="train", dataset=default_dataset)],
     )
     mocker.patch(
         "arkindex_worker.worker.base.BaseWorker.find_extras_directory",
         return_value=tmp_path,
     )
-    mock_process = mocker.patch("arkindex_worker.worker.DatasetWorker.process_dataset")
+    mock_process = mocker.patch("arkindex_worker.worker.DatasetWorker.process_set")
 
+    archive_path = (
+        FIXTURES_DIR
+        / "extract_parent_archives"
+        / "first_parent"
+        / "arkindex_data.tar.zst"
+    )
     responses.add(
         responses.GET,
-        f"http://testserver/api/v1/datasets/{default_dataset.id}/",
+        f"http://testserver/api/v1/task/{default_dataset.task_id}/artifacts/",
         status=200,
-        json=default_dataset,
+        json=[default_artifact],
+    )
+    responses.add(
+        responses.GET,
+        f"http://testserver/api/v1/task/{default_dataset.task_id}/artifact/dataset_id.tar.zst",
+        status=200,
+        body=archive_path.read_bytes(),
+        content_type="application/zstd",
     )
-
-    extra_calls = []
-    extra_logs = []
-    if generator:
-        extra_logs += [
-            (logging.INFO, "Building Dataset (dataset_id) (1/1)"),
-            (
-                logging.WARNING,
-                "Cannot update dataset as this worker is in read-only mode",
-            ),
-            (logging.INFO, "Completed Dataset (dataset_id) (1/1)"),
-            (
-                logging.WARNING,
-                "Cannot update dataset as this worker is in read-only mode",
-            ),
-        ]
-    else:
-        archive_path = (
-            FIXTURES_DIR
-            / "extract_parent_archives"
-            / "first_parent"
-            / "arkindex_data.tar.zst"
-        )
-        responses.add(
-            responses.GET,
-            f"http://testserver/api/v1/task/{default_dataset.task_id}/artifacts/",
-            status=200,
-            json=[default_artifact],
-        )
-        responses.add(
-            responses.GET,
-            f"http://testserver/api/v1/task/{default_dataset.task_id}/artifact/dataset_id.tar.zst",
-            status=200,
-            body=archive_path.read_bytes(),
-            content_type="application/zstd",
-        )
-        extra_calls += [
-            (
-                "GET",
-                f"http://testserver/api/v1/task/{default_dataset.task_id}/artifacts/",
-            ),
-            (
-                "GET",
-                f"http://testserver/api/v1/task/{default_dataset.task_id}/artifact/dataset_id.tar.zst",
-            ),
-        ]
-        extra_logs += [
-            (logging.INFO, "Downloading data for Dataset (dataset_id) (1/1)"),
-        ]
 
     mock_dev_dataset_worker.run()
 
     assert mock_process.call_count == 1
 
-    assert len(responses.calls) == 1 + len(extra_calls)
+    assert len(responses.calls) == 2
     assert [(call.request.method, call.request.url) for call in responses.calls] == [
-        ("GET", f"http://testserver/api/v1/datasets/{default_dataset.id}/")
-    ] + extra_calls
+        (
+            "GET",
+            f"http://testserver/api/v1/task/{default_dataset.task_id}/artifacts/",
+        ),
+        (
+            "GET",
+            f"http://testserver/api/v1/task/{default_dataset.task_id}/artifact/dataset_id.tar.zst",
+        ),
+    ]
 
     assert [(level, message) for _, level, message in caplog.record_tuples] == [
         (logging.WARNING, "Running without any extra configuration"),
-        (logging.INFO, "Processing Dataset (dataset_id) (1/1)"),
-        *extra_logs,
-        (logging.INFO, "Ran on 1 dataset: 1 completed, 0 failed"),
+        (
+            logging.INFO,
+            "Retrieving data for Set (train) from Dataset (dataset_id) (1/1)",
+        ),
+        (logging.INFO, "Downloading artifact for Dataset (dataset_id)"),
+        (logging.INFO, "Processing Set (train) from Dataset (dataset_id) (1/1)"),
+        (logging.INFO, "Ran on 1 set: 1 completed, 0 failed"),
     ]