[PLT-1492] Fix 'flaky' tests spotted during a prod run (#1846)

Author: vbrodsky

libs/labelbox/tests/conftest.py

@@ -1248,6 +1248,21 @@ def teardown_ontology_feature_schemas(ontology: Ontology):
 class ModuleTearDownHelpers(TearDownHelpers): ...
 
 
+class LabelHelpers:
+    def wait_for_labels(self, project, number_of_labels=1):
+        timeout_seconds = 10
+        while True:
+            labels = list(project.labels())
+            if len(labels) >= number_of_labels:
+                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)
+
+
 @pytest.fixture
 def teardown_helpers():
     return TearDownHelpers()
@@ -1256,3 +1271,8 @@ def teardown_helpers():
 @pytest.fixture(scope="module")
 def module_teardown_helpers():
     return TearDownHelpers()
+
+
+@pytest.fixture
+def label_helpers():
+    return LabelHelpers()

libs/labelbox/tests/data/annotation_import/test_generic_data_types.py

@@ -40,11 +40,6 @@ def validate_iso_format(date_string: str):
         (MediaType.Conversational, MediaType.Conversational),
         (MediaType.Document, MediaType.Document),
         (MediaType.Dicom, MediaType.Dicom),
-        (
-            MediaType.LLMPromptResponseCreation,
-            MediaType.LLMPromptResponseCreation,
-        ),
-        (MediaType.LLMPromptCreation, MediaType.LLMPromptCreation),
         (OntologyKind.ResponseCreation, OntologyKind.ResponseCreation),
         (OntologyKind.ModelEvaluation, OntologyKind.ModelEvaluation),
     ],
@@ -114,6 +109,53 @@ def test_import_media_types(
         assert exported_annotations == expected_data
 
 
+@pytest.mark.parametrize(
+    "configured_project, media_type",
+    [
+        (
+            MediaType.LLMPromptResponseCreation,
+            MediaType.LLMPromptResponseCreation,
+        ),
+        (MediaType.LLMPromptCreation, MediaType.LLMPromptCreation),
+    ],
+    indirect=["configured_project"],
+)
+def test_import_media_types_llm(
+    client: Client,
+    configured_project: Project,
+    annotations_by_media_type,
+    exports_v2_by_media_type,
+    export_v2_test_helpers,
+    helpers,
+    media_type,
+    wait_for_label_processing,
+):
+    annotations_ndjson = list(
+        itertools.chain.from_iterable(annotations_by_media_type[media_type])
+    )
+
+    label_import = lb.LabelImport.create_from_objects(
+        client,
+        configured_project.uid,
+        f"test-import-{media_type}",
+        annotations_ndjson,
+    )
+    label_import.wait_until_done()
+
+    assert label_import.state == AnnotationImportState.FINISHED
+    assert len(label_import.errors) == 0
+
+    all_annotations = sorted([a["uuid"] for a in annotations_ndjson])
+    successful_annotations = sorted(
+        [
+            status["uuid"]
+            for status in label_import.statuses
+            if status["status"] == "SUCCESS"
+        ]
+    )
+    assert successful_annotations == all_annotations
+
+
 @pytest.mark.parametrize(
     "configured_project_by_global_key, media_type",
     [

libs/labelbox/tests/integration/schema/test_user_group.py

@@ -147,35 +147,19 @@ def test_cannot_update_group_id(user_group):
 def test_get_user_groups_with_creation_deletion(client):
     user_group = None
     try:
-        # Get all user groups
-        user_groups = list(UserGroup(client).get_user_groups())
-
         # manual delete for iterators
         group_name = data.name()
         user_group = UserGroup(client)
         user_group.name = group_name
         user_group.create()
 
         user_groups_post_creation = list(UserGroup(client).get_user_groups())
+        assert user_group in user_groups_post_creation
 
-        # Verify that at least one user group is returned
-        assert len(user_groups_post_creation) > 0
-        assert len(user_groups_post_creation) == len(user_groups) + 1
-
-        # Verify that each user group has a valid ID and name
-        for ug in user_groups_post_creation:
-            assert ug.id is not None
-            assert ug.name is not None
-
-        user_group.delete()
         user_group = None
 
         user_groups_post_deletion = list(UserGroup(client).get_user_groups())
-
-        assert (
-            len(user_groups_post_deletion) == len(user_groups_post_creation) - 1
-        )
-
+        assert user_group not in user_groups_post_deletion
     finally:
         if user_group:
             user_group.delete()

libs/labelbox/tests/integration/test_label.py

@@ -41,11 +41,13 @@ def test_label_update(configured_project_with_label):
     assert label.label == "something else"
 
 
-def test_label_filter_order(configured_project_with_label):
+def test_label_filter_order(configured_project_with_label, label_helpers):
     project, _, _, label = configured_project_with_label
 
     l1 = label
     project.create_label()
+    label_helpers.wait_for_labels(project, 2)
+
     l2 = next(project.labels())
 
     assert set(project.labels()) == {l1, l2}