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

tests/test_elements_worker/test_task.py

@@ -1,10 +1,10 @@
 import uuid
 
 import pytest
-from apistar.exceptions import ErrorResponse
 
+from arkindex.exceptions import ErrorResponse
 from arkindex_worker.models import Artifact
-from tests.conftest import FIXTURES_DIR
+from tests import FIXTURES_DIR
 from tests.test_elements_worker import BASE_API_CALLS
 
 TASK_ID = uuid.UUID("cafecafe-cafe-cafe-cafe-cafecafecafe")
@@ -33,22 +33,17 @@ def test_list_artifacts_api_error(responses, mock_dataset_worker):
     responses.add(
         responses.GET,
         f"http://testserver/api/v1/task/{TASK_ID}/artifacts/",
-        status=500,
+        status=418,
     )
 
     with pytest.raises(ErrorResponse):
         mock_dataset_worker.list_artifacts(task_id=TASK_ID)
 
-    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/")
     ]
 
 
@@ -153,7 +148,7 @@ def test_download_artifact_api_error(responses, mock_dataset_worker, default_art
     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):
@@ -161,16 +156,11 @@ def test_download_artifact_api_error(responses, mock_dataset_worker, default_art
             task_id=TASK_ID, artifact=default_artifact
         )
 
-    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}/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"),
+        ("GET", f"http://testserver/api/v1/task/{TASK_ID}/artifact/dataset_id.tar.zst")
     ]
 
 

tests/test_elements_worker/test_training.py

@@ -8,7 +8,7 @@ from arkindex_worker.worker import BaseWorker
 from arkindex_worker.worker.training import TrainingMixin, create_archive
 
 
-@pytest.fixture()
+@pytest.fixture
 def mock_training_worker(monkeypatch):
     class TrainingWorker(BaseWorker, TrainingMixin):
         """
@@ -22,7 +22,7 @@ def mock_training_worker(monkeypatch):
     return training_worker
 
 
-@pytest.fixture()
+@pytest.fixture
 def default_model_version():
     return {
         "id": "model_version_id",
@@ -181,7 +181,10 @@ def test_validate_model_version_not_created(mock_training_worker):
 
 @pytest.mark.parametrize("deletion_failed", [True, False])
 def test_validate_model_version_hash_conflict(
-    mock_training_worker, default_model_version, caplog, deletion_failed
+    mock_training_worker,
+    default_model_version,
+    caplog,
+    deletion_failed,
 ):
     mock_training_worker.model_version = {"id": "another_id"}
     args = {
@@ -190,11 +193,11 @@ def test_validate_model_version_hash_conflict(
         "size": 30,
     }
     mock_training_worker.api_client.add_error_response(
-        "ValidateModelVersion",
+        "PartialUpdateModelVersion",
         id="another_id",
         status_code=409,
-        body=args,
-        content=default_model_version,
+        body={"state": "available", **args},
+        content={"id": ["model_version_id"]},
     )
     if deletion_failed:
         mock_training_worker.api_client.add_error_response(
@@ -209,6 +212,11 @@ def test_validate_model_version_hash_conflict(
             id="another_id",
             response="No content",
         )
+    mock_training_worker.api_client.add_response(
+        "RetrieveModelVersion",
+        id="model_version_id",
+        response=default_model_version,
+    )
 
     mock_training_worker.validate_model_version(**args)
     assert mock_training_worker.model_version == default_model_version
@@ -231,6 +239,7 @@ def test_validate_model_version_hash_conflict(
         ),
         (logging.WARNING, "Removing the pending model version."),
         *error_msg,
+        (logging.INFO, "Retrieving the existing model version."),
         (logging.INFO, "Model version model_version_id is now available."),
     ]
 
@@ -243,9 +252,9 @@ def test_validate_model_version(mock_training_worker, default_model_version, cap
         "size": 30,
     }
     mock_training_worker.api_client.add_response(
-        "ValidateModelVersion",
+        "PartialUpdateModelVersion",
         id="model_version_id",
-        body=args,
+        body={"state": "available", **args},
         response=default_model_version,
     )