arkindex-base-worker 0.4.0__py3-none-any.whl → 0.4.0a1__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",
@@ -179,12 +179,44 @@ def test_validate_model_version_not_created(mock_training_worker):
         mock_training_worker.validate_model_version(hash="a", size=1, archive_hash="b")
 
 
+@pytest.mark.parametrize("status_code", [403, 500])
+def test_validate_model_version_catch_errors(
+    mocker, mock_training_worker, caplog, status_code
+):
+    mocker.patch(
+        "arkindex.client.ArkindexClient.request.retry.retry", return_value=False
+    )
+
+    mock_training_worker.model_version = {"id": "model_version_id"}
+    args = {
+        "hash": "hash",
+        "archive_hash": "archive_hash",
+        "size": 30,
+    }
+    mock_training_worker.api_client.add_error_response(
+        "ValidateModelVersion",
+        id="model_version_id",
+        status_code=status_code,
+        body=args,
+    )
+
+    mock_training_worker.validate_model_version(**args)
+    assert mock_training_worker.model_version == {"id": "model_version_id"}
+    assert [
+        (level, message)
+        for module, level, message in caplog.record_tuples
+        if module == "arkindex_worker"
+    ] == [
+        (
+            logging.WARNING,
+            "An error occurred while validating model version model_version_id, please check its status.",
+        ),
+    ]
+
+
 @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 = {
@@ -193,11 +225,11 @@ def test_validate_model_version_hash_conflict(
         "size": 30,
     }
     mock_training_worker.api_client.add_error_response(
-        "PartialUpdateModelVersion",
+        "ValidateModelVersion",
         id="another_id",
         status_code=409,
-        body={"state": "available", **args},
-        content={"id": ["model_version_id"]},
+        body=args,
+        content=default_model_version,
     )
     if deletion_failed:
         mock_training_worker.api_client.add_error_response(
@@ -212,11 +244,6 @@ 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
@@ -239,7 +266,6 @@ 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."),
     ]
 
@@ -252,9 +278,9 @@ def test_validate_model_version(mock_training_worker, default_model_version, cap
         "size": 30,
     }
     mock_training_worker.api_client.add_response(
-        "PartialUpdateModelVersion",
+        "ValidateModelVersion",
         id="model_version_id",
-        body={"state": "available", **args},
+        body=args,
         response=default_model_version,
     )