arkindex-base-worker 0.3.6rc5__py3-none-any.whl → 0.3.7.post1__py3-none-any.whl

tests/test_elements_worker/test_cli.py

@@ -24,7 +24,7 @@ def test_cli_default(monkeypatch):
         )
     )
 
-    monkeypatch.setenv("TASK_ELEMENTS", path)
+    monkeypatch.setenv("TASK_ELEMENTS", str(path))
     monkeypatch.setattr(sys, "argv", ["worker"])
     worker = ElementsWorker()
     worker.configure()

tests/test_elements_worker/test_dataset.py

@@ -4,13 +4,13 @@ import logging
 import pytest
 from apistar.exceptions import ErrorResponse
 
-from arkindex_worker.models import Dataset
+from arkindex_worker.models import Dataset, Element, Set
 from arkindex_worker.worker.dataset import DatasetState
 from tests.conftest import PROCESS_ID
 from tests.test_elements_worker import BASE_API_CALLS
 
 
-def test_list_process_datasets_readonly_error(mock_dataset_worker):
+def test_list_process_sets_readonly_error(mock_dataset_worker):
     # Set worker in read_only mode
     mock_dataset_worker.worker_run_id = None
     assert mock_dataset_worker.is_read_only
@@ -18,73 +18,91 @@ def test_list_process_datasets_readonly_error(mock_dataset_worker):
     with pytest.raises(
         AssertionError, match="This helper is not available in read-only mode."
     ):
-        mock_dataset_worker.list_process_datasets()
+        mock_dataset_worker.list_process_sets()
 
 
-def test_list_process_datasets_api_error(responses, mock_dataset_worker):
+def test_list_process_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_process_datasets())
+        next(mock_dataset_worker.list_process_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_process_datasets(
+def test_list_process_sets(
     responses,
     mock_dataset_worker,
 ):
     expected_results = [
         {
-            "id": "dataset_1",
-            "name": "Dataset 1",
-            "description": "My first great dataset",
-            "sets": ["train", "val", "test"],
-            "state": "open",
-            "corpus_id": "corpus_id",
-            "creator": "test@teklia.com",
-            "task_id": "task_id_1",
+            "id": "set_1",
+            "dataset": {
+                "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_1",
+            },
+            "set_name": "train",
         },
         {
-            "id": "dataset_2",
-            "name": "Dataset 2",
-            "description": "My second great dataset",
-            "sets": ["train", "val"],
-            "state": "complete",
-            "corpus_id": "corpus_id",
-            "creator": "test@teklia.com",
-            "task_id": "task_id_2",
+            "id": "set_2",
+            "dataset": {
+                "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_1",
+            },
+            "set_name": "val",
         },
         {
-            "id": "dataset_3",
-            "name": "Dataset 3 (TRASHME)",
-            "description": "My third dataset, in error",
-            "sets": ["nonsense", "random set"],
-            "state": "error",
-            "corpus_id": "corpus_id",
-            "creator": "test@teklia.com",
-            "task_id": "task_id_3",
+            "id": "set_3",
+            "dataset": {
+                "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_2",
+            },
+            "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,
@@ -93,50 +111,54 @@ def test_list_process_datasets(
         },
     )
 
-    for idx, dataset in enumerate(mock_dataset_worker.list_process_datasets()):
-        assert isinstance(dataset, Dataset)
-        assert dataset == expected_results[idx]
+    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(
     ("payload", "error"),
     [
-        # Dataset
+        # Set
         (
-            {"dataset": None},
-            "dataset shouldn't be null and should be a Dataset",
+            {"dataset_set": None},
+            "dataset_set shouldn't be null and should be a Set",
         ),
         (
-            {"dataset": "not Dataset type"},
-            "dataset shouldn't be null and should be a Dataset",
+            {"dataset_set": "not Set type"},
+            "dataset_set shouldn't be null and should be a Set",
         ),
     ],
 )
-def test_list_dataset_elements_wrong_param_dataset(mock_dataset_worker, payload, error):
+def test_list_set_elements_wrong_param_dataset_set(mock_dataset_worker, payload, error):
     with pytest.raises(AssertionError, match=error):
-        mock_dataset_worker.list_dataset_elements(**payload)
+        mock_dataset_worker.list_set_elements(**payload)
 
 
-def test_list_dataset_elements_api_error(
-    responses, mock_dataset_worker, default_dataset
+def test_list_set_elements_api_error(
+    responses, mock_dataset_worker, default_dataset, default_train_set
 ):
+    query_params = f"?set={default_train_set.name}&with_count=true"
     responses.add(
         responses.GET,
-        f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/",
+        f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/{query_params}",
         status=500,
     )
 
     with pytest.raises(
         Exception, match="Stopping pagination as data will be incomplete"
     ):
-        next(mock_dataset_worker.list_dataset_elements(dataset=default_dataset))
+        next(mock_dataset_worker.list_set_elements(dataset_set=default_train_set))
 
     assert len(responses.calls) == len(BASE_API_CALLS) + 5
     assert [
@@ -145,69 +167,40 @@ def test_list_dataset_elements_api_error(
         # The API call is retried 5 times
         (
             "GET",
-            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?with_count=true",
+            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/{query_params}",
         ),
         (
             "GET",
-            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?with_count=true",
+            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/{query_params}",
         ),
         (
             "GET",
-            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?with_count=true",
+            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/{query_params}",
         ),
         (
             "GET",
-            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?with_count=true",
+            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/{query_params}",
         ),
         (
             "GET",
-            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?with_count=true",
+            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/{query_params}",
         ),
     ]
 
 
-def test_list_dataset_elements(
+def test_list_set_elements(
     responses,
     mock_dataset_worker,
     default_dataset,
+    default_train_set,
 ):
     expected_results = [
         {
-            "set": "set_1",
-            "element": {
-                "id": "0000",
-                "type": "page",
-                "name": "Test",
-                "corpus": {},
-                "thumbnail_url": None,
-                "zone": {},
-                "best_classes": None,
-                "has_children": None,
-                "worker_version_id": None,
-                "worker_run_id": None,
-            },
-        },
-        {
-            "set": "set_1",
-            "element": {
-                "id": "1111",
-                "type": "page",
-                "name": "Test 2",
-                "corpus": {},
-                "thumbnail_url": None,
-                "zone": {},
-                "best_classes": None,
-                "has_children": None,
-                "worker_version_id": None,
-                "worker_run_id": None,
-            },
-        },
-        {
-            "set": "set_2",
+            "set": "train",
             "element": {
-                "id": "2222",
+                "id": "element_1",
                 "type": "page",
-                "name": "Test 3",
+                "name": "1",
                 "corpus": {},
                 "thumbnail_url": None,
                 "zone": {},
@@ -216,41 +209,29 @@ def test_list_dataset_elements(
                 "worker_version_id": None,
                 "worker_run_id": None,
             },
-        },
-        {
-            "set": "set_3",
-            "element": {
-                "id": "3333",
-                "type": "page",
-                "name": "Test 4",
-                "corpus": {},
-                "thumbnail_url": None,
-                "zone": {},
-                "best_classes": None,
-                "has_children": None,
-                "worker_version_id": None,
-                "worker_run_id": None,
-            },
-        },
+        }
     ]
+    expected_results.append({**expected_results[-1]})
+    expected_results[-1]["element"]["id"] = "element_2"
+    expected_results[-1]["element"]["name"] = "2"
+
+    query_params = f"?set={default_train_set.name}&with_count=true"
     responses.add(
         responses.GET,
-        f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/",
+        f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/{query_params}",
         status=200,
         json={
-            "count": 4,
+            "count": 2,
             "next": None,
             "results": expected_results,
         },
     )
 
     for idx, element in enumerate(
-        mock_dataset_worker.list_dataset_elements(dataset=default_dataset)
+        mock_dataset_worker.list_set_elements(dataset_set=default_train_set)
     ):
-        assert element == (
-            expected_results[idx]["set"],
-            expected_results[idx]["element"],
-        )
+        assert isinstance(element, Element)
+        assert element == expected_results[idx]["element"]
 
     assert len(responses.calls) == len(BASE_API_CALLS) + 1
     assert [
@@ -258,8 +239,8 @@ def test_list_dataset_elements(
     ] == BASE_API_CALLS + [
         (
             "GET",
-            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?with_count=true",
-        ),
+            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/{query_params}",
+        )
     ]