import logging
import pytest
from tests import check_db, check_json
from worker_init_elements.worker import INIT_PAGE_SIZE
@pytest.mark.parametrize("use_cache", [True, False])
def test_run_process(use_cache, mock_worker):
mock_worker.use_cache = use_cache
mock_worker.api_client.add_response(
"RetrieveCorpus",
id=mock_worker.process_information["corpus"],
response={
"id": mock_worker.process_information["corpus"],
"types": [{"id": "A", "slug": "class"}, {"id": "B", "slug": "student"}],
},
)
mock_worker.api_client.add_response(
"ListProcessElements",
id=mock_worker.process_information["id"],
with_image=mock_worker.use_cache,
allow_missing_data=True,
page_size=INIT_PAGE_SIZE,
response=[
{
"id": "11111111-1111-1111-1111-111111111111",
"type_id": "A",
"name": "Class 1",
"confidence": None,
**(
{
"image_id": "cafecafe-cafe-cafe-cafe-cafecafecafe",
"image_width": 42,
"image_height": 1337,
"image_url": "http://cafe",
"polygon": [[0, 0], [0, 1], [1, 1], [1, 0], [0, 0]],
"rotation_angle": 42,
"mirrored": False,
}
if mock_worker.use_cache
else {}
),
},
{
"id": "22222222-2222-2222-2222-222222222222",
"type_id": "A",
"name": "Class 2",
"confidence": None,
**(
{
"image_id": "beefbeef-beef-beef-beef-beefbeefbeef",
"image_width": 42,
"image_height": 1337,
"image_url": "http://beef",
"polygon": [[0, 0], [0, 2], [2, 2], [2, 0], [0, 0]],
"rotation_angle": 0,
"mirrored": True,
}
if mock_worker.use_cache
else {}
),
},
{
"id": "33333333-3333-3333-3333-333333333333",
"type_id": "A",
"name": "Class 3",
"confidence": 0.42,
**(
{
"rotation_angle": 17,
"mirrored": True,
"image_id": None,
"image_width": None,
"image_height": None,
"image_url": None,
"polygon": None,
}
if mock_worker.use_cache
else {}
),
},
],
)
mock_worker.process()
check_json(
json_path=mock_worker.work_dir / "elements.json",
elements=[
{"id": "11111111-1111-1111-1111-111111111111", "type": "class"},
{"id": "22222222-2222-2222-2222-222222222222", "type": "class"},
{"id": "33333333-3333-3333-3333-333333333333", "type": "class"},
],
)
db_path = mock_worker.work_dir / "db.sqlite"
assert db_path.is_file() is use_cache
if use_cache:
check_db(
db_path=db_path,
elements=[
{
"id": "11111111111111111111111111111111",
"image_id": "cafecafecafecafecafecafecafecafe",
"initial": 1,
"parent_id": None,
"polygon": "[[0, 0], [0, 1], [1, 1], [1, 0], [0, 0]]",
"type": "class",
"worker_version_id": None,
"worker_run_id": None,
"rotation_angle": 42,
"mirrored": False,
"confidence": None,
},
{
"id": "22222222222222222222222222222222",
"image_id": "beefbeefbeefbeefbeefbeefbeefbeef",
"initial": 1,
"parent_id": None,
"polygon": "[[0, 0], [0, 2], [2, 2], [2, 0], [0, 0]]",
"type": "class",
"worker_version_id": None,
"worker_run_id": None,
"rotation_angle": 0,
"mirrored": True,
"confidence": None,
},
{
"id": "33333333333333333333333333333333",
"image_id": None,
"initial": 1,
"parent_id": None,
"polygon": None,
"type": "class",
"worker_version_id": None,
"worker_run_id": None,
"rotation_angle": 17,
"mirrored": True,
"confidence": 0.42,
},
],
images=[
{
"id": "beefbeefbeefbeefbeefbeefbeefbeef",
"width": 42,
"height": 1337,
"url": "http://beef",
},
{
"id": "cafecafecafecafecafecafecafecafe",
"width": 42,
"height": 1337,
"url": "http://cafe",
},
],
)
def test_run_distributed(mock_worker):
mock_worker.use_cache = True
mock_worker.chunks_number = 4
mock_worker.api_client.add_response(
"RetrieveCorpus",
id=mock_worker.process_information["corpus"],
response={
"id": mock_worker.process_information["corpus"],
"types": [{"id": "A", "slug": "class"}, {"id": "B", "slug": "student"}],
},
)
mock_worker.api_client.add_response(
"ListProcessElements",
id=mock_worker.process_information["id"],
with_image=True,
allow_missing_data=True,
page_size=INIT_PAGE_SIZE,
response=[
{
"id": "22222222-2222-2222-2222-222222222222",
"type_id": "A",
"name": "Class 2",
"image_id": None,
"image_width": None,
"image_height": None,
"image_url": None,
"polygon": None,
"rotation_angle": 0,
"mirrored": False,
"confidence": None,
},
{
"id": "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa",
"type_id": "B",
"name": "Student 1",
"image_id": "cafecafe-cafe-cafe-cafe-cafecafecafe",
"image_width": 42,
"image_height": 1337,
"image_url": "http://cafe",
"polygon": [[0, 0], [0, 1], [1, 1], [1, 0], [0, 0]],
"rotation_angle": 0,
"mirrored": False,
"confidence": None,
},
{
"id": "bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb",
"type_id": "B",
"name": "Student 2",
"image_id": "beefbeef-beef-beef-beef-beefbeefbeef",
"image_width": 42,
"image_height": 1337,
"image_url": "http://beef",
"polygon": [[0, 0], [0, 2], [2, 2], [2, 0], [0, 0]],
"rotation_angle": 0,
"mirrored": False,
"confidence": None,
},
{
"id": "cccccccc-cccc-cccc-cccc-cccccccccccc",
"type_id": "B",
"name": "Student 3",
"image_id": "beefbeef-beef-beef-beef-beefbeefbeef",
"image_width": 42,
"image_height": 1337,
"image_url": "http://beef",
"polygon": [[0, 0], [0, 2], [2, 2], [2, 0], [0, 0]],
"rotation_angle": 0,
"mirrored": False,
"confidence": None,
},
{
"id": "dddddddd-dddd-dddd-dddd-dddddddddddd",
"type_id": "B",
"name": "Student 4",
"image_id": "cafecafe-cafe-cafe-cafe-cafecafecafe",
"image_width": 42,
"image_height": 1337,
"image_url": "http://cafe",
"polygon": [[0, 0], [0, 1], [1, 1], [1, 0], [0, 0]],
"rotation_angle": 0,
"mirrored": False,
"confidence": None,
},
{
"id": "eeeeeeee-eeee-eeee-eeee-eeeeeeeeeeee",
"type_id": "B",
"name": "Student 5",
"image_id": None,
"image_width": None,
"image_height": None,
"image_url": None,
"polygon": None,
"rotation_angle": 0,
"mirrored": False,
"confidence": None,
},
{
"id": "ffffffff-ffff-ffff-ffff-ffffffffffff",
"type_id": "B",
"name": "Student 6",
"image_id": None,
"image_width": None,
"image_height": None,
"image_url": None,
"polygon": None,
"rotation_angle": 0,
"mirrored": False,
"confidence": None,
},
],
)
mock_worker.process()
check_json(
json_path=mock_worker.work_dir / "elements_chunk_1.json",
elements=[
{"id": "22222222-2222-2222-2222-222222222222", "type": "class"},
{"id": "dddddddd-dddd-dddd-dddd-dddddddddddd", "type": "student"},
],
)
check_json(
json_path=mock_worker.work_dir / "elements_chunk_2.json",
elements=[
{"id": "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa", "type": "student"},
{"id": "eeeeeeee-eeee-eeee-eeee-eeeeeeeeeeee", "type": "student"},
],
)
check_json(
json_path=mock_worker.work_dir / "elements_chunk_3.json",
elements=[
{"id": "bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb", "type": "student"},
{"id": "ffffffff-ffff-ffff-ffff-ffffffffffff", "type": "student"},
],
)
check_json(
json_path=mock_worker.work_dir / "elements_chunk_4.json",
elements=[{"id": "cccccccc-cccc-cccc-cccc-cccccccccccc", "type": "student"}],
)
check_db(
db_path=mock_worker.work_dir / "db_1.sqlite",
elements=[
{
"id": "22222222222222222222222222222222",
"image_id": None,
"initial": 1,
"parent_id": None,
"polygon": None,
"type": "class",
"worker_version_id": None,
"worker_run_id": None,
"rotation_angle": 0,
"mirrored": False,
"confidence": None,
},
{
"id": "dddddddddddddddddddddddddddddddd",
"image_id": "cafecafecafecafecafecafecafecafe",
"initial": 1,
"parent_id": None,
"polygon": "[[0, 0], [0, 1], [1, 1], [1, 0], [0, 0]]",
"type": "student",
"worker_version_id": None,
"worker_run_id": None,
"rotation_angle": 0,
"mirrored": False,
"confidence": None,
},
],
images=[
{
"id": "cafecafecafecafecafecafecafecafe",
"width": 42,
"height": 1337,
"url": "http://cafe",
},
],
)
check_db(
db_path=mock_worker.work_dir / "db_2.sqlite",
elements=[
{
"id": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
"image_id": "cafecafecafecafecafecafecafecafe",
"initial": 1,
"parent_id": None,
"polygon": "[[0, 0], [0, 1], [1, 1], [1, 0], [0, 0]]",
"type": "student",
"worker_version_id": None,
"worker_run_id": None,
"rotation_angle": 0,
"mirrored": False,
"confidence": None,
},
{
"id": "eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee",
"image_id": None,
"initial": 1,
"parent_id": None,
"polygon": None,
"type": "student",
"worker_version_id": None,
"worker_run_id": None,
"rotation_angle": 0,
"mirrored": False,
"confidence": None,
},
],
images=[
{
"id": "cafecafecafecafecafecafecafecafe",
"width": 42,
"height": 1337,
"url": "http://cafe",
},
],
)
check_db(
db_path=mock_worker.work_dir / "db_3.sqlite",
elements=[
{
"id": "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb",
"image_id": "beefbeefbeefbeefbeefbeefbeefbeef",
"initial": 1,
"parent_id": None,
"polygon": "[[0, 0], [0, 2], [2, 2], [2, 0], [0, 0]]",
"type": "student",
"worker_version_id": None,
"worker_run_id": None,
"rotation_angle": 0,
"mirrored": False,
"confidence": None,
},
{
"id": "ffffffffffffffffffffffffffffffff",
"image_id": None,
"initial": 1,
"parent_id": None,
"polygon": None,
"type": "student",
"worker_version_id": None,
"worker_run_id": None,
"rotation_angle": 0,
"mirrored": False,
"confidence": None,
},
],
images=[
{
"id": "beefbeefbeefbeefbeefbeefbeefbeef",
"width": 42,
"height": 1337,
"url": "http://beef",
},
],
)
check_db(
db_path=mock_worker.work_dir / "db_4.sqlite",
elements=[
{
"id": "cccccccccccccccccccccccccccccccc",
"image_id": "beefbeefbeefbeefbeefbeefbeefbeef",
"initial": 1,
"parent_id": None,
"polygon": "[[0, 0], [0, 2], [2, 2], [2, 0], [0, 0]]",
"type": "student",
"worker_version_id": None,
"worker_run_id": None,
"rotation_angle": 0,
"mirrored": False,
"confidence": None,
},
],
images=[
{
"id": "beefbeefbeefbeefbeefbeefbeefbeef",
"width": 42,
"height": 1337,
"url": "http://beef",
},
],
)
def test_not_enough_elements(mock_worker):
mock_worker.chunks_number = 5
mock_worker.api_client.add_response(
"RetrieveCorpus",
id=mock_worker.process_information["corpus"],
response={
"id": mock_worker.process_information["corpus"],
"types": [{"id": "A", "slug": "class"}, {"id": "B", "slug": "student"}],
},
)
mock_worker.api_client.add_response(
"ListProcessElements",
id=mock_worker.process_information["id"],
with_image=False,
allow_missing_data=True,
page_size=INIT_PAGE_SIZE,
response=[
{
"id": "22222222-2222-2222-2222-222222222222",
"type_id": "A",
"name": "Class 2",
"confidence": 1.0,
}
],
)
with pytest.raises(AssertionError, match="Too few elements have been retrieved"):
mock_worker.process()
def test_run_duplicates(mock_worker, caplog):
caplog.set_level(logging.WARNING)
mock_worker.use_cache = True
mock_worker.api_client.add_response(
"RetrieveCorpus",
id=mock_worker.process_information["corpus"],
response={
"id": mock_worker.process_information["corpus"],
"types": [{"id": "A", "slug": "class"}, {"id": "B", "slug": "student"}],
},
)
mock_worker.api_client.add_response(
"ListProcessElements",
id=mock_worker.process_information["id"],
with_image=True,
allow_missing_data=True,
page_size=INIT_PAGE_SIZE,
response=[
{
"id": "11111111-1111-1111-1111-111111111111",
"type_id": "A",
"name": "Class 1",
"image_id": "cafecafe-cafe-cafe-cafe-cafecafecafe",
"image_width": 42,
"image_height": 1337,
"image_url": "http://cafe",
"polygon": [[0, 0], [0, 1], [1, 1], [1, 0], [0, 0]],
"rotation_angle": 0,
"mirrored": False,
"confidence": None,
},
{
"id": "22222222-2222-2222-2222-222222222222",
"type_id": "A",
"name": "Class 2",
"image_id": "cafecafe-cafe-cafe-cafe-cafecafecafe",
"image_width": 42,
"image_height": 1337,
"image_url": "http://cafe",
"polygon": [[0, 0], [0, 1], [1, 1], [1, 0], [0, 0]],
"rotation_angle": 0,
"mirrored": False,
"confidence": None,
},
{
"id": "22222222-2222-2222-2222-222222222222",
"type_id": "A",
"name": "Class 2",
"image_id": "cafecafe-cafe-cafe-cafe-cafecafecafe",
"image_width": 42,
"image_height": 1337,
"image_url": "http://cafe",
"polygon": [[0, 0], [0, 1], [1, 1], [1, 0], [0, 0]],
"rotation_angle": 0,
"mirrored": False,
"confidence": None,
},
{
"id": "33333333-3333-3333-3333-333333333333",
"type_id": "A",
"name": "Class 3",
"image_id": None,
"image_width": None,
"image_height": None,
"image_url": None,
"polygon": None,
"rotation_angle": 0,
"mirrored": False,
"confidence": None,
},
],
)
mock_worker.process()
check_json(
json_path=mock_worker.work_dir / "elements.json",
elements=[
{"id": "11111111-1111-1111-1111-111111111111", "type": "class"},
{"id": "22222222-2222-2222-2222-222222222222", "type": "class"},
{"id": "33333333-3333-3333-3333-333333333333", "type": "class"},
],
)
check_db(
db_path=mock_worker.work_dir / "db.sqlite",
elements=[
{
"id": "11111111111111111111111111111111",
"image_id": "cafecafecafecafecafecafecafecafe",
"initial": 1,
"parent_id": None,
"polygon": "[[0, 0], [0, 1], [1, 1], [1, 0], [0, 0]]",
"type": "class",
"worker_version_id": None,
"worker_run_id": None,
"rotation_angle": 0,
"mirrored": False,
"confidence": None,
},
{
"id": "22222222222222222222222222222222",
"image_id": "cafecafecafecafecafecafecafecafe",
"initial": 1,
"parent_id": None,
"polygon": "[[0, 0], [0, 1], [1, 1], [1, 0], [0, 0]]",
"type": "class",
"worker_version_id": None,
"worker_run_id": None,
"rotation_angle": 0,
"mirrored": False,
"confidence": None,
},
{
"id": "33333333333333333333333333333333",
"image_id": None,
"initial": 1,
"parent_id": None,
"polygon": None,
"type": "class",
"worker_version_id": None,
"worker_run_id": None,
"rotation_angle": 0,
"mirrored": False,
"confidence": None,
},
],
images=[
{
"id": "cafecafecafecafecafecafecafecafe",
"width": 42,
"height": 1337,
"url": "http://cafe",
},
],
)
assert [(record.levelname, record.message) for record in caplog.records] == [
("WARNING", "1 duplicate elements have been ignored.")
]
def test_run_empty(mock_worker, caplog):
caplog.set_level(logging.WARNING)
mock_worker.api_client.add_response(
"RetrieveCorpus",
id=mock_worker.process_information["corpus"],
response={
"id": mock_worker.process_information["corpus"],
"types": [{"id": "A", "slug": "class"}, {"id": "B", "slug": "student"}],
},
)
mock_worker.api_client.add_response(
"ListProcessElements",
id=mock_worker.process_information["id"],
with_image=False,
allow_missing_data=True,
page_size=INIT_PAGE_SIZE,
response=[],
)
with pytest.raises(SystemExit) as ctx:
mock_worker.process()
assert ctx.value.code == 1
assert list(mock_worker.work_dir.rglob("*")) == []
assert [(record.levelname, record.message) for record in caplog.records] == [
("ERROR", "No elements found, aborting workflow."),
]