Fix tests that were flaky due to label creation

Author: Kevin Kim

tests/integration/conftest.py

@@ -119,6 +119,49 @@ def get_invites(client):
     return invites
 
 
+def wait_for_data_row_processing(client, data_row):
+    """
+    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
+    """
+    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)
+
+
+def wait_for_label_processing(project):
+    """
+    Do not use. Only for testing.
+
+    Returns Label after waiting for it to finish processing.
+    If `project.labels()` is called before label is fully processed,
+    it may return an empty set
+    """
+    timeout_seconds = 10
+    while True:
+        label = project.labels().get_one()
+        if label is not None:
+            return label
+        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)
+
+
 @pytest.fixture
 def queries():
     return SimpleNamespace(cancel_invite=cancel_invite,
@@ -207,7 +250,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()
 
@@ -348,8 +391,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)
+
     yield [project, dataset, datarow, label]
 
     for label in project.labels():

tests/integration/test_data_row_media_attributes.py

@@ -1,10 +1,10 @@
 from time import sleep
+from conftest import wait_for_data_row_processing
 
 
-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):
+    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

tests/integration/test_data_row_metadata.py

@@ -4,6 +4,7 @@
 import pytest
 import uuid
 
+from conftest import wait_for_data_row_processing
 from labelbox import DataRow, Dataset
 from labelbox.schema.data_row_metadata import DataRowMetadataField, DataRowMetadata, DataRowMetadataKind, DeleteDataRowMetadata, \
     DataRowMetadataOntology, _parse_metadata_schema
@@ -90,10 +91,9 @@ 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):
+    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 == []

tests/integration/test_export.py

@@ -1,13 +1,16 @@
 from time import sleep
 import uuid
 
+from conftest import wait_for_data_row_processing
 from labelbox.data.annotation_types.annotation import ObjectAnnotation
 from labelbox.schema.annotation_import import LabelImport
 
 
 def test_export_annotations_nested_checklist(
         client, configured_project_with_complex_ontology):
     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 +50,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)