[AL-0] Fix tests that were flaky due to label creation

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

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

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 == []

tests/integration/test_export.py

@@ -1,13 +1,15 @@
-from time import sleep
 import uuid
 
 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):
+        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)

tests/integration/test_project_setup.py

@@ -1,7 +1,6 @@
 from datetime import datetime, timedelta, timezone
 import json
 import time
-import time
 
 import pytest