Skip to content
Snippets Groups Projects
Verified Commit dcd12d6c authored by Yoann Schneider's avatar Yoann Schneider :tennis:
Browse files

assert ml_class_id is a valid uuid

parent 10527c74
No related branches found
No related tags found
1 merge request!278Use MLClass when using the CreateClassifications helper
Pipeline #79930 passed
Stage: test
Stage: build
Stage: release
This commit is part of merge request !278. Comments created here will be created in the context of that merge request.
Hide whitespace changes
Inline Side-by-side
arkindex_worker/worker/classification.py
+ 9
− 0
View file @ dcd12d6c
...@@ -4,6 +4,7 @@ ElementsWorker methods for classifications and ML classes. ...@@ -4,6 +4,7 @@ ElementsWorker methods for classifications and ML classes.
""" """
from typing import Dict, List, Optional, Union from typing import Dict, List, Optional, Union
from uuid import UUID
from apistar.exceptions import ErrorResponse from apistar.exceptions import ErrorResponse
from peewee import IntegrityError from peewee import IntegrityError
...@@ -209,6 +210,14 @@ class ClassificationMixin(object): ...@@ -209,6 +210,14 @@ class ClassificationMixin(object):
ml_class_id, str ml_class_id, str
), f"Classification at index {index} in classifications: ml_class_id shouldn't be null and should be of type str" ), f"Classification at index {index} in classifications: ml_class_id shouldn't be null and should be of type str"
# Make sure it's a valid UUID
try:
UUID(ml_class_id)
except ValueError:
raise ValueError(
f"Classification at index {index} in classifications: ml_class_id is not a valid uuid."
)
confidence = classification.get("confidence") confidence = classification.get("confidence")
assert ( assert (
confidence is not None confidence is not None
......
tests/test_elements_worker/test_classifications.py
+ 54
− 29
View file @ dcd12d6c
# -*- coding: utf-8 -*- # -*- coding: utf-8 -*-
import json import json
from uuid import UUID from uuid import UUID, uuid4
import pytest import pytest
from apistar.exceptions import ErrorResponse from apistar.exceptions import ErrorResponse
...@@ -624,7 +624,7 @@ def test_create_classifications_wrong_classifications(mock_elements_worker): ...@@ -624,7 +624,7 @@ def test_create_classifications_wrong_classifications(mock_elements_worker):
element=elt, element=elt,
classifications=[ classifications=[
{ {
"ml_class_id": "uuid1", "ml_class_id": str(uuid4()),
"confidence": 0.75, "confidence": 0.75,
"high_confidence": False, "high_confidence": False,
}, },
...@@ -644,7 +644,7 @@ def test_create_classifications_wrong_classifications(mock_elements_worker): ...@@ -644,7 +644,7 @@ def test_create_classifications_wrong_classifications(mock_elements_worker):
element=elt, element=elt,
classifications=[ classifications=[
{ {
"ml_class_id": "uuid1", "ml_class_id": str(uuid4()),
"confidence": 0.75, "confidence": 0.75,
"high_confidence": False, "high_confidence": False,
}, },
...@@ -665,7 +665,7 @@ def test_create_classifications_wrong_classifications(mock_elements_worker): ...@@ -665,7 +665,7 @@ def test_create_classifications_wrong_classifications(mock_elements_worker):
element=elt, element=elt,
classifications=[ classifications=[
{ {
"ml_class_id": "uuid1", "ml_class_id": str(uuid4()),
"confidence": 0.75, "confidence": 0.75,
"high_confidence": False, "high_confidence": False,
}, },
...@@ -681,17 +681,38 @@ def test_create_classifications_wrong_classifications(mock_elements_worker): ...@@ -681,17 +681,38 @@ def test_create_classifications_wrong_classifications(mock_elements_worker):
== "Classification at index 1 in classifications: ml_class_id shouldn't be null and should be of type str" == "Classification at index 1 in classifications: ml_class_id shouldn't be null and should be of type str"
) )
with pytest.raises(ValueError) as e:
mock_elements_worker.create_classifications(
element=elt,
classifications=[
{
"ml_class_id": str(uuid4()),
"confidence": 0.75,
"high_confidence": False,
},
{
"ml_class_id": "not_an_uuid",
"confidence": 0.25,
"high_confidence": False,
},
],
)
assert (
str(e.value)
== "Classification at index 1 in classifications: ml_class_id is not a valid uuid."
)
with pytest.raises(AssertionError) as e: with pytest.raises(AssertionError) as e:
mock_elements_worker.create_classifications( mock_elements_worker.create_classifications(
element=elt, element=elt,
classifications=[ classifications=[
{ {
"ml_class_id": "uuid1", "ml_class_id": str(uuid4()),
"confidence": 0.75, "confidence": 0.75,
"high_confidence": False, "high_confidence": False,
}, },
{ {
"ml_class_id": "uuid2", "ml_class_id": str(uuid4()),
"high_confidence": False, "high_confidence": False,
}, },
], ],
...@@ -706,12 +727,12 @@ def test_create_classifications_wrong_classifications(mock_elements_worker): ...@@ -706,12 +727,12 @@ def test_create_classifications_wrong_classifications(mock_elements_worker):
element=elt, element=elt,
classifications=[ classifications=[
{ {
"ml_class_id": "uuid1", "ml_class_id": str(uuid4()),
"confidence": 0.75, "confidence": 0.75,
"high_confidence": False, "high_confidence": False,
}, },
{ {
"ml_class_id": "uuid2", "ml_class_id": str(uuid4()),
"confidence": None, "confidence": None,
"high_confidence": False, "high_confidence": False,
}, },
...@@ -727,12 +748,12 @@ def test_create_classifications_wrong_classifications(mock_elements_worker): ...@@ -727,12 +748,12 @@ def test_create_classifications_wrong_classifications(mock_elements_worker):
element=elt, element=elt,
classifications=[ classifications=[
{ {
"ml_class_id": "uuid1", "ml_class_id": str(uuid4()),
"confidence": 0.75, "confidence": 0.75,
"high_confidence": False, "high_confidence": False,
}, },
{ {
"ml_class_id": "uuid2", "ml_class_id": str(uuid4()),
"confidence": "wrong confidence", "confidence": "wrong confidence",
"high_confidence": False, "high_confidence": False,
}, },
...@@ -748,12 +769,12 @@ def test_create_classifications_wrong_classifications(mock_elements_worker): ...@@ -748,12 +769,12 @@ def test_create_classifications_wrong_classifications(mock_elements_worker):
element=elt, element=elt,
classifications=[ classifications=[
{ {
"ml_class_id": "uuid1", "ml_class_id": str(uuid4()),
"confidence": 0.75, "confidence": 0.75,
"high_confidence": False, "high_confidence": False,
}, },
{ {
"ml_class_id": "uuid2", "ml_class_id": str(uuid4()),
"confidence": 0, "confidence": 0,
"high_confidence": False, "high_confidence": False,
}, },
...@@ -769,12 +790,12 @@ def test_create_classifications_wrong_classifications(mock_elements_worker): ...@@ -769,12 +790,12 @@ def test_create_classifications_wrong_classifications(mock_elements_worker):
element=elt, element=elt,
classifications=[ classifications=[
{ {
"ml_class_id": "uuid1", "ml_class_id": str(uuid4()),
"confidence": 0.75, "confidence": 0.75,
"high_confidence": False, "high_confidence": False,
}, },
{ {
"ml_class_id": "uuid2", "ml_class_id": str(uuid4()),
"confidence": 2.00, "confidence": 2.00,
"high_confidence": False, "high_confidence": False,
}, },
...@@ -790,12 +811,12 @@ def test_create_classifications_wrong_classifications(mock_elements_worker): ...@@ -790,12 +811,12 @@ def test_create_classifications_wrong_classifications(mock_elements_worker):
element=elt, element=elt,
classifications=[ classifications=[
{ {
"ml_class_id": "uuid1", "ml_class_id": str(uuid4()),
"confidence": 0.75, "confidence": 0.75,
"high_confidence": False, "high_confidence": False,
}, },
{ {
"ml_class_id": "uuid2", "ml_class_id": str(uuid4()),
"confidence": 0.25, "confidence": 0.25,
"high_confidence": "wrong high_confidence", "high_confidence": "wrong high_confidence",
}, },
...@@ -816,12 +837,12 @@ def test_create_classifications_api_error(responses, mock_elements_worker): ...@@ -816,12 +837,12 @@ def test_create_classifications_api_error(responses, mock_elements_worker):
elt = Element({"id": "12341234-1234-1234-1234-123412341234"}) elt = Element({"id": "12341234-1234-1234-1234-123412341234"})
classes = [ classes = [
{ {
"ml_class_id": "uuid1", "ml_class_id": str(uuid4()),
"confidence": 0.75, "confidence": 0.75,
"high_confidence": False, "high_confidence": False,
}, },
{ {
"ml_class_id": "uuid2", "ml_class_id": str(uuid4()),
"confidence": 0.25, "confidence": 0.25,
"high_confidence": False, "high_confidence": False,
}, },
...@@ -847,19 +868,21 @@ def test_create_classifications_api_error(responses, mock_elements_worker): ...@@ -847,19 +868,21 @@ def test_create_classifications_api_error(responses, mock_elements_worker):
def test_create_classifications(responses, mock_elements_worker_with_cache): def test_create_classifications(responses, mock_elements_worker_with_cache):
# Set MLClass in cache # Set MLClass in cache
portrait_uuid = str(uuid4())
landscape_uuid = str(uuid4())
mock_elements_worker_with_cache.classes[ mock_elements_worker_with_cache.classes[
mock_elements_worker_with_cache.corpus_id mock_elements_worker_with_cache.corpus_id
] = {"portrait": "uuid1", "landscape": "uuid2"} ] = {"portrait": portrait_uuid, "landscape": landscape_uuid}
elt = CachedElement.create(id="12341234-1234-1234-1234-123412341234", type="thing") elt = CachedElement.create(id="12341234-1234-1234-1234-123412341234", type="thing")
classes = [ classes = [
{ {
"ml_class_id": "uuid1", "ml_class_id": portrait_uuid,
"confidence": 0.75, "confidence": 0.75,
"high_confidence": False, "high_confidence": False,
}, },
{ {
"ml_class_id": "uuid2", "ml_class_id": landscape_uuid,
"confidence": 0.25, "confidence": 0.25,
"high_confidence": False, "high_confidence": False,
}, },
...@@ -875,14 +898,14 @@ def test_create_classifications(responses, mock_elements_worker_with_cache): ...@@ -875,14 +898,14 @@ def test_create_classifications(responses, mock_elements_worker_with_cache):
"classifications": [ "classifications": [
{ {
"id": "00000000-0000-0000-0000-000000000000", "id": "00000000-0000-0000-0000-000000000000",
"ml_class": "uuid1", "ml_class": portrait_uuid,
"confidence": 0.75, "confidence": 0.75,
"high_confidence": False, "high_confidence": False,
"state": "pending", "state": "pending",
}, },
{ {
"id": "11111111-1111-1111-1111-111111111111", "id": "11111111-1111-1111-1111-111111111111",
"ml_class": "uuid2", "ml_class": landscape_uuid,
"confidence": 0.25, "confidence": 0.25,
"high_confidence": False, "high_confidence": False,
"state": "pending", "state": "pending",
...@@ -936,15 +959,17 @@ def test_create_classifications_not_in_cache( ...@@ -936,15 +959,17 @@ def test_create_classifications_not_in_cache(
CreateClassifications using ID that are not in `.classes` attribute. CreateClassifications using ID that are not in `.classes` attribute.
Will load corpus MLClass to insert the corresponding name in Cache. Will load corpus MLClass to insert the corresponding name in Cache.
""" """
portrait_uuid = str(uuid4())
landscape_uuid = str(uuid4())
elt = CachedElement.create(id="12341234-1234-1234-1234-123412341234", type="thing") elt = CachedElement.create(id="12341234-1234-1234-1234-123412341234", type="thing")
classes = [ classes = [
{ {
"ml_class_id": "uuid1", "ml_class_id": portrait_uuid,
"confidence": 0.75, "confidence": 0.75,
"high_confidence": False, "high_confidence": False,
}, },
{ {
"ml_class_id": "uuid2", "ml_class_id": landscape_uuid,
"confidence": 0.25, "confidence": 0.25,
"high_confidence": False, "high_confidence": False,
}, },
...@@ -960,14 +985,14 @@ def test_create_classifications_not_in_cache( ...@@ -960,14 +985,14 @@ def test_create_classifications_not_in_cache(
"classifications": [ "classifications": [
{ {
"id": "00000000-0000-0000-0000-000000000000", "id": "00000000-0000-0000-0000-000000000000",
"ml_class": "uuid1", "ml_class": portrait_uuid,
"confidence": 0.75, "confidence": 0.75,
"high_confidence": False, "high_confidence": False,
"state": "pending", "state": "pending",
}, },
{ {
"id": "11111111-1111-1111-1111-111111111111", "id": "11111111-1111-1111-1111-111111111111",
"ml_class": "uuid2", "ml_class": landscape_uuid,
"confidence": 0.25, "confidence": 0.25,
"high_confidence": False, "high_confidence": False,
"state": "pending", "state": "pending",
...@@ -984,10 +1009,10 @@ def test_create_classifications_not_in_cache( ...@@ -984,10 +1009,10 @@ def test_create_classifications_not_in_cache(
"next": None, "next": None,
"results": [ "results": [
{ {
"id": "uuid1", "id": portrait_uuid,
"name": "portrait", "name": "portrait",
}, },
{"id": "uuid2", "name": "landscape"}, {"id": landscape_uuid, "name": "landscape"},
], ],
}, },
) )
......
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment