arkindex-base-worker 0.4.0rc2__py3-none-any.whl → 0.4.0rc4__py3-none-any.whl

tests/test_elements_worker/test_worker.py

@@ -1,89 +1,526 @@
 import json
 import sys
+from argparse import Namespace
+from uuid import UUID
 
 import pytest
-from apistar.exceptions import ErrorResponse
 
-from arkindex_worker.cache import CachedElement
+from arkindex.exceptions import ErrorResponse
+from arkindex_worker.cache import (
+    SQL_VERSION,
+    CachedElement,
+    create_version_table,
+    init_cache_db,
+)
+from arkindex_worker.models import Element
 from arkindex_worker.worker import ActivityState, ElementsWorker
+from arkindex_worker.worker.dataset import DatasetState
+from arkindex_worker.worker.process import ProcessMode
 
 from . import BASE_API_CALLS
 
-TEST_VERSION_ID = "test_123"
-TEST_SLUG = "some_slug"
 
+def test_database_arg(mocker, mock_elements_worker, tmp_path):
+    database_path = tmp_path / "my_database.sqlite"
+    init_cache_db(database_path)
+    create_version_table()
+
+    mocker.patch(
+        "arkindex_worker.worker.base.argparse.ArgumentParser.parse_args",
+        return_value=Namespace(
+            element=["volumeid", "pageid"],
+            verbose=False,
+            elements_list=None,
+            database=database_path,
+            dev=False,
+            set=[],
+        ),
+    )
 
-def test_get_worker_version(fake_dummy_worker):
-    api_client = fake_dummy_worker.api_client
+    worker = ElementsWorker(support_cache=True)
+    worker.configure()
 
-    response = {"worker": {"slug": TEST_SLUG}}
+    assert worker.use_cache is True
+    assert worker.cache_path == database_path
 
-    api_client.add_response("RetrieveWorkerVersion", response, id=TEST_VERSION_ID)
 
-    with pytest.deprecated_call(match="WorkerVersion usage is deprecated."):
-        res = fake_dummy_worker.get_worker_version(TEST_VERSION_ID)
+def test_database_arg_cache_missing_version_table(
+    mocker, mock_elements_worker, tmp_path
+):
+    database_path = tmp_path / "my_database.sqlite"
+    database_path.touch()
+
+    mocker.patch(
+        "arkindex_worker.worker.base.argparse.ArgumentParser.parse_args",
+        return_value=Namespace(
+            element=["volumeid", "pageid"],
+            verbose=False,
+            elements_list=None,
+            database=database_path,
+            dev=False,
+            set=[],
+        ),
+    )
 
-    assert res == response
-    assert fake_dummy_worker._worker_version_cache[TEST_VERSION_ID] == response
+    worker = ElementsWorker(support_cache=True)
+    with pytest.raises(
+        AssertionError,
+        match=f"The SQLite database {database_path} does not have the correct cache version, it should be {SQL_VERSION}",
+    ):
+        worker.configure()
 
 
-def test_get_worker_version__uses_cache(fake_dummy_worker):
-    api_client = fake_dummy_worker.api_client
+def test_readonly(responses, mock_elements_worker):
+    """Test readonly worker does not trigger any API calls"""
 
-    response = {"worker": {"slug": TEST_SLUG}}
+    # Setup the worker as read-only
+    mock_elements_worker.worker_run_id = None
+    assert mock_elements_worker.is_read_only is True
 
-    api_client.add_response("RetrieveWorkerVersion", response, id=TEST_VERSION_ID)
+    out = mock_elements_worker.update_activity("1234-deadbeef", ActivityState.Processed)
 
-    with pytest.deprecated_call(match="WorkerVersion usage is deprecated."):
-        response_1 = fake_dummy_worker.get_worker_version(TEST_VERSION_ID)
+    # update_activity returns False in very specific cases
+    assert out is True
+    assert len(responses.calls) == len(BASE_API_CALLS)
+    assert [
+        (call.request.method, call.request.url) for call in responses.calls
+    ] == BASE_API_CALLS
 
-    with pytest.deprecated_call(match="WorkerVersion usage is deprecated."):
-        response_2 = fake_dummy_worker.get_worker_version(TEST_VERSION_ID)
 
-    assert response_1 == response
-    assert response_1 == response_2
+def test_get_elements_elements_list_arg_wrong_type(
+    monkeypatch, tmp_path, mock_elements_worker
+):
+    elements_path = tmp_path / "elements.json"
+    elements_path.write_text("{}")
 
-    # assert that only one call to the API
-    assert len(api_client.history) == 1
-    assert not api_client.responses
+    monkeypatch.setenv("TASK_ELEMENTS", str(elements_path))
+    worker = ElementsWorker()
+    worker.configure()
 
+    with pytest.raises(AssertionError, match="Elements list must be a list"):
+        worker.get_elements()
 
-def test_get_worker_version_slug(mocker, fake_dummy_worker):
-    fake_dummy_worker.get_worker_version = mocker.MagicMock()
-    fake_dummy_worker.get_worker_version.return_value = {
-        "id": TEST_VERSION_ID,
-        "worker": {"slug": "mock_slug"},
-    }
 
-    with pytest.deprecated_call(match="WorkerVersion usage is deprecated."):
-        slug = fake_dummy_worker.get_worker_version_slug(TEST_VERSION_ID)
-    assert slug == "mock_slug"
+def test_get_elements_elements_list_arg_empty_list(
+    monkeypatch, tmp_path, mock_elements_worker
+):
+    elements_path = tmp_path / "elements.json"
+    elements_path.write_text("[]")
 
+    monkeypatch.setenv("TASK_ELEMENTS", str(elements_path))
+    worker = ElementsWorker()
+    worker.configure()
+
+    with pytest.raises(AssertionError, match="No elements in elements list"):
+        worker.get_elements()
 
-def test_get_worker_version_slug_none(fake_dummy_worker):
-    # WARNING: pytest.deprecated_call must be placed BEFORE pytest.raises, otherwise `match` argument won't be checked
-    with (
-        pytest.deprecated_call(match="WorkerVersion usage is deprecated."),
-        pytest.raises(ValueError, match="No worker version ID"),
+
+def test_get_elements_elements_list_arg_missing_id(
+    monkeypatch, tmp_path, mock_elements_worker
+):
+    elements_path = tmp_path / "elements.json"
+    elements_path.write_text(json.dumps([{"type": "volume"}]))
+
+    monkeypatch.setenv("TASK_ELEMENTS", str(elements_path))
+    worker = ElementsWorker()
+    worker.configure()
+
+    elt_list = worker.get_elements()
+
+    assert elt_list == []
+
+
+def test_get_elements_elements_list_arg_not_uuid(
+    monkeypatch, tmp_path, mock_elements_worker
+):
+    elements_path = tmp_path / "elements.json"
+    elements_path.write_text(
+        json.dumps(
+            [
+                {"id": "volumeid", "type": "volume"},
+                {"id": "pageid", "type": "page"},
+                {"id": "actid", "type": "act"},
+                {"id": "surfaceid", "type": "surface"},
+            ]
+        )
+    )
+
+    monkeypatch.setenv("TASK_ELEMENTS", str(elements_path))
+    worker = ElementsWorker()
+    worker.configure()
+
+    with pytest.raises(
+        Exception,
+        match="These element IDs are invalid: volumeid, pageid, actid, surfaceid",
     ):
-        fake_dummy_worker.get_worker_version_slug(None)
+        worker.get_elements()
+
+
+def test_get_elements_elements_list_arg(monkeypatch, tmp_path, mock_elements_worker):
+    elements_path = tmp_path / "elements.json"
+    elements_path.write_text(
+        json.dumps(
+            [
+                {"id": "11111111-1111-1111-1111-111111111111", "type": "volume"},
+                {"id": "22222222-2222-2222-2222-222222222222", "type": "page"},
+                {"id": "33333333-3333-3333-3333-333333333333", "type": "act"},
+            ]
+        )
+    )
 
+    monkeypatch.setenv("TASK_ELEMENTS", str(elements_path))
+    worker = ElementsWorker()
+    worker.configure()
 
-def test_readonly(responses, mock_elements_worker):
-    """Test readonly worker does not trigger any API calls"""
+    elt_list = worker.get_elements()
 
-    # Setup the worker as read-only
-    mock_elements_worker.worker_run_id = None
-    assert mock_elements_worker.is_read_only is True
+    assert elt_list == [
+        "11111111-1111-1111-1111-111111111111",
+        "22222222-2222-2222-2222-222222222222",
+        "33333333-3333-3333-3333-333333333333",
+    ]
 
-    out = mock_elements_worker.update_activity("1234-deadbeef", ActivityState.Processed)
 
-    # update_activity returns False in very specific cases
-    assert out is True
-    assert len(responses.calls) == len(BASE_API_CALLS)
-    assert [
-        (call.request.method, call.request.url) for call in responses.calls
-    ] == BASE_API_CALLS
+def test_get_elements_element_arg_not_uuid(mocker, mock_elements_worker):
+    mocker.patch(
+        "arkindex_worker.worker.base.argparse.ArgumentParser.parse_args",
+        return_value=Namespace(
+            element=["volumeid", "pageid"],
+            config={},
+            verbose=False,
+            elements_list=None,
+            database=None,
+            dev=True,
+            set=[],
+        ),
+    )
+
+    worker = ElementsWorker()
+    worker.configure()
+
+    with pytest.raises(
+        Exception, match="These element IDs are invalid: volumeid, pageid"
+    ):
+        worker.get_elements()
+
+
+def test_get_elements_element_arg(mocker, mock_elements_worker):
+    mocker.patch(
+        "arkindex_worker.worker.base.argparse.ArgumentParser.parse_args",
+        return_value=Namespace(
+            element=[
+                "11111111-1111-1111-1111-111111111111",
+                "22222222-2222-2222-2222-222222222222",
+            ],
+            config={},
+            verbose=False,
+            elements_list=None,
+            database=None,
+            dev=True,
+            set=[],
+        ),
+    )
+
+    worker = ElementsWorker()
+    worker.configure()
+
+    elt_list = worker.get_elements()
+
+    assert elt_list == [
+        "11111111-1111-1111-1111-111111111111",
+        "22222222-2222-2222-2222-222222222222",
+    ]
+
+
+def test_get_elements_dataset_set_arg(responses, mocker, mock_elements_worker):
+    mocker.patch(
+        "arkindex_worker.worker.base.argparse.ArgumentParser.parse_args",
+        return_value=Namespace(
+            element=[],
+            config={},
+            verbose=False,
+            elements_list=None,
+            database=None,
+            dev=True,
+            set=[(UUID("11111111-1111-1111-1111-111111111111"), "train")],
+        ),
+    )
+
+    # Mock RetrieveDataset call
+    responses.add(
+        responses.GET,
+        "http://testserver/api/v1/datasets/11111111-1111-1111-1111-111111111111/",
+        status=200,
+        json={
+            "id": "11111111-1111-1111-1111-111111111111",
+            "name": "My dataset",
+            "description": "A dataset about cats.",
+            "sets": ["train", "dev", "test"],
+            "state": DatasetState.Complete.value,
+        },
+        content_type="application/json",
+    )
+
+    # Mock ListSetElements call
+    element = {
+        "id": "22222222-2222-2222-2222-222222222222",
+        "type": "page",
+        "name": "1",
+        "corpus": {
+            "id": "11111111-1111-1111-1111-111111111111",
+        },
+        "thumbnail_url": "http://example.com",
+        "zone": {
+            "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+            "polygon": [[0, 0], [0, 0], [0, 0]],
+            "image": {
+                "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+                "path": "string",
+                "width": 0,
+                "height": 0,
+                "url": "http://example.com",
+                "s3_url": "string",
+                "status": "checked",
+                "server": {
+                    "display_name": "string",
+                    "url": "http://example.com",
+                    "max_width": 2147483647,
+                    "max_height": 2147483647,
+                },
+            },
+            "url": "http://example.com",
+        },
+        "rotation_angle": 0,
+        "mirrored": False,
+        "created": "2019-08-24T14:15:22Z",
+        "classes": [
+            {
+                "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+                "ml_class": {
+                    "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+                    "name": "string",
+                },
+                "state": "pending",
+                "confidence": 0,
+                "high_confidence": True,
+                "worker_run": {
+                    "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+                    "summary": "string",
+                },
+            }
+        ],
+        "metadata": [
+            {
+                "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+                "type": "text",
+                "name": "string",
+                "value": "string",
+                "dates": [{"type": "exact", "year": 0, "month": 1, "day": 1}],
+            }
+        ],
+        "transcriptions": [
+            {
+                "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+                "text": "string",
+                "confidence": 0,
+                "orientation": "horizontal-lr",
+                "worker_run": {
+                    "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+                    "summary": "string",
+                },
+            }
+        ],
+        "has_children": True,
+        "worker_run": {
+            "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+            "summary": "string",
+        },
+        "confidence": 1,
+    }
+    responses.add(
+        responses.GET,
+        "http://testserver/api/v1/datasets/11111111-1111-1111-1111-111111111111/elements/?set=train&with_count=true",
+        status=200,
+        json={
+            "next": None,
+            "previous": None,
+            "results": [
+                {
+                    "set": "train",
+                    "element": element,
+                }
+            ],
+            "count": 1,
+        },
+        content_type="application/json",
+    )
+
+    worker = ElementsWorker()
+    worker.configure()
+
+    elt_list = worker.get_elements()
+
+    assert elt_list == [
+        Element(**element),
+    ]
+
+
+def test_get_elements_dataset_set_api(responses, mocker, mock_elements_worker):
+    # Mock ListProcessSets call
+    responses.add(
+        responses.GET,
+        "http://testserver/api/v1/process/aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeffff/sets/",
+        status=200,
+        json={
+            "next": None,
+            "previous": None,
+            "results": [
+                {
+                    "id": "33333333-3333-3333-3333-333333333333",
+                    "dataset": {"id": "11111111-1111-1111-1111-111111111111"},
+                    "set_name": "train",
+                }
+            ],
+            "count": 1,
+        },
+        content_type="application/json",
+    )
+
+    # Mock ListSetElements call
+    element = {
+        "id": "22222222-2222-2222-2222-222222222222",
+        "type": "page",
+        "name": "1",
+        "corpus": {
+            "id": "11111111-1111-1111-1111-111111111111",
+        },
+        "thumbnail_url": "http://example.com",
+        "zone": {
+            "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+            "polygon": [[0, 0], [0, 0], [0, 0]],
+            "image": {
+                "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+                "path": "string",
+                "width": 0,
+                "height": 0,
+                "url": "http://example.com",
+                "s3_url": "string",
+                "status": "checked",
+                "server": {
+                    "display_name": "string",
+                    "url": "http://example.com",
+                    "max_width": 2147483647,
+                    "max_height": 2147483647,
+                },
+            },
+            "url": "http://example.com",
+        },
+        "rotation_angle": 0,
+        "mirrored": False,
+        "created": "2019-08-24T14:15:22Z",
+        "classes": [
+            {
+                "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+                "ml_class": {
+                    "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+                    "name": "string",
+                },
+                "state": "pending",
+                "confidence": 0,
+                "high_confidence": True,
+                "worker_run": {
+                    "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+                    "summary": "string",
+                },
+            }
+        ],
+        "metadata": [
+            {
+                "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+                "type": "text",
+                "name": "string",
+                "value": "string",
+                "dates": [{"type": "exact", "year": 0, "month": 1, "day": 1}],
+            }
+        ],
+        "transcriptions": [
+            {
+                "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+                "text": "string",
+                "confidence": 0,
+                "orientation": "horizontal-lr",
+                "worker_run": {
+                    "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+                    "summary": "string",
+                },
+            }
+        ],
+        "has_children": True,
+        "worker_run": {
+            "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08",
+            "summary": "string",
+        },
+        "confidence": 1,
+    }
+    responses.add(
+        responses.GET,
+        "http://testserver/api/v1/datasets/11111111-1111-1111-1111-111111111111/elements/?set=train&with_count=true",
+        status=200,
+        json={
+            "next": None,
+            "previous": None,
+            "results": [
+                {
+                    "set": "train",
+                    "element": element,
+                }
+            ],
+            "count": 1,
+        },
+        content_type="application/json",
+    )
+
+    # Update ProcessMode to Dataset
+    mock_elements_worker.process_information["mode"] = ProcessMode.Dataset
+
+    elt_list = mock_elements_worker.get_elements()
+
+    assert elt_list == [
+        Element(**element),
+    ]
+
+
+def test_get_elements_both_args_error(mocker, mock_elements_worker, tmp_path):
+    elements_path = tmp_path / "elements.json"
+    elements_path.write_text(
+        json.dumps(
+            [
+                {"id": "volumeid", "type": "volume"},
+                {"id": "pageid", "type": "page"},
+                {"id": "actid", "type": "act"},
+                {"id": "surfaceid", "type": "surface"},
+            ]
+        )
+    )
+    mocker.patch(
+        "arkindex_worker.worker.base.argparse.ArgumentParser.parse_args",
+        return_value=Namespace(
+            element=["anotherid", "againanotherid"],
+            verbose=False,
+            elements_list=elements_path.open(),
+            database=None,
+            dev=False,
+            set=[],
+        ),
+    )
+
+    worker = ElementsWorker()
+    worker.configure()
+
+    with pytest.raises(
+        AssertionError, match="elements-list and element CLI args shouldn't be both set"
+    ):
+        worker.get_elements()
 
 
 @pytest.mark.usefixtures("_mock_worker_run_api")
@@ -114,43 +551,6 @@ def test_activities_dev_mode(mocker):
     assert worker.store_activity is False
 
 
-@pytest.mark.usefixtures("_mock_worker_run_api")
-def test_update_call(responses, mock_elements_worker):
-    """Test an update call with feature enabled triggers an API call"""
-    responses.add(
-        responses.PUT,
-        "http://testserver/api/v1/workers/versions/56785678-5678-5678-5678-567856785678/activity/",
-        status=200,
-        json={
-            "element_id": "1234-deadbeef",
-            "process_id": "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeffff",
-            "state": "processed",
-        },
-    )
-
-    out = mock_elements_worker.update_activity("1234-deadbeef", ActivityState.Processed)
-
-    # Check the response received by worker
-    assert out is True
-
-    assert len(responses.calls) == len(BASE_API_CALLS) + 1
-    assert [
-        (call.request.method, call.request.url) for call in responses.calls
-    ] == BASE_API_CALLS + [
-        (
-            "PUT",
-            "http://testserver/api/v1/workers/versions/56785678-5678-5678-5678-567856785678/activity/",
-        ),
-    ]
-
-    # Check the request sent by worker
-    assert json.loads(responses.calls[-1].request.body) == {
-        "element_id": "1234-deadbeef",
-        "process_id": "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeffff",
-        "state": "processed",
-    }
-
-
 @pytest.mark.usefixtures("_mock_activity_calls")
 @pytest.mark.parametrize(
     ("process_exception", "final_state"),
@@ -312,3 +712,40 @@ def test_start_activity_error(
     assert logger.error.call_args_list == [
         mocker.call("Ran on 1 element: 0 completed, 1 failed")
     ]
+
+
+@pytest.mark.usefixtures("_mock_worker_run_api")
+def test_update_activity(responses, mock_elements_worker):
+    """Test an update call with feature enabled triggers an API call"""
+    responses.add(
+        responses.PUT,
+        "http://testserver/api/v1/workers/versions/56785678-5678-5678-5678-567856785678/activity/",
+        status=200,
+        json={
+            "element_id": "1234-deadbeef",
+            "process_id": "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeffff",
+            "state": "processed",
+        },
+    )
+
+    out = mock_elements_worker.update_activity("1234-deadbeef", ActivityState.Processed)
+
+    # Check the response received by worker
+    assert out is True
+
+    assert len(responses.calls) == len(BASE_API_CALLS) + 1
+    assert [
+        (call.request.method, call.request.url) for call in responses.calls
+    ] == BASE_API_CALLS + [
+        (
+            "PUT",
+            "http://testserver/api/v1/workers/versions/56785678-5678-5678-5678-567856785678/activity/",
+        ),
+    ]
+
+    # Check the request sent by worker
+    assert json.loads(responses.calls[-1].request.body) == {
+        "element_id": "1234-deadbeef",
+        "process_id": "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeffff",
+        "state": "processed",
+    }