arkindex-base-worker 0.3.6rc4__py3-none-any.whl → 0.3.7__py3-none-any.whl

tests/test_elements_worker/test_classifications.py

@@ -1,6 +1,6 @@
 import json
 import re
-from uuid import UUID, uuid4
+from uuid import UUID
 
 import pytest
 from apistar.exceptions import ErrorResponse
@@ -10,6 +10,10 @@ from arkindex_worker.models import Element
 
 from . import BASE_API_CALLS
 
+# Special string used to know if the `arg_name` passed in
+# `pytest.mark.parametrize` should be removed from the payload
+DELETE_PARAMETER = "DELETE_PARAMETER"
+
 
 def test_get_ml_class_id_load_classes(responses, mock_elements_worker):
     corpus_id = "11111111-1111-1111-1111-111111111111"
@@ -190,54 +194,116 @@ def test_retrieve_ml_class_not_in_cache(responses, mock_elements_worker):
     ]
 
 
-def test_create_classification_wrong_element(mock_elements_worker):
-    with pytest.raises(
-        AssertionError,
-        match="element shouldn't be null and should be an Element or CachedElement",
-    ):
-        mock_elements_worker.create_classification(
-            element=None,
-            ml_class="a_class",
-            confidence=0.42,
-            high_confidence=True,
-        )
-
-    with pytest.raises(
-        AssertionError,
-        match="element shouldn't be null and should be an Element or CachedElement",
-    ):
+@pytest.mark.parametrize(
+    ("arg_name", "data", "error_message"),
+    [
+        # Wrong element
+        (
+            "element",
+            None,
+            "element shouldn't be null and should be an Element or CachedElement",
+        ),
+        (
+            "element",
+            "not element type",
+            "element shouldn't be null and should be an Element or CachedElement",
+        ),
+        # Wrong ml_class
+        (
+            "ml_class",
+            None,
+            "ml_class shouldn't be null and should be of type str",
+        ),
+        (
+            "ml_class",
+            1234,
+            "ml_class shouldn't be null and should be of type str",
+        ),
+        # Wrong confidence
+        (
+            "confidence",
+            None,
+            "confidence shouldn't be null and should be a float in [0..1] range",
+        ),
+        (
+            "confidence",
+            "wrong confidence",
+            "confidence shouldn't be null and should be a float in [0..1] range",
+        ),
+        (
+            "confidence",
+            0,
+            "confidence shouldn't be null and should be a float in [0..1] range",
+        ),
+        (
+            "confidence",
+            2.00,
+            "confidence shouldn't be null and should be a float in [0..1] range",
+        ),
+        # Wrong high_confidence
+        (
+            "high_confidence",
+            None,
+            "high_confidence shouldn't be null and should be of type bool",
+        ),
+        (
+            "high_confidence",
+            "wrong high_confidence",
+            "high_confidence shouldn't be null and should be of type bool",
+        ),
+    ],
+)
+def test_create_classification_wrong_data(
+    arg_name, data, error_message, mock_elements_worker
+):
+    mock_elements_worker.classes = {"a_class": "0000"}
+    with pytest.raises(AssertionError, match=re.escape(error_message)):
         mock_elements_worker.create_classification(
-            element="not element type",
-            ml_class="a_class",
-            confidence=0.42,
-            high_confidence=True,
+            **{
+                "element": Element({"id": "12341234-1234-1234-1234-123412341234"}),
+                "ml_class": "a_class",
+                "confidence": 0.42,
+                "high_confidence": True,
+                # Overwrite with wrong data
+                arg_name: data,
+            }
         )
 
 
-def test_create_classification_wrong_ml_class(mock_elements_worker, responses):
+def test_create_classification_api_error(responses, mock_elements_worker):
+    mock_elements_worker.classes = {"a_class": "0000"}
     elt = Element({"id": "12341234-1234-1234-1234-123412341234"})
+    responses.add(
+        responses.POST,
+        "http://testserver/api/v1/classifications/",
+        status=500,
+    )
 
-    with pytest.raises(
-        AssertionError, match="ml_class shouldn't be null and should be of type str"
-    ):
+    with pytest.raises(ErrorResponse):
         mock_elements_worker.create_classification(
             element=elt,
-            ml_class=None,
+            ml_class="a_class",
             confidence=0.42,
             high_confidence=True,
         )
 
-    with pytest.raises(
-        AssertionError, match="ml_class shouldn't be null and should be of type str"
-    ):
-        mock_elements_worker.create_classification(
-            element=elt,
-            ml_class=1234,
-            confidence=0.42,
-            high_confidence=True,
-        )
+    assert len(responses.calls) == len(BASE_API_CALLS) + 5
+    assert [
+        (call.request.method, call.request.url) for call in responses.calls
+    ] == BASE_API_CALLS + [
+        # We retry 5 times the API call
+        ("POST", "http://testserver/api/v1/classifications/"),
+        ("POST", "http://testserver/api/v1/classifications/"),
+        ("POST", "http://testserver/api/v1/classifications/"),
+        ("POST", "http://testserver/api/v1/classifications/"),
+        ("POST", "http://testserver/api/v1/classifications/"),
+    ]
+
+
+def test_create_classification_create_ml_class(mock_elements_worker, responses):
+    elt = Element({"id": "12341234-1234-1234-1234-123412341234"})
 
-    # Automatically create a missing class !
+    # Automatically create a missing class!
     responses.add(
         responses.POST,
         "http://testserver/api/v1/corpus/11111111-1111-1111-1111-111111111111/classes/",
@@ -259,9 +325,6 @@ def test_create_classification_wrong_ml_class(mock_elements_worker, responses):
     )
 
     # Check a class & classification has been created
-    for call in responses.calls:
-        print(call.request.url, call.request.body)
-
     assert [
         (call.request.url, json.loads(call.request.body))
         for call in responses.calls[-2:]
@@ -283,119 +346,6 @@ def test_create_classification_wrong_ml_class(mock_elements_worker, responses):
     ]
 
 
-def test_create_classification_wrong_confidence(mock_elements_worker):
-    mock_elements_worker.classes = {"a_class": "0000"}
-    elt = Element({"id": "12341234-1234-1234-1234-123412341234"})
-    with pytest.raises(
-        AssertionError,
-        match=re.escape(
-            "confidence shouldn't be null and should be a float in [0..1] range"
-        ),
-    ):
-        mock_elements_worker.create_classification(
-            element=elt,
-            ml_class="a_class",
-            confidence=None,
-            high_confidence=True,
-        )
-
-    with pytest.raises(
-        AssertionError,
-        match=re.escape(
-            "confidence shouldn't be null and should be a float in [0..1] range"
-        ),
-    ):
-        mock_elements_worker.create_classification(
-            element=elt,
-            ml_class="a_class",
-            confidence="wrong confidence",
-            high_confidence=True,
-        )
-
-    with pytest.raises(
-        AssertionError,
-        match=re.escape(
-            "confidence shouldn't be null and should be a float in [0..1] range"
-        ),
-    ):
-        mock_elements_worker.create_classification(
-            element=elt,
-            ml_class="a_class",
-            confidence=0,
-            high_confidence=True,
-        )
-
-    with pytest.raises(
-        AssertionError,
-        match=re.escape(
-            "confidence shouldn't be null and should be a float in [0..1] range"
-        ),
-    ):
-        mock_elements_worker.create_classification(
-            element=elt,
-            ml_class="a_class",
-            confidence=2.00,
-            high_confidence=True,
-        )
-
-
-def test_create_classification_wrong_high_confidence(mock_elements_worker):
-    mock_elements_worker.classes = {"a_class": "0000"}
-    elt = Element({"id": "12341234-1234-1234-1234-123412341234"})
-
-    with pytest.raises(
-        AssertionError,
-        match="high_confidence shouldn't be null and should be of type bool",
-    ):
-        mock_elements_worker.create_classification(
-            element=elt,
-            ml_class="a_class",
-            confidence=0.42,
-            high_confidence=None,
-        )
-
-    with pytest.raises(
-        AssertionError,
-        match="high_confidence shouldn't be null and should be of type bool",
-    ):
-        mock_elements_worker.create_classification(
-            element=elt,
-            ml_class="a_class",
-            confidence=0.42,
-            high_confidence="wrong high_confidence",
-        )
-
-
-def test_create_classification_api_error(responses, mock_elements_worker):
-    mock_elements_worker.classes = {"a_class": "0000"}
-    elt = Element({"id": "12341234-1234-1234-1234-123412341234"})
-    responses.add(
-        responses.POST,
-        "http://testserver/api/v1/classifications/",
-        status=500,
-    )
-
-    with pytest.raises(ErrorResponse):
-        mock_elements_worker.create_classification(
-            element=elt,
-            ml_class="a_class",
-            confidence=0.42,
-            high_confidence=True,
-        )
-
-    assert len(responses.calls) == len(BASE_API_CALLS) + 5
-    assert [
-        (call.request.method, call.request.url) for call in responses.calls
-    ] == BASE_API_CALLS + [
-        # We retry 5 times the API call
-        ("POST", "http://testserver/api/v1/classifications/"),
-        ("POST", "http://testserver/api/v1/classifications/"),
-        ("POST", "http://testserver/api/v1/classifications/"),
-        ("POST", "http://testserver/api/v1/classifications/"),
-        ("POST", "http://testserver/api/v1/classifications/"),
-    ]
-
-
 def test_create_classification(responses, mock_elements_worker):
     mock_elements_worker.classes = {"a_class": "0000"}
     elt = Element({"id": "12341234-1234-1234-1234-123412341234"})
@@ -519,303 +469,165 @@ def test_create_classification_duplicate_worker_run(responses, mock_elements_wor
     }
 
 
-def test_create_classifications_wrong_element(mock_elements_worker):
-    with pytest.raises(
-        AssertionError,
-        match="element shouldn't be null and should be an Element or CachedElement",
-    ):
+@pytest.mark.parametrize(
+    ("arg_name", "data", "error_message"),
+    [
+        (
+            "element",
+            None,
+            "element shouldn't be null and should be an Element or CachedElement",
+        ),
+        (
+            "element",
+            "not element type",
+            "element shouldn't be null and should be an Element or CachedElement",
+        ),
+        (
+            "classifications",
+            None,
+            "classifications shouldn't be null and should be of type list",
+        ),
+        (
+            "classifications",
+            1234,
+            "classifications shouldn't be null and should be of type list",
+        ),
+    ],
+)
+def test_create_classifications_wrong_data(
+    arg_name, data, error_message, mock_elements_worker
+):
+    with pytest.raises(AssertionError, match=error_message):
         mock_elements_worker.create_classifications(
-            element=None,
-            classifications=[
-                {
-                    "ml_class_id": "uuid1",
-                    "confidence": 0.75,
-                    "high_confidence": False,
-                },
-                {
-                    "ml_class_id": "uuid2",
-                    "confidence": 0.25,
-                    "high_confidence": False,
-                },
-            ],
+            **{
+                "element": Element({"id": "12341234-1234-1234-1234-123412341234"}),
+                "classifications": [
+                    {
+                        "ml_class": "cat",
+                        "confidence": 0.75,
+                        "high_confidence": False,
+                    },
+                    {
+                        "ml_class": "dog",
+                        "confidence": 0.25,
+                        "high_confidence": False,
+                    },
+                ],
+                # Overwrite with wrong data
+                arg_name: data,
+            },
         )
 
+
+@pytest.mark.parametrize(
+    ("arg_name", "data", "error_message"),
+    [
+        # Wrong classifications > ml_class
+        (
+            "ml_class",
+            DELETE_PARAMETER,
+            "ml_class shouldn't be null and should be of type str",
+        ),
+        (
+            "ml_class",
+            None,
+            "ml_class shouldn't be null and should be of type str",
+        ),
+        (
+            "ml_class",
+            1234,
+            "ml_class shouldn't be null and should be of type str",
+        ),
+        # Wrong classifications > confidence
+        (
+            "confidence",
+            DELETE_PARAMETER,
+            "confidence shouldn't be null and should be a float in [0..1] range",
+        ),
+        (
+            "confidence",
+            None,
+            "confidence shouldn't be null and should be a float in [0..1] range",
+        ),
+        (
+            "confidence",
+            "wrong confidence",
+            "confidence shouldn't be null and should be a float in [0..1] range",
+        ),
+        (
+            "confidence",
+            0,
+            "confidence shouldn't be null and should be a float in [0..1] range",
+        ),
+        (
+            "confidence",
+            2.00,
+            "confidence shouldn't be null and should be a float in [0..1] range",
+        ),
+        # Wrong classifications > high_confidence
+        (
+            "high_confidence",
+            "wrong high_confidence",
+            "high_confidence should be of type bool",
+        ),
+    ],
+)
+def test_create_classifications_wrong_classifications_data(
+    arg_name, data, error_message, mock_elements_worker
+):
+    all_data = {
+        "element": Element({"id": "12341234-1234-1234-1234-123412341234"}),
+        "classifications": [
+            {
+                "ml_class": "cat",
+                "confidence": 0.75,
+                "high_confidence": False,
+            },
+            {
+                "ml_class": "dog",
+                "confidence": 0.25,
+                "high_confidence": False,
+                # Overwrite with wrong data
+                arg_name: data,
+            },
+        ],
+    }
+    if data == DELETE_PARAMETER:
+        del all_data["classifications"][1][arg_name]
+
     with pytest.raises(
         AssertionError,
-        match="element shouldn't be null and should be an Element or CachedElement",
+        match=re.escape(
+            f"Classification at index 1 in classifications: {error_message}"
+        ),
     ):
-        mock_elements_worker.create_classifications(
-            element="not element type",
-            classifications=[
-                {
-                    "ml_class_id": "uuid1",
-                    "confidence": 0.75,
-                    "high_confidence": False,
-                },
-                {
-                    "ml_class_id": "uuid2",
-                    "confidence": 0.25,
-                    "high_confidence": False,
-                },
-            ],
-        )
+        mock_elements_worker.create_classifications(**all_data)
 
 
-def test_create_classifications_wrong_classifications(mock_elements_worker):
+def test_create_classifications_api_error(responses, mock_elements_worker):
+    mock_elements_worker.classes = {"cat": "0000", "dog": "1111"}
+    responses.add(
+        responses.POST,
+        "http://testserver/api/v1/classification/bulk/",
+        status=500,
+    )
     elt = Element({"id": "12341234-1234-1234-1234-123412341234"})
+    classes = [
+        {
+            "ml_class": "cat",
+            "confidence": 0.75,
+            "high_confidence": False,
+        },
+        {
+            "ml_class": "dog",
+            "confidence": 0.25,
+            "high_confidence": False,
+        },
+    ]
 
-    with pytest.raises(
-        AssertionError,
-        match="classifications shouldn't be null and should be of type list",
-    ):
+    with pytest.raises(ErrorResponse):
         mock_elements_worker.create_classifications(
-            element=elt,
-            classifications=None,
-        )
-
-    with pytest.raises(
-        AssertionError,
-        match="classifications shouldn't be null and should be of type list",
-    ):
-        mock_elements_worker.create_classifications(
-            element=elt,
-            classifications=1234,
-        )
-
-    with pytest.raises(
-        AssertionError,
-        match="Classification at index 1 in classifications: ml_class_id shouldn't be null and should be of type str",
-    ):
-        mock_elements_worker.create_classifications(
-            element=elt,
-            classifications=[
-                {
-                    "ml_class_id": str(uuid4()),
-                    "confidence": 0.75,
-                    "high_confidence": False,
-                },
-                {
-                    "ml_class_id": 0.25,
-                    "high_confidence": False,
-                },
-            ],
-        )
-
-    with pytest.raises(
-        AssertionError,
-        match="Classification at index 1 in classifications: ml_class_id shouldn't be null and should be of type str",
-    ):
-        mock_elements_worker.create_classifications(
-            element=elt,
-            classifications=[
-                {
-                    "ml_class_id": str(uuid4()),
-                    "confidence": 0.75,
-                    "high_confidence": False,
-                },
-                {
-                    "ml_class_id": None,
-                    "confidence": 0.25,
-                    "high_confidence": False,
-                },
-            ],
-        )
-
-    with pytest.raises(
-        AssertionError,
-        match="Classification at index 1 in classifications: ml_class_id shouldn't be null and should be of type str",
-    ):
-        mock_elements_worker.create_classifications(
-            element=elt,
-            classifications=[
-                {
-                    "ml_class_id": str(uuid4()),
-                    "confidence": 0.75,
-                    "high_confidence": False,
-                },
-                {
-                    "ml_class_id": 1234,
-                    "confidence": 0.25,
-                    "high_confidence": False,
-                },
-            ],
-        )
-
-    with pytest.raises(
-        ValueError,
-        match="Classification at index 1 in classifications: ml_class_id is not a valid uuid.",
-    ):
-        mock_elements_worker.create_classifications(
-            element=elt,
-            classifications=[
-                {
-                    "ml_class_id": str(uuid4()),
-                    "confidence": 0.75,
-                    "high_confidence": False,
-                },
-                {
-                    "ml_class_id": "not_an_uuid",
-                    "confidence": 0.25,
-                    "high_confidence": False,
-                },
-            ],
-        )
-
-    with pytest.raises(
-        AssertionError,
-        match=re.escape(
-            "Classification at index 1 in classifications: confidence shouldn't be null and should be a float in [0..1] range"
-        ),
-    ):
-        mock_elements_worker.create_classifications(
-            element=elt,
-            classifications=[
-                {
-                    "ml_class_id": str(uuid4()),
-                    "confidence": 0.75,
-                    "high_confidence": False,
-                },
-                {
-                    "ml_class_id": str(uuid4()),
-                    "high_confidence": False,
-                },
-            ],
-        )
-
-    with pytest.raises(
-        AssertionError,
-        match=re.escape(
-            "Classification at index 1 in classifications: confidence shouldn't be null and should be a float in [0..1] range"
-        ),
-    ):
-        mock_elements_worker.create_classifications(
-            element=elt,
-            classifications=[
-                {
-                    "ml_class_id": str(uuid4()),
-                    "confidence": 0.75,
-                    "high_confidence": False,
-                },
-                {
-                    "ml_class_id": str(uuid4()),
-                    "confidence": None,
-                    "high_confidence": False,
-                },
-            ],
-        )
-
-    with pytest.raises(
-        AssertionError,
-        match=re.escape(
-            "Classification at index 1 in classifications: confidence shouldn't be null and should be a float in [0..1] range"
-        ),
-    ):
-        mock_elements_worker.create_classifications(
-            element=elt,
-            classifications=[
-                {
-                    "ml_class_id": str(uuid4()),
-                    "confidence": 0.75,
-                    "high_confidence": False,
-                },
-                {
-                    "ml_class_id": str(uuid4()),
-                    "confidence": "wrong confidence",
-                    "high_confidence": False,
-                },
-            ],
-        )
-
-    with pytest.raises(
-        AssertionError,
-        match=re.escape(
-            "Classification at index 1 in classifications: confidence shouldn't be null and should be a float in [0..1] range"
-        ),
-    ):
-        mock_elements_worker.create_classifications(
-            element=elt,
-            classifications=[
-                {
-                    "ml_class_id": str(uuid4()),
-                    "confidence": 0.75,
-                    "high_confidence": False,
-                },
-                {
-                    "ml_class_id": str(uuid4()),
-                    "confidence": 0,
-                    "high_confidence": False,
-                },
-            ],
-        )
-
-    with pytest.raises(
-        AssertionError,
-        match=re.escape(
-            "Classification at index 1 in classifications: confidence shouldn't be null and should be a float in [0..1] range"
-        ),
-    ):
-        mock_elements_worker.create_classifications(
-            element=elt,
-            classifications=[
-                {
-                    "ml_class_id": str(uuid4()),
-                    "confidence": 0.75,
-                    "high_confidence": False,
-                },
-                {
-                    "ml_class_id": str(uuid4()),
-                    "confidence": 2.00,
-                    "high_confidence": False,
-                },
-            ],
-        )
-
-    with pytest.raises(
-        AssertionError,
-        match=re.escape(
-            "Classification at index 1 in classifications: high_confidence should be of type bool"
-        ),
-    ):
-        mock_elements_worker.create_classifications(
-            element=elt,
-            classifications=[
-                {
-                    "ml_class_id": str(uuid4()),
-                    "confidence": 0.75,
-                    "high_confidence": False,
-                },
-                {
-                    "ml_class_id": str(uuid4()),
-                    "confidence": 0.25,
-                    "high_confidence": "wrong high_confidence",
-                },
-            ],
-        )
-
-
-def test_create_classifications_api_error(responses, mock_elements_worker):
-    responses.add(
-        responses.POST,
-        "http://testserver/api/v1/classification/bulk/",
-        status=500,
-    )
-    elt = Element({"id": "12341234-1234-1234-1234-123412341234"})
-    classes = [
-        {
-            "ml_class_id": str(uuid4()),
-            "confidence": 0.75,
-            "high_confidence": False,
-        },
-        {
-            "ml_class_id": str(uuid4()),
-            "confidence": 0.25,
-            "high_confidence": False,
-        },
-    ]
-
-    with pytest.raises(ErrorResponse):
-        mock_elements_worker.create_classifications(
-            element=elt, classifications=classes
+            element=elt, classifications=classes
         )
 
     assert len(responses.calls) == len(BASE_API_CALLS) + 5
@@ -831,57 +643,96 @@ def test_create_classifications_api_error(responses, mock_elements_worker):
     ]
 
 
-def test_create_classifications(responses, mock_elements_worker_with_cache):
-    # Set MLClass in cache
-    portrait_uuid = str(uuid4())
-    landscape_uuid = str(uuid4())
-    mock_elements_worker_with_cache.classes = {
-        "portrait": portrait_uuid,
-        "landscape": landscape_uuid,
-    }
-
-    elt = CachedElement.create(id="12341234-1234-1234-1234-123412341234", type="thing")
-    classes = [
-        {
-            "ml_class_id": portrait_uuid,
-            "confidence": 0.75,
-            "high_confidence": False,
-        },
-        {
-            "ml_class_id": landscape_uuid,
-            "confidence": 0.25,
-            "high_confidence": False,
-        },
-    ]
+def test_create_classifications_create_ml_class(mock_elements_worker, responses):
+    elt = Element({"id": "12341234-1234-1234-1234-123412341234"})
 
+    # Automatically create a missing class!
+    responses.add(
+        responses.POST,
+        "http://testserver/api/v1/corpus/11111111-1111-1111-1111-111111111111/classes/",
+        status=201,
+        json={"id": "new-ml-class-1234"},
+    )
     responses.add(
         responses.POST,
         "http://testserver/api/v1/classification/bulk/",
-        status=200,
+        status=201,
         json={
             "parent": str(elt.id),
             "worker_run_id": "56785678-5678-5678-5678-567856785678",
             "classifications": [
                 {
                     "id": "00000000-0000-0000-0000-000000000000",
-                    "ml_class": portrait_uuid,
+                    "ml_class": "new-ml-class-1234",
                     "confidence": 0.75,
                     "high_confidence": False,
                     "state": "pending",
                 },
-                {
-                    "id": "11111111-1111-1111-1111-111111111111",
-                    "ml_class": landscape_uuid,
-                    "confidence": 0.25,
-                    "high_confidence": False,
-                    "state": "pending",
-                },
             ],
         },
     )
+    mock_elements_worker.classes = {"another_class": "0000"}
+    mock_elements_worker.create_classifications(
+        element=elt,
+        classifications=[
+            {
+                "ml_class": "a_class",
+                "confidence": 0.75,
+                "high_confidence": False,
+            }
+        ],
+    )
 
-    mock_elements_worker_with_cache.create_classifications(
-        element=elt, classifications=classes
+    # Check a class & classification has been created
+    assert len(responses.calls) == len(BASE_API_CALLS) + 2
+    assert [
+        (call.request.method, call.request.url) for call in responses.calls
+    ] == BASE_API_CALLS + [
+        (
+            "POST",
+            "http://testserver/api/v1/corpus/11111111-1111-1111-1111-111111111111/classes/",
+        ),
+        ("POST", "http://testserver/api/v1/classification/bulk/"),
+    ]
+
+    assert json.loads(responses.calls[-2].request.body) == {"name": "a_class"}
+    assert json.loads(responses.calls[-1].request.body) == {
+        "parent": "12341234-1234-1234-1234-123412341234",
+        "worker_run_id": "56785678-5678-5678-5678-567856785678",
+        "classifications": [
+            {
+                "ml_class": "new-ml-class-1234",
+                "confidence": 0.75,
+                "high_confidence": False,
+            }
+        ],
+    }
+
+
+def test_create_classifications(responses, mock_elements_worker):
+    mock_elements_worker.classes = {"portrait": "0000", "landscape": "1111"}
+    elt = Element({"id": "12341234-1234-1234-1234-123412341234"})
+    responses.add(
+        responses.POST,
+        "http://testserver/api/v1/classification/bulk/",
+        status=200,
+        json={"classifications": []},
+    )
+
+    mock_elements_worker.create_classifications(
+        element=elt,
+        classifications=[
+            {
+                "ml_class": "portrait",
+                "confidence": 0.75,
+                "high_confidence": False,
+            },
+            {
+                "ml_class": "landscape",
+                "confidence": 0.25,
+                "high_confidence": False,
+            },
+        ],
     )
 
     assert len(responses.calls) == len(BASE_API_CALLS) + 1
@@ -894,52 +745,24 @@ def test_create_classifications(responses, mock_elements_worker_with_cache):
     assert json.loads(responses.calls[-1].request.body) == {
         "parent": str(elt.id),
         "worker_run_id": "56785678-5678-5678-5678-567856785678",
-        "classifications": classes,
+        "classifications": [
+            {
+                "confidence": 0.75,
+                "high_confidence": False,
+                "ml_class": "0000",
+            },
+            {
+                "confidence": 0.25,
+                "high_confidence": False,
+                "ml_class": "1111",
+            },
+        ],
     }
 
-    # Check that created classifications were properly stored in SQLite cache
-    assert list(CachedClassification.select()) == [
-        CachedClassification(
-            id=UUID("00000000-0000-0000-0000-000000000000"),
-            element_id=UUID(elt.id),
-            class_name="portrait",
-            confidence=0.75,
-            state="pending",
-            worker_run_id=UUID("56785678-5678-5678-5678-567856785678"),
-        ),
-        CachedClassification(
-            id=UUID("11111111-1111-1111-1111-111111111111"),
-            element_id=UUID(elt.id),
-            class_name="landscape",
-            confidence=0.25,
-            state="pending",
-            worker_run_id=UUID("56785678-5678-5678-5678-567856785678"),
-        ),
-    ]
-
 
-def test_create_classifications_not_in_cache(
-    responses, mock_elements_worker_with_cache
-):
-    """
-    CreateClassifications using ID that are not in `.classes` attribute.
-    Will load corpus MLClass to insert the corresponding name in Cache.
-    """
-    portrait_uuid = str(uuid4())
-    landscape_uuid = str(uuid4())
+def test_create_classifications_with_cache(responses, mock_elements_worker_with_cache):
+    mock_elements_worker_with_cache.classes = {"portrait": "0000", "landscape": "1111"}
     elt = CachedElement.create(id="12341234-1234-1234-1234-123412341234", type="thing")
-    classes = [
-        {
-            "ml_class_id": portrait_uuid,
-            "confidence": 0.75,
-            "high_confidence": False,
-        },
-        {
-            "ml_class_id": landscape_uuid,
-            "confidence": 0.25,
-            "high_confidence": False,
-        },
-    ]
 
     responses.add(
         responses.POST,
@@ -951,14 +774,14 @@ def test_create_classifications_not_in_cache(
             "classifications": [
                 {
                     "id": "00000000-0000-0000-0000-000000000000",
-                    "ml_class": portrait_uuid,
+                    "ml_class": "0000",
                     "confidence": 0.75,
                     "high_confidence": False,
                     "state": "pending",
                 },
                 {
                     "id": "11111111-1111-1111-1111-111111111111",
-                    "ml_class": landscape_uuid,
+                    "ml_class": "1111",
                     "confidence": 0.25,
                     "high_confidence": False,
                     "state": "pending",
@@ -966,42 +789,45 @@ def test_create_classifications_not_in_cache(
             ],
         },
     )
-    responses.add(
-        responses.GET,
-        f"http://testserver/api/v1/corpus/{mock_elements_worker_with_cache.corpus_id}/classes/",
-        status=200,
-        json={
-            "count": 2,
-            "next": None,
-            "results": [
-                {
-                    "id": portrait_uuid,
-                    "name": "portrait",
-                },
-                {"id": landscape_uuid, "name": "landscape"},
-            ],
-        },
-    )
 
     mock_elements_worker_with_cache.create_classifications(
-        element=elt, classifications=classes
+        element=elt,
+        classifications=[
+            {
+                "ml_class": "portrait",
+                "confidence": 0.75,
+                "high_confidence": False,
+            },
+            {
+                "ml_class": "landscape",
+                "confidence": 0.25,
+                "high_confidence": False,
+            },
+        ],
     )
 
-    assert len(responses.calls) == len(BASE_API_CALLS) + 2
+    assert len(responses.calls) == len(BASE_API_CALLS) + 1
     assert [
         (call.request.method, call.request.url) for call in responses.calls
     ] == BASE_API_CALLS + [
         ("POST", "http://testserver/api/v1/classification/bulk/"),
-        (
-            "GET",
-            f"http://testserver/api/v1/corpus/{mock_elements_worker_with_cache.corpus_id}/classes/",
-        ),
     ]
 
-    assert json.loads(responses.calls[-2].request.body) == {
+    assert json.loads(responses.calls[-1].request.body) == {
         "parent": str(elt.id),
         "worker_run_id": "56785678-5678-5678-5678-567856785678",
-        "classifications": classes,
+        "classifications": [
+            {
+                "confidence": 0.75,
+                "high_confidence": False,
+                "ml_class": "0000",
+            },
+            {
+                "confidence": 0.25,
+                "high_confidence": False,
+                "ml_class": "1111",
+            },
+        ],
     }
 
     # Check that created classifications were properly stored in SQLite cache