Fix flaky annotation test by waiting on prompt response data rows to be created

Author: Val Brodsky

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

@@ -686,14 +686,21 @@ def _create_prompt_response_project(client: Client, rand_gen, media_type, normal
     prompt_response_project.connect_ontology(ontology)
 
     # We have to export to get data row ids
-    result = export_v2_test_helpers.run_project_export_v2_task(prompt_response_project)
-    
-    data_row_ids = [dr["data_row"]["id"] for dr in result]
-    global_keys = [dr["data_row"]["global_key"] for dr in result]
-    
+    data_row_ids = []
+    global_keys = []
+    timeout = 0
+    while len(data_row_ids) < 1 and timeout < 3:
+        result = export_v2_test_helpers.run_project_export_v2_task(prompt_response_project)
+
+        data_row_ids.extend([dr["data_row"]["id"] for dr in result])
+        global_keys.extend([dr["data_row"]["global_key"] for dr in result])
+
+        time.sleep(1)
+        timeout += 1
+
     prompt_response_project.data_row_ids = data_row_ids
     prompt_response_project.global_keys = global_keys
-    
+
     return prompt_response_project, ontology
 
 def _create_project(client: Client, rand_gen, data_row_json_by_media_type, media_type, normalized_ontology_by_media_type) -> Tuple[Project, Ontology, Dataset]:

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

@@ -66,8 +66,8 @@ def test_generic_data_row_type_by_data_row_id(
         (MediaType.Video, GenericDataRowData),
         (MediaType.Conversational, GenericDataRowData),
         (MediaType.Document, GenericDataRowData),
-        # (MediaType.LLMPromptResponseCreation, GenericDataRowData),
-        # (MediaType.LLMPromptCreation, GenericDataRowData),
+        (MediaType.LLMPromptResponseCreation, GenericDataRowData),
+        (MediaType.LLMPromptCreation, GenericDataRowData),
         (OntologyKind.ResponseCreation, GenericDataRowData)
     ],
 )
@@ -100,8 +100,8 @@ def test_generic_data_row_type_by_global_key(
         (MediaType.Conversational, MediaType.Conversational),
         (MediaType.Document, MediaType.Document),
         (MediaType.Dicom, MediaType.Dicom),
-        # (MediaType.LLMPromptResponseCreation, MediaType.LLMPromptResponseCreation),
-        # (MediaType.LLMPromptCreation, MediaType.LLMPromptCreation),
+        (MediaType.LLMPromptResponseCreation, MediaType.LLMPromptResponseCreation),
+        (MediaType.LLMPromptCreation, MediaType.LLMPromptCreation),
         (OntologyKind.ResponseCreation, OntologyKind.ResponseCreation)
     ],
     indirect=["configured_project"]
@@ -220,8 +220,8 @@ def test_import_media_types_by_global_key(
         (MediaType.Conversational, MediaType.Conversational),
         (MediaType.Document, MediaType.Document),
         (MediaType.Dicom, MediaType.Dicom),
-        # (MediaType.LLMPromptResponseCreation, MediaType.LLMPromptResponseCreation),
-        # (MediaType.LLMPromptCreation, MediaType.LLMPromptCreation),
+        (MediaType.LLMPromptResponseCreation, MediaType.LLMPromptResponseCreation),
+        (MediaType.LLMPromptCreation, MediaType.LLMPromptCreation),
         (OntologyKind.ResponseCreation, OntologyKind.ResponseCreation)
     ],
     indirect=["configured_project"]