arkindex-base-worker 0.3.7rc9__py3-none-any.whl → 0.4.0__py3-none-any.whl

tests/test_dataset_worker.py

@@ -3,18 +3,21 @@ import uuid
 from argparse import ArgumentTypeError
 
 import pytest
-from apistar.exceptions import ErrorResponse
 
+from arkindex.exceptions import ErrorResponse
 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 arkindex_worker.worker.dataset import (
+    DatasetState,
+    MissingDatasetArchive,
+    check_dataset_set,
+)
+from tests import FIXTURES_DIR, PROCESS_ID
 from tests.test_elements_worker import BASE_API_CALLS
 
 RANDOM_UUID = uuid.uuid4()
 
 
-@pytest.fixture()
+@pytest.fixture
 def tmp_archive(tmp_path):
     archive = tmp_path / "test_archive.tar.zst"
     archive.touch()
@@ -63,22 +66,17 @@ def test_download_dataset_artifact_list_api_error(
     responses.add(
         responses.GET,
         f"http://testserver/api/v1/task/{task_id}/artifacts/",
-        status=500,
+        status=418,
     )
 
     with pytest.raises(ErrorResponse):
         mock_dataset_worker.download_dataset_artifact(default_dataset)
 
-    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 + [
-        # The API call is retried 5 times
-        ("GET", f"http://testserver/api/v1/task/{task_id}/artifacts/"),
-        ("GET", f"http://testserver/api/v1/task/{task_id}/artifacts/"),
-        ("GET", f"http://testserver/api/v1/task/{task_id}/artifacts/"),
-        ("GET", f"http://testserver/api/v1/task/{task_id}/artifacts/"),
-        ("GET", f"http://testserver/api/v1/task/{task_id}/artifacts/"),
+        ("GET", f"http://testserver/api/v1/task/{task_id}/artifacts/")
     ]
 
 
@@ -116,22 +114,17 @@ def test_download_dataset_artifact_download_api_error(
     responses.add(
         responses.GET,
         f"http://testserver/api/v1/task/{task_id}/artifact/dataset_id.tar.zst",
-        status=500,
+        status=418,
     )
 
     with pytest.raises(ErrorResponse):
         mock_dataset_worker.download_dataset_artifact(default_dataset)
 
-    assert len(responses.calls) == len(BASE_API_CALLS) + 6
+    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/task/{task_id}/artifacts/"),
-        # The API call is retried 5 times
-        ("GET", f"http://testserver/api/v1/task/{task_id}/artifact/dataset_id.tar.zst"),
-        ("GET", f"http://testserver/api/v1/task/{task_id}/artifact/dataset_id.tar.zst"),
-        ("GET", f"http://testserver/api/v1/task/{task_id}/artifact/dataset_id.tar.zst"),
-        ("GET", f"http://testserver/api/v1/task/{task_id}/artifact/dataset_id.tar.zst"),
         ("GET", f"http://testserver/api/v1/task/{task_id}/artifact/dataset_id.tar.zst"),
     ]
 
@@ -284,7 +277,7 @@ def test_list_sets_api_error(responses, mock_dataset_worker):
     responses.add(
         responses.GET,
         f"http://testserver/api/v1/process/{PROCESS_ID}/sets/",
-        status=500,
+        status=418,
     )
 
     with pytest.raises(
@@ -393,20 +386,15 @@ def test_list_sets_retrieve_dataset_api_error(
     responses.add(
         responses.GET,
         f"http://testserver/api/v1/datasets/{default_dataset.id}/",
-        status=500,
+        status=418,
     )
 
     with pytest.raises(ErrorResponse):
         next(mock_dev_dataset_worker.list_sets())
 
-    assert len(responses.calls) == 5
+    assert len(responses.calls) == 1
     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}/"),
+        ("GET", f"http://testserver/api/v1/datasets/{default_dataset.id}/")
     ]
 
 
@@ -494,22 +482,17 @@ def test_run_download_dataset_artifact_api_error(
     responses.add(
         responses.GET,
         f"http://testserver/api/v1/task/{default_dataset.task_id}/artifacts/",
-        status=500,
+        status=418,
     )
 
     with pytest.raises(SystemExit):
         mock_dataset_worker.run()
 
-    assert len(responses.calls) == len(BASE_API_CALLS) * 2 + 5
+    assert len(responses.calls) == len(BASE_API_CALLS) * 2 + 1
     assert [
         (call.request.method, call.request.url) for call in responses.calls
     ] == BASE_API_CALLS * 2 + [
-        # We retry 5 times the API call
-        ("GET", f"http://testserver/api/v1/task/{default_dataset.task_id}/artifacts/"),
-        ("GET", f"http://testserver/api/v1/task/{default_dataset.task_id}/artifacts/"),
-        ("GET", f"http://testserver/api/v1/task/{default_dataset.task_id}/artifacts/"),
-        ("GET", f"http://testserver/api/v1/task/{default_dataset.task_id}/artifacts/"),
-        ("GET", f"http://testserver/api/v1/task/{default_dataset.task_id}/artifacts/"),
+        ("GET", f"http://testserver/api/v1/task/{default_dataset.task_id}/artifacts/")
     ]
 
     assert [(level, message) for _, level, message in caplog.record_tuples] == [
@@ -519,16 +502,9 @@ def test_run_download_dataset_artifact_api_error(
             "Retrieving data for Set (train) from Dataset (dataset_id) (1/1)",
         ),
         (logging.INFO, "Downloading artifact for Dataset (dataset_id)"),
-        *[
-            (
-                logging.INFO,
-                f"Retrying arkindex_worker.worker.base.BaseWorker.request in {retry} seconds as it raised ErrorResponse: .",
-            )
-            for retry in [3.0, 4.0, 8.0, 16.0]
-        ],
         (
             logging.WARNING,
-            "An API error occurred while processing Set (train) from Dataset (dataset_id): 500 Internal Server Error - None",
+            "An API error occurred while processing Set (train) from Dataset (dataset_id): 418 I'm a Teapot - None",
         ),
         (
             logging.ERROR,

tests/test_elements_worker/{test_classifications.py → test_classification.py}

@@ -3,10 +3,12 @@ import re
 from uuid import UUID
 
 import pytest
-from apistar.exceptions import ErrorResponse
 
+from arkindex.exceptions import ErrorResponse
 from arkindex_worker.cache import CachedClassification, CachedElement
 from arkindex_worker.models import Element
+from arkindex_worker.utils import DEFAULT_BATCH_SIZE
+from tests import CORPUS_ID
 
 from . import BASE_API_CALLS
 
@@ -15,11 +17,96 @@ from . import BASE_API_CALLS
 DELETE_PARAMETER = "DELETE_PARAMETER"
 
 
+def test_load_corpus_classes_api_error(responses, mock_elements_worker):
+    responses.add(
+        responses.GET,
+        f"http://testserver/api/v1/corpus/{CORPUS_ID}/classes/",
+        status=418,
+    )
+
+    assert not mock_elements_worker.classes
+    with pytest.raises(
+        Exception, match="Stopping pagination as data will be incomplete"
+    ):
+        mock_elements_worker.load_corpus_classes()
+
+    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 do 5 retries
+        (
+            "GET",
+            f"http://testserver/api/v1/corpus/{CORPUS_ID}/classes/",
+        ),
+        (
+            "GET",
+            f"http://testserver/api/v1/corpus/{CORPUS_ID}/classes/",
+        ),
+        (
+            "GET",
+            f"http://testserver/api/v1/corpus/{CORPUS_ID}/classes/",
+        ),
+        (
+            "GET",
+            f"http://testserver/api/v1/corpus/{CORPUS_ID}/classes/",
+        ),
+        (
+            "GET",
+            f"http://testserver/api/v1/corpus/{CORPUS_ID}/classes/",
+        ),
+    ]
+    assert not mock_elements_worker.classes
+
+
+def test_load_corpus_classes(responses, mock_elements_worker):
+    responses.add(
+        responses.GET,
+        f"http://testserver/api/v1/corpus/{CORPUS_ID}/classes/",
+        status=200,
+        json={
+            "count": 3,
+            "next": None,
+            "results": [
+                {
+                    "id": "0000",
+                    "name": "good",
+                },
+                {
+                    "id": "1111",
+                    "name": "average",
+                },
+                {
+                    "id": "2222",
+                    "name": "bad",
+                },
+            ],
+        },
+    )
+
+    assert not mock_elements_worker.classes
+    mock_elements_worker.load_corpus_classes()
+
+    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}/classes/",
+        ),
+    ]
+    assert mock_elements_worker.classes == {
+        "good": "0000",
+        "average": "1111",
+        "bad": "2222",
+    }
+
+
 def test_get_ml_class_id_load_classes(responses, mock_elements_worker):
-    corpus_id = "11111111-1111-1111-1111-111111111111"
     responses.add(
         responses.GET,
-        f"http://testserver/api/v1/corpus/{corpus_id}/classes/",
+        f"http://testserver/api/v1/corpus/{CORPUS_ID}/classes/",
         status=200,
         json={
             "count": 1,
@@ -42,7 +129,7 @@ def test_get_ml_class_id_load_classes(responses, mock_elements_worker):
     ] == BASE_API_CALLS + [
         (
             "GET",
-            f"http://testserver/api/v1/corpus/{corpus_id}/classes/",
+            f"http://testserver/api/v1/corpus/{CORPUS_ID}/classes/",
         ),
     ]
     assert mock_elements_worker.classes == {"good": "0000"}
@@ -51,12 +138,11 @@ def test_get_ml_class_id_load_classes(responses, mock_elements_worker):
 
 def test_get_ml_class_id_inexistant_class(mock_elements_worker, responses):
     # A missing class is now created automatically
-    corpus_id = "11111111-1111-1111-1111-111111111111"
     mock_elements_worker.classes = {"good": "0000"}
 
     responses.add(
         responses.POST,
-        f"http://testserver/api/v1/corpus/{corpus_id}/classes/",
+        f"http://testserver/api/v1/corpus/{CORPUS_ID}/classes/",
         status=201,
         json={"id": "new-ml-class-1234"},
     )
@@ -82,12 +168,10 @@ def test_get_ml_class_id(mock_elements_worker):
 
 
 def test_get_ml_class_reload(responses, mock_elements_worker):
-    corpus_id = "11111111-1111-1111-1111-111111111111"
-
     # Add some initial classes
     responses.add(
         responses.GET,
-        f"http://testserver/api/v1/corpus/{corpus_id}/classes/",
+        f"http://testserver/api/v1/corpus/{CORPUS_ID}/classes/",
         json={
             "count": 1,
             "next": None,
@@ -103,7 +187,7 @@ def test_get_ml_class_reload(responses, mock_elements_worker):
     # Invalid response when trying to create class2
     responses.add(
         responses.POST,
-        f"http://testserver/api/v1/corpus/{corpus_id}/classes/",
+        f"http://testserver/api/v1/corpus/{CORPUS_ID}/classes/",
         status=400,
         json={"non_field_errors": "Already exists"},
     )
@@ -111,7 +195,7 @@ def test_get_ml_class_reload(responses, mock_elements_worker):
     # Add both classes (class2 is created by another process)
     responses.add(
         responses.GET,
-        f"http://testserver/api/v1/corpus/{corpus_id}/classes/",
+        f"http://testserver/api/v1/corpus/{CORPUS_ID}/classes/",
         json={
             "count": 2,
             "next": None,
@@ -141,15 +225,15 @@ def test_get_ml_class_reload(responses, mock_elements_worker):
     ] == BASE_API_CALLS + [
         (
             "GET",
-            f"http://testserver/api/v1/corpus/{corpus_id}/classes/",
+            f"http://testserver/api/v1/corpus/{CORPUS_ID}/classes/",
         ),
         (
             "POST",
-            f"http://testserver/api/v1/corpus/{corpus_id}/classes/",
+            f"http://testserver/api/v1/corpus/{CORPUS_ID}/classes/",
         ),
         (
             "GET",
-            f"http://testserver/api/v1/corpus/{corpus_id}/classes/",
+            f"http://testserver/api/v1/corpus/{CORPUS_ID}/classes/",
         ),
     ]
 
@@ -169,7 +253,7 @@ def test_retrieve_ml_class_not_in_cache(responses, mock_elements_worker):
     """
     responses.add(
         responses.GET,
-        f"http://testserver/api/v1/corpus/{mock_elements_worker.corpus_id}/classes/",
+        f"http://testserver/api/v1/corpus/{CORPUS_ID}/classes/",
         status=200,
         json={
             "count": 1,
@@ -189,7 +273,7 @@ def test_retrieve_ml_class_not_in_cache(responses, mock_elements_worker):
     ] == BASE_API_CALLS + [
         (
             "GET",
-            f"http://testserver/api/v1/corpus/{mock_elements_worker.corpus_id}/classes/",
+            f"http://testserver/api/v1/corpus/{CORPUS_ID}/classes/",
         ),
     ]
 
@@ -276,7 +360,7 @@ def test_create_classification_api_error(responses, mock_elements_worker):
     responses.add(
         responses.POST,
         "http://testserver/api/v1/classifications/",
-        status=500,
+        status=418,
     )
 
     with pytest.raises(ErrorResponse):
@@ -287,17 +371,10 @@ def test_create_classification_api_error(responses, mock_elements_worker):
             high_confidence=True,
         )
 
-    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
-        ("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/"),
-    ]
+    ] == BASE_API_CALLS + [("POST", "http://testserver/api/v1/classifications/")]
 
 
 def test_create_classification_create_ml_class(mock_elements_worker, responses):
@@ -306,7 +383,7 @@ def test_create_classification_create_ml_class(mock_elements_worker, responses):
     # Automatically create a missing class!
     responses.add(
         responses.POST,
-        "http://testserver/api/v1/corpus/11111111-1111-1111-1111-111111111111/classes/",
+        f"http://testserver/api/v1/corpus/{CORPUS_ID}/classes/",
         status=201,
         json={"id": "new-ml-class-1234"},
     )
@@ -330,7 +407,7 @@ def test_create_classification_create_ml_class(mock_elements_worker, responses):
         for call in responses.calls[-2:]
     ] == [
         (
-            "http://testserver/api/v1/corpus/11111111-1111-1111-1111-111111111111/classes/",
+            f"http://testserver/api/v1/corpus/{CORPUS_ID}/classes/",
             {"name": "a_class"},
         ),
         (
@@ -609,7 +686,7 @@ def test_create_classifications_api_error(responses, mock_elements_worker):
     responses.add(
         responses.POST,
         "http://testserver/api/v1/classification/bulk/",
-        status=500,
+        status=418,
     )
     elt = Element({"id": "12341234-1234-1234-1234-123412341234"})
     classes = [
@@ -630,17 +707,10 @@ def test_create_classifications_api_error(responses, mock_elements_worker):
             element=elt, classifications=classes
         )
 
-    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
-        ("POST", "http://testserver/api/v1/classification/bulk/"),
-        ("POST", "http://testserver/api/v1/classification/bulk/"),
-        ("POST", "http://testserver/api/v1/classification/bulk/"),
-        ("POST", "http://testserver/api/v1/classification/bulk/"),
-        ("POST", "http://testserver/api/v1/classification/bulk/"),
-    ]
+    ] == BASE_API_CALLS + [("POST", "http://testserver/api/v1/classification/bulk/")]
 
 
 def test_create_classifications_create_ml_class(mock_elements_worker, responses):
@@ -649,7 +719,7 @@ def test_create_classifications_create_ml_class(mock_elements_worker, responses)
     # Automatically create a missing class!
     responses.add(
         responses.POST,
-        "http://testserver/api/v1/corpus/11111111-1111-1111-1111-111111111111/classes/",
+        f"http://testserver/api/v1/corpus/{CORPUS_ID}/classes/",
         status=201,
         json={"id": "new-ml-class-1234"},
     )
@@ -690,7 +760,7 @@ def test_create_classifications_create_ml_class(mock_elements_worker, responses)
     ] == BASE_API_CALLS + [
         (
             "POST",
-            "http://testserver/api/v1/corpus/11111111-1111-1111-1111-111111111111/classes/",
+            f"http://testserver/api/v1/corpus/{CORPUS_ID}/classes/",
         ),
         ("POST", "http://testserver/api/v1/classification/bulk/"),
     ]
@@ -709,7 +779,8 @@ def test_create_classifications_create_ml_class(mock_elements_worker, responses)
     }
 
 
-def test_create_classifications(responses, mock_elements_worker):
+@pytest.mark.parametrize("batch_size", [DEFAULT_BATCH_SIZE, 1])
+def test_create_classifications(batch_size, responses, mock_elements_worker):
     mock_elements_worker.classes = {"portrait": "0000", "landscape": "1111"}
     elt = Element({"id": "12341234-1234-1234-1234-123412341234"})
     responses.add(
@@ -733,62 +804,98 @@ def test_create_classifications(responses, mock_elements_worker):
                 "high_confidence": False,
             },
         ],
+        batch_size=batch_size,
     )
 
-    assert len(responses.calls) == len(BASE_API_CALLS) + 1
+    bulk_api_calls = [("POST", "http://testserver/api/v1/classification/bulk/")]
+    if batch_size != DEFAULT_BATCH_SIZE:
+        bulk_api_calls.append(("POST", "http://testserver/api/v1/classification/bulk/"))
+
+    assert len(responses.calls) == len(BASE_API_CALLS) + len(bulk_api_calls)
     assert [
         (call.request.method, call.request.url) for call in responses.calls
-    ] == BASE_API_CALLS + [
-        ("POST", "http://testserver/api/v1/classification/bulk/"),
-    ]
+    ] == BASE_API_CALLS + bulk_api_calls
 
-    assert json.loads(responses.calls[-1].request.body) == {
+    first_cl = {"confidence": 0.75, "high_confidence": False, "ml_class": "0000"}
+    second_cl = {"confidence": 0.25, "high_confidence": False, "ml_class": "1111"}
+    empty_payload = {
         "parent": str(elt.id),
         "worker_run_id": "56785678-5678-5678-5678-567856785678",
-        "classifications": [
-            {
-                "confidence": 0.75,
-                "high_confidence": False,
-                "ml_class": "0000",
-            },
-            {
-                "confidence": 0.25,
-                "high_confidence": False,
-                "ml_class": "1111",
-            },
-        ],
+        "classifications": [],
     }
 
+    bodies = []
+    first_call_idx = None
+    if batch_size > 1:
+        first_call_idx = -1
+        bodies.append({**empty_payload, "classifications": [first_cl, second_cl]})
+    else:
+        first_call_idx = -2
+        bodies.append({**empty_payload, "classifications": [first_cl]})
+        bodies.append({**empty_payload, "classifications": [second_cl]})
+
+    assert [
+        json.loads(bulk_call.request.body)
+        for bulk_call in responses.calls[first_call_idx:]
+    ] == bodies
 
-def test_create_classifications_with_cache(responses, mock_elements_worker_with_cache):
+
+@pytest.mark.parametrize("batch_size", [DEFAULT_BATCH_SIZE, 1])
+def test_create_classifications_with_cache(
+    batch_size, 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")
 
-    responses.add(
-        responses.POST,
-        "http://testserver/api/v1/classification/bulk/",
-        status=200,
-        json={
-            "parent": str(elt.id),
-            "worker_run_id": "56785678-5678-5678-5678-567856785678",
-            "classifications": [
-                {
-                    "id": "00000000-0000-0000-0000-000000000000",
-                    "ml_class": "0000",
-                    "confidence": 0.75,
-                    "high_confidence": False,
-                    "state": "pending",
-                },
-                {
-                    "id": "11111111-1111-1111-1111-111111111111",
-                    "ml_class": "1111",
-                    "confidence": 0.25,
-                    "high_confidence": False,
-                    "state": "pending",
+    if batch_size > 1:
+        responses.add(
+            responses.POST,
+            "http://testserver/api/v1/classification/bulk/",
+            status=200,
+            json={
+                "parent": str(elt.id),
+                "worker_run_id": "56785678-5678-5678-5678-567856785678",
+                "classifications": [
+                    {
+                        "id": "00000000-0000-0000-0000-000000000000",
+                        "ml_class": "0000",
+                        "confidence": 0.75,
+                        "high_confidence": False,
+                        "state": "pending",
+                    },
+                    {
+                        "id": "11111111-1111-1111-1111-111111111111",
+                        "ml_class": "1111",
+                        "confidence": 0.25,
+                        "high_confidence": False,
+                        "state": "pending",
+                    },
+                ],
+            },
+        )
+    else:
+        for cl_id, cl_class, cl_conf in [
+            ("00000000-0000-0000-0000-000000000000", "0000", 0.75),
+            ("11111111-1111-1111-1111-111111111111", "1111", 0.25),
+        ]:
+            responses.add(
+                responses.POST,
+                "http://testserver/api/v1/classification/bulk/",
+                status=200,
+                json={
+                    "parent": str(elt.id),
+                    "worker_run_id": "56785678-5678-5678-5678-567856785678",
+                    "classifications": [
+                        {
+                            "id": cl_id,
+                            "ml_class": cl_class,
+                            "confidence": cl_conf,
+                            "high_confidence": False,
+                            "state": "pending",
+                        },
+                    ],
                 },
-            ],
-        },
-    )
+            )
 
     mock_elements_worker_with_cache.create_classifications(
         element=elt,
@@ -804,32 +911,41 @@ def test_create_classifications_with_cache(responses, mock_elements_worker_with_
                 "high_confidence": False,
             },
         ],
+        batch_size=batch_size,
     )
 
-    assert len(responses.calls) == len(BASE_API_CALLS) + 1
+    bulk_api_calls = [("POST", "http://testserver/api/v1/classification/bulk/")]
+    if batch_size != DEFAULT_BATCH_SIZE:
+        bulk_api_calls.append(("POST", "http://testserver/api/v1/classification/bulk/"))
+
+    assert len(responses.calls) == len(BASE_API_CALLS) + len(bulk_api_calls)
     assert [
         (call.request.method, call.request.url) for call in responses.calls
-    ] == BASE_API_CALLS + [
-        ("POST", "http://testserver/api/v1/classification/bulk/"),
-    ]
+    ] == BASE_API_CALLS + bulk_api_calls
 
-    assert json.loads(responses.calls[-1].request.body) == {
+    first_cl = {"confidence": 0.75, "high_confidence": False, "ml_class": "0000"}
+    second_cl = {"confidence": 0.25, "high_confidence": False, "ml_class": "1111"}
+    empty_payload = {
         "parent": str(elt.id),
         "worker_run_id": "56785678-5678-5678-5678-567856785678",
-        "classifications": [
-            {
-                "confidence": 0.75,
-                "high_confidence": False,
-                "ml_class": "0000",
-            },
-            {
-                "confidence": 0.25,
-                "high_confidence": False,
-                "ml_class": "1111",
-            },
-        ],
+        "classifications": [],
     }
 
+    bodies = []
+    first_call_idx = None
+    if batch_size > 1:
+        first_call_idx = -1
+        bodies.append({**empty_payload, "classifications": [first_cl, second_cl]})
+    else:
+        first_call_idx = -2
+        bodies.append({**empty_payload, "classifications": [first_cl]})
+        bodies.append({**empty_payload, "classifications": [second_cl]})
+
+    assert [
+        json.loads(bulk_call.request.body)
+        for bulk_call in responses.calls[first_call_idx:]
+    ] == bodies
+
     # Check that created classifications were properly stored in SQLite cache
     assert list(CachedClassification.select()) == [
         CachedClassification(