Skip to content

Commit 3d8b596

Browse files
author
Kevin Kim
committed
Fix tests that were flaky due to label creation
1 parent 0ea6ef9 commit 3d8b596

File tree

4 files changed

+58
-14
lines changed

4 files changed

+58
-14
lines changed

tests/integration/conftest.py

Lines changed: 46 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,49 @@ def get_invites(client):
119119
return invites
120120

121121

122+
def wait_for_data_row_processing(client, data_row):
123+
"""
124+
Do not use. Only for testing.
125+
126+
Returns DataRow after waiting for it to finish processing media_attributes.
127+
Some tests, specifically ones that rely on label export, rely on
128+
DataRow be fully processed with media_attributes
129+
"""
130+
data_row_id = data_row.uid
131+
timeout_seconds = 60
132+
while True:
133+
data_row = client.get_data_row(data_row_id)
134+
if data_row.media_attributes:
135+
return data_row
136+
timeout_seconds -= 2
137+
if timeout_seconds <= 0:
138+
raise TimeoutError(
139+
f"Timed out waiting for DataRow '{data_row_id}' to finish processing media_attributes"
140+
)
141+
time.sleep(2)
142+
143+
144+
def wait_for_label_processing(project):
145+
"""
146+
Do not use. Only for testing.
147+
148+
Returns Label after waiting for it to finish processing.
149+
If `project.labels()` is called before label is fully processed,
150+
it may return an empty set
151+
"""
152+
timeout_seconds = 10
153+
while True:
154+
label = project.labels().get_one()
155+
if label is not None:
156+
return label
157+
timeout_seconds -= 2
158+
if timeout_seconds <= 0:
159+
raise TimeoutError(
160+
f"Timed out waiting for label for project '{project.uid}' to finish processing"
161+
)
162+
time.sleep(2)
163+
164+
122165
@pytest.fixture
123166
def queries():
124167
return SimpleNamespace(cancel_invite=cancel_invite,
@@ -207,7 +250,7 @@ def datarow(dataset, image_url):
207250
},
208251
])
209252
task.wait_till_done()
210-
dr = next(dataset.data_rows())
253+
dr = dataset.data_rows().get_one()
211254
yield dr
212255
dr.delete()
213256

@@ -348,8 +391,8 @@ def create_label():
348391

349392
project.create_label = create_label
350393
project.create_label()
351-
label = project.labels().get_one()
352-
assert label is not None, "Cannot fetch created label"
394+
label = wait_for_label_processing(project)
395+
353396
yield [project, dataset, datarow, label]
354397

355398
for label in project.labels():

tests/integration/test_data_row_media_attributes.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
11
from time import sleep
2+
from conftest import wait_for_data_row_processing
23

34

4-
def test_export_empty_media_attributes(configured_project_with_label):
5-
project, _, _, _ = configured_project_with_label
6-
# Wait for exporter to retrieve latest labels
7-
sleep(10)
5+
def test_export_empty_media_attributes(client, configured_project_with_label):
6+
project, _, data_row, _ = configured_project_with_label
7+
data_row = wait_for_data_row_processing(client, data_row)
88
labels = list(project.label_generator())
99
assert len(
1010
labels

tests/integration/test_data_row_metadata.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import pytest
55
import uuid
66

7+
from conftest import wait_for_data_row_processing
78
from labelbox import DataRow, Dataset
89
from labelbox.schema.data_row_metadata import DataRowMetadataField, DataRowMetadata, DataRowMetadataKind, DeleteDataRowMetadata, \
910
DataRowMetadataOntology, _parse_metadata_schema
@@ -90,10 +91,9 @@ def make_named_metadata(dr_id) -> DataRowMetadata:
9091
return metadata
9192

9293

93-
def test_export_empty_metadata(configured_project_with_label):
94-
project, _, _, _ = configured_project_with_label
95-
# Wait for exporter to retrieve latest labels
96-
sleep(10)
94+
def test_export_empty_metadata(client, configured_project_with_label):
95+
project, _, data_row, _ = configured_project_with_label
96+
data_row = wait_for_data_row_processing(client, data_row)
9797
labels = project.label_generator()
9898
label = next(labels)
9999
assert label.data.metadata == []

tests/integration/test_export.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,16 @@
11
from time import sleep
22
import uuid
33

4+
from conftest import wait_for_data_row_processing
45
from labelbox.data.annotation_types.annotation import ObjectAnnotation
56
from labelbox.schema.annotation_import import LabelImport
67

78

89
def test_export_annotations_nested_checklist(
910
client, configured_project_with_complex_ontology):
1011
project, data_row = configured_project_with_complex_ontology
12+
data_row = wait_for_data_row_processing(client, data_row)
13+
1114
ontology = project.ontology().normalized
1215

1316
tool = ontology["tools"][0]
@@ -47,9 +50,7 @@ def test_export_annotations_nested_checklist(
4750
task = LabelImport.create_from_objects(client, project.uid,
4851
f'label-import-{uuid.uuid4()}', data)
4952
task.wait_until_done()
50-
# Wait for exporter to retrieve latest labels
51-
sleep(10)
52-
labels = project.label_generator().as_list()
53+
labels = project.label_generator()
5354
object_annotation = [
5455
annot for annot in next(labels).annotations
5556
if isinstance(annot, ObjectAnnotation)

0 commit comments

Comments
 (0)