diff --git a/arkindex_worker/worker/classification.py b/arkindex_worker/worker/classification.py index 3e979c043cceac8d847fd99d291797ee8e74e5cb..8bc3b7ba4ff9d811e6a707e084cfa42be7a35338 100644 --- a/arkindex_worker/worker/classification.py +++ b/arkindex_worker/worker/classification.py @@ -4,6 +4,7 @@ import os from apistar.exceptions import ErrorResponse from arkindex_worker import logger +from arkindex_worker.cache import CachedElement from arkindex_worker.models import Element @@ -65,8 +66,8 @@ class ClassificationMixin(object): Create a classification on the given element through API """ assert element and isinstance( - element, Element - ), "element shouldn't be null and should be of type Element" + element, (Element, CachedElement) + ), "element shouldn't be null and should be an Element or CachedElement" assert ml_class and isinstance( ml_class, str ), "ml_class shouldn't be null and should be of type str" diff --git a/tests/test_elements_worker/test_classifications.py b/tests/test_elements_worker/test_classifications.py index 20a677071451b76cc329816fd91787f7bcd571c0..4ea4eee87ff82282393d32a659390b91d0e7b33b 100644 --- a/tests/test_elements_worker/test_classifications.py +++ b/tests/test_elements_worker/test_classifications.py @@ -159,7 +159,10 @@ def test_create_classification_wrong_element(mock_elements_worker): confidence=0.42, high_confidence=True, ) - assert str(e.value) == "element shouldn't be null and should be of type Element" + assert ( + str(e.value) + == "element shouldn't be null and should be an Element or CachedElement" + ) with pytest.raises(AssertionError) as e: mock_elements_worker.create_classification( @@ -168,7 +171,10 @@ def test_create_classification_wrong_element(mock_elements_worker): confidence=0.42, high_confidence=True, ) - assert str(e.value) == "element shouldn't be null and should be of type Element" + assert ( + str(e.value) + == "element shouldn't be null and should be an Element or CachedElement" + ) def test_create_classification_wrong_ml_class(mock_elements_worker, responses):