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

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,
     )