diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index fd2900061..497a894dc 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -207,7 +207,7 @@ def datarow(dataset, image_url): }, ]) task.wait_till_done() - dr = next(dataset.data_rows()) + dr = dataset.data_rows().get_one() yield dr dr.delete() @@ -304,7 +304,7 @@ def configured_project(project, client, rand_gen, image_url): @pytest.fixture def configured_project_with_label(client, rand_gen, image_url, project, dataset, - datarow): + datarow, wait_for_label_processing): """Project with a connected dataset, having one datarow Project contains an ontology with 1 bbox tool Additionally includes a create_label method for any needed extra labels @@ -348,8 +348,8 @@ def create_label(): project.create_label = create_label project.create_label() - label = project.labels().get_one() - assert label is not None, "Cannot fetch created label" + label = wait_for_label_processing(project)[0] + yield [project, dataset, datarow, label] for label in project.labels(): @@ -407,3 +407,56 @@ def configured_project_with_complex_ontology(client, rand_gen, image_url): yield [project, data_row] dataset.delete() project.delete() + + +@pytest.fixture +def wait_for_data_row_processing(): + """ + Do not use. Only for testing. + + Returns DataRow after waiting for it to finish processing media_attributes. + Some tests, specifically ones that rely on label export, rely on + DataRow be fully processed with media_attributes + """ + + def func(client, data_row): + data_row_id = data_row.uid + timeout_seconds = 60 + while True: + data_row = client.get_data_row(data_row_id) + if data_row.media_attributes: + return data_row + timeout_seconds -= 2 + if timeout_seconds <= 0: + raise TimeoutError( + f"Timed out waiting for DataRow '{data_row_id}' to finish processing media_attributes" + ) + time.sleep(2) + + return func + + +@pytest.fixture +def wait_for_label_processing(): + """ + Do not use. Only for testing. + + Returns project's labels as a list after waiting for them to finish processing. + If `project.labels()` is called before label is fully processed, + it may return an empty set + """ + + def func(project): + timeout_seconds = 10 + while True: + labels = list(project.labels()) + if len(labels) > 0: + return labels + timeout_seconds -= 2 + if timeout_seconds <= 0: + raise TimeoutError( + f"Timed out waiting for label for project '{project.uid}' to finish processing" + ) + time.sleep(2) + + return func diff --git a/tests/integration/test_data_row_media_attributes.py b/tests/integration/test_data_row_media_attributes.py index 3b78744c4..4e75513d4 100644 --- a/tests/integration/test_data_row_media_attributes.py +++ b/tests/integration/test_data_row_media_attributes.py @@ -1,10 +1,7 @@ -from time import sleep - - -def test_export_empty_media_attributes(configured_project_with_label): - project, _, _, _ = configured_project_with_label - # Wait for exporter to retrieve latest labels - sleep(10) +def test_export_empty_media_attributes(client, configured_project_with_label, + wait_for_data_row_processing): + project, _, data_row, _ = configured_project_with_label + data_row = wait_for_data_row_processing(client, data_row) labels = list(project.label_generator()) assert len( labels diff --git a/tests/integration/test_data_row_metadata.py b/tests/integration/test_data_row_metadata.py index 232bc718f..26eb0447d 100644 --- a/tests/integration/test_data_row_metadata.py +++ b/tests/integration/test_data_row_metadata.py @@ -1,5 +1,4 @@ from datetime import datetime -from time import sleep import pytest import uuid @@ -90,10 +89,10 @@ def make_named_metadata(dr_id) -> DataRowMetadata: return metadata -def test_export_empty_metadata(configured_project_with_label): - project, _, _, _ = configured_project_with_label - # Wait for exporter to retrieve latest labels - sleep(10) +def test_export_empty_metadata(client, configured_project_with_label, + wait_for_data_row_processing): + project, _, data_row, _ = configured_project_with_label + data_row = wait_for_data_row_processing(client, data_row) labels = project.label_generator() label = next(labels) assert label.data.metadata == [] diff --git a/tests/integration/test_export.py b/tests/integration/test_export.py index d7f0689e1..9e35436a2 100644 --- a/tests/integration/test_export.py +++ b/tests/integration/test_export.py @@ -1,4 +1,3 @@ -from time import sleep import uuid from labelbox.data.annotation_types.annotation import ObjectAnnotation @@ -6,8 +5,11 @@ def test_export_annotations_nested_checklist( - client, configured_project_with_complex_ontology): + client, configured_project_with_complex_ontology, + wait_for_data_row_processing): project, data_row = configured_project_with_complex_ontology + data_row = wait_for_data_row_processing(client, data_row) + ontology = project.ontology().normalized tool = ontology["tools"][0] @@ -47,9 +49,7 @@ def test_export_annotations_nested_checklist( task = LabelImport.create_from_objects(client, project.uid, f'label-import-{uuid.uuid4()}', data) task.wait_until_done() - # Wait for exporter to retrieve latest labels - sleep(10) - labels = project.label_generator().as_list() + labels = project.label_generator() object_annotation = [ annot for annot in next(labels).annotations if isinstance(annot, ObjectAnnotation) diff --git a/tests/integration/test_project_setup.py b/tests/integration/test_project_setup.py index d55a1731c..8d81ba043 100644 --- a/tests/integration/test_project_setup.py +++ b/tests/integration/test_project_setup.py @@ -1,7 +1,6 @@ from datetime import datetime, timedelta, timezone import json import time -import time import pytest