arkindex-base-worker 0.3.7rc4__py3-none-any.whl → 0.5.0a1__py3-none-any.whl

tests/test_elements_worker/test_cli.py

@@ -2,7 +2,6 @@ import json
 import sys
 import tempfile
 from pathlib import Path
-from uuid import UUID
 
 import pytest
 
@@ -58,13 +57,6 @@ def test_cli_arg_elements_list_given(mocker):
     path.unlink()
 
 
-def test_cli_arg_element_one_given_not_uuid(mocker):
-    mocker.patch.object(sys, "argv", ["worker", "--element", "1234"])
-    worker = ElementsWorker()
-    with pytest.raises(SystemExit):
-        worker.configure()
-
-
 @pytest.mark.usefixtures("_mock_worker_run_api")
 def test_cli_arg_element_one_given(mocker):
     mocker.patch.object(
@@ -73,7 +65,7 @@ def test_cli_arg_element_one_given(mocker):
     worker = ElementsWorker()
     worker.configure()
 
-    assert worker.args.element == [UUID("12341234-1234-1234-1234-123412341234")]
+    assert worker.args.element == ["12341234-1234-1234-1234-123412341234"]
     # elements_list is None because TASK_ELEMENTS environment variable isn't set
     assert not worker.args.elements_list
 
@@ -94,8 +86,8 @@ def test_cli_arg_element_many_given(mocker):
     worker.configure()
 
     assert worker.args.element == [
-        UUID("12341234-1234-1234-1234-123412341234"),
-        UUID("43214321-4321-4321-4321-432143214321"),
+        "12341234-1234-1234-1234-123412341234",
+        "43214321-4321-4321-4321-432143214321",
     ]
     # elements_list is None because TASK_ELEMENTS environment variable isn't set
     assert not worker.args.elements_list

tests/test_elements_worker/test_corpus.py

@@ -0,0 +1,168 @@
+import re
+import uuid
+
+import pytest
+
+from arkindex.exceptions import ErrorResponse
+from arkindex_worker.worker.corpus import CorpusExportState
+from tests import CORPUS_ID
+from tests.test_elements_worker import BASE_API_CALLS
+
+
+def test_download_export_not_a_uuid(responses, mock_elements_worker):
+    with pytest.raises(ValueError, match="export_id is not a valid uuid."):
+        mock_elements_worker.download_export("mon export")
+
+
+def test_download_export(responses, mock_elements_worker):
+    responses.add(
+        responses.GET,
+        "http://testserver/api/v1/export/aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeffff/",
+        status=302,
+        body=b"some SQLite export",
+        content_type="application/x-sqlite3",
+        stream=True,
+    )
+
+    export = mock_elements_worker.download_export(
+        "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeffff"
+    )
+    assert export.name == "/tmp/aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeffff"
+
+    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",
+            "http://testserver/api/v1/export/aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeffff/",
+        ),
+    ]
+
+
+def mock_list_exports_call(responses, export_id):
+    responses.add(
+        responses.GET,
+        f"http://testserver/api/v1/corpus/{CORPUS_ID}/export/",
+        status=200,
+        json={
+            "count": len(CorpusExportState),
+            "next": None,
+            "results": [
+                {
+                    "id": str(uuid.uuid4())
+                    if state != CorpusExportState.Done
+                    else export_id,
+                    "created": "2019-08-24T14:15:22Z",
+                    "updated": "2019-08-24T14:15:22Z",
+                    "corpus_id": CORPUS_ID,
+                    "user": {
+                        "id": 0,
+                        "email": "user@example.com",
+                        "display_name": "User",
+                    },
+                    "state": state.value,
+                    "source": "default",
+                }
+                for state in CorpusExportState
+            ],
+        },
+    )
+
+
+def test_download_latest_export_list_error(responses, mock_elements_worker):
+    responses.add(
+        responses.GET,
+        f"http://testserver/api/v1/corpus/{CORPUS_ID}/export/",
+        status=418,
+    )
+
+    with pytest.raises(
+        Exception, match="Stopping pagination as data will be incomplete"
+    ):
+        mock_elements_worker.download_latest_export()
+
+    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/corpus/{CORPUS_ID}/export/"),
+        ("GET", f"http://testserver/api/v1/corpus/{CORPUS_ID}/export/"),
+        ("GET", f"http://testserver/api/v1/corpus/{CORPUS_ID}/export/"),
+        ("GET", f"http://testserver/api/v1/corpus/{CORPUS_ID}/export/"),
+        ("GET", f"http://testserver/api/v1/corpus/{CORPUS_ID}/export/"),
+    ]
+
+
+def test_download_latest_export_no_available_exports(responses, mock_elements_worker):
+    responses.add(
+        responses.GET,
+        f"http://testserver/api/v1/corpus/{CORPUS_ID}/export/",
+        status=200,
+        json={
+            "count": 0,
+            "next": None,
+            "results": [],
+        },
+    )
+
+    with pytest.raises(
+        AssertionError,
+        match=re.escape(
+            f'No available exports found for the corpus ({CORPUS_ID}) with state "Done".'
+        ),
+    ):
+        mock_elements_worker.download_latest_export()
+
+    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/corpus/{CORPUS_ID}/export/"),
+    ]
+
+
+def test_download_latest_export_download_error(responses, mock_elements_worker):
+    export_id = str(uuid.uuid4())
+    mock_list_exports_call(responses, export_id)
+    responses.add(
+        responses.GET,
+        f"http://testserver/api/v1/export/{export_id}/",
+        status=418,
+    )
+
+    with pytest.raises(ErrorResponse):
+        mock_elements_worker.download_latest_export()
+
+    assert len(responses.calls) == len(BASE_API_CALLS) + 2
+    assert [
+        (call.request.method, call.request.url) for call in responses.calls
+    ] == BASE_API_CALLS + [
+        ("GET", f"http://testserver/api/v1/corpus/{CORPUS_ID}/export/"),
+        ("GET", f"http://testserver/api/v1/export/{export_id}/"),
+    ]
+
+
+def test_download_latest_export(responses, mock_elements_worker):
+    export_id = str(uuid.uuid4())
+    mock_list_exports_call(responses, export_id)
+    responses.add(
+        responses.GET,
+        f"http://testserver/api/v1/export/{export_id}/",
+        status=302,
+        body=b"some SQLite export",
+        content_type="application/x-sqlite3",
+        stream=True,
+    )
+
+    export = mock_elements_worker.download_latest_export()
+    assert export.name == f"/tmp/{export_id}"
+
+    assert len(responses.calls) == len(BASE_API_CALLS) + 2
+    assert [
+        (call.request.method, call.request.url) for call in responses.calls
+    ] == BASE_API_CALLS + [
+        ("GET", f"http://testserver/api/v1/corpus/{CORPUS_ID}/export/"),
+        ("GET", f"http://testserver/api/v1/export/{export_id}/"),
+    ]

tests/test_elements_worker/test_dataset.py

@@ -2,15 +2,15 @@ import json
 import logging
 
 import pytest
-from apistar.exceptions import ErrorResponse
 
-from arkindex_worker.models import Dataset
+from arkindex.exceptions import ErrorResponse
+from arkindex_worker.models import Dataset, Element, Set
 from arkindex_worker.worker.dataset import DatasetState
-from tests.conftest import PROCESS_ID
+from tests 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,85 +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/",
-        status=500,
+        f"http://testserver/api/v1/process/{PROCESS_ID}/sets/",
+        status=418,
     )
 
     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": "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,
@@ -105,67 +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]["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(
     ("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)
 
 
-@pytest.mark.parametrize(
-    "sets",
-    [
-        ["set_1"],
-        ["set_1", "set_2", "set_3"],
-        ["set_1", "set_2", "set_3", "set_4"],
-    ],
-)
-def test_list_dataset_elements_api_error(
-    responses, mock_dataset_worker, sets, default_dataset
+def test_list_set_elements_api_error(
+    responses, mock_dataset_worker, default_dataset, default_train_set
 ):
-    default_dataset.selected_sets = sets
-    query_params = (
-        "?with_count=true"
-        if sets == default_dataset.sets
-        else "?set=set_1&with_count=true"
-    )
+    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/{query_params}",
-        status=500,
+        status=418,
     )
 
     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 [
@@ -195,99 +188,60 @@ def test_list_dataset_elements_api_error(
     ]
 
 
-@pytest.mark.parametrize(
-    "sets",
-    [
-        ["set_1"],
-        ["set_1", "set_2", "set_3"],
-        ["set_1", "set_2", "set_3", "set_4"],
-    ],
-)
-def test_list_dataset_elements(
+def test_list_set_elements(
     responses,
     mock_dataset_worker,
-    sets,
     default_dataset,
+    default_train_set,
 ):
-    default_dataset.selected_sets = sets
-
-    dataset_elements = []
-    for split in default_dataset.sets:
-        index = split[-1]
-        dataset_elements.append(
-            {
-                "set": split,
-                "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,
-                },
-            }
-        )
-        if split == "set_1":
-            dataset_elements.append({**dataset_elements[-1]})
-            dataset_elements[-1]["element"]["name"] = f"Test {index} (bis)"
-
-    # All sets are selected, we call the unfiltered endpoint once
-    if default_dataset.sets == default_dataset.selected_sets:
-        expected_results = dataset_elements
-        responses.add(
-            responses.GET,
-            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?with_count=true",
-            status=200,
-            json={
-                "count": len(expected_results),
-                "next": None,
-                "results": expected_results,
+    expected_results = [
+        {
+            "set": "train",
+            "element": {
+                "id": "element_1",
+                "type": "page",
+                "name": "1",
+                "corpus": {},
+                "thumbnail_url": None,
+                "zone": {},
+                "best_classes": None,
+                "has_children": None,
+                "worker_version_id": None,
+                "worker_run_id": None,
             },
-        )
-        expected_calls = [
-            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?with_count=true"
-        ]
-
-    # Not all sets are selected, we call the filtered endpoint multiple times, once per set
-    else:
-        expected_results, expected_calls = [], []
-        for selected_set in default_dataset.selected_sets:
-            partial_results = [
-                element
-                for element in dataset_elements
-                if element["set"] == selected_set
-            ]
-            expected_results += partial_results
-            responses.add(
-                responses.GET,
-                f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?set={selected_set}&with_count=true",
-                status=200,
-                json={
-                    "count": len(partial_results),
-                    "next": None,
-                    "results": partial_results,
-                },
-            )
-            expected_calls += [
-                f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/?set={selected_set}&with_count=true"
-            ]
+        }
+    ]
+    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/{query_params}",
+        status=200,
+        json={
+            "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) + len(expected_calls)
+    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", expected_call) for expected_call in expected_calls]
+    ] == BASE_API_CALLS + [
+        (
+            "GET",
+            f"http://testserver/api/v1/datasets/{default_dataset.id}/elements/{query_params}",
+        )
+    ]
 
 
 @pytest.mark.parametrize(
@@ -367,7 +321,7 @@ def test_update_dataset_state_api_error(
     responses.add(
         responses.PATCH,
         f"http://testserver/api/v1/datasets/{default_dataset.id}/",
-        status=500,
+        status=418,
     )
 
     with pytest.raises(ErrorResponse):
@@ -376,16 +330,11 @@ def test_update_dataset_state_api_error(
             state=DatasetState.Building,
         )
 
-    assert len(responses.calls) == len(BASE_API_CALLS) + 5
+    assert len(responses.calls) == len(BASE_API_CALLS) + 1
     assert [
         (call.request.method, call.request.url) for call in responses.calls
     ] == BASE_API_CALLS + [
-        # We retry 5 times the API call
-        ("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}/"),
+        ("PATCH", f"http://testserver/api/v1/datasets/{default_dataset.id}/")
     ]