PR: added extra test to verify project labels in a model run

Author: Val Brodsky

labelbox/schema/model_run.py

@@ -55,9 +55,11 @@ def upsert_labels(self,
         Args:
             label_ids (list): label ids to insert
             project_id (string): project uuid, all project labels will be uploaded
+                Either label_ids OR project_id is required but NOT both
             timeout_seconds (float): Max waiting time, in seconds.
         Returns:
             ID of newly generated async task
+
         """
 
         use_label_ids = label_ids is not None and len(label_ids) > 0
@@ -71,13 +73,13 @@ def upsert_labels(self,
             raise ValueError("Must only one of label ids, project id")
 
         if use_label_ids:
-            return self._upsert_labels_by_label_ids(label_ids)
+            return self._upsert_labels_by_label_ids(label_ids, timeout_seconds)
         else:  # use_project_id
-            return self._upsert_labels_by_project_id(project_id)
+            return self._upsert_labels_by_project_id(project_id,
+                                                     timeout_seconds)
 
-    def _upsert_labels_by_label_ids(self,
-                                    label_ids: List[str],
-                                    timeout_seconds=3600):
+    def _upsert_labels_by_label_ids(self, label_ids: List[str],
+                                    timeout_seconds: int):
         mutation_name = 'createMEAModelRunLabelRegistrationTask'
         create_task_query_str = """mutation createMEAModelRunLabelRegistrationTaskPyApi($modelRunId: ID!, $labelIds : [ID!]!) {
         %s(where : { id : $modelRunId}, data : {labelIds: $labelIds})}
@@ -99,9 +101,8 @@ def _upsert_labels_by_label_ids(self,
             }})['MEALabelRegistrationTaskStatus'],
                                      timeout_seconds=timeout_seconds)
 
-    def _upsert_labels_by_project_id(self,
-                                     project_id: str,
-                                     timeout_seconds=3600):
+    def _upsert_labels_by_project_id(self, project_id: str,
+                                     timeout_seconds: int):
         mutation_name = 'createMEAModelRunProjectLabelRegistrationTask'
         create_task_query_str = """mutation createMEAModelRunProjectLabelRegistrationTaskPyApi($modelRunId: ID!, $projectId : ID!) {
         %s(where : { modelRunId : $modelRunId, projectId: $projectId})}

tests/integration/annotation_import/conftest.py

@@ -441,9 +441,8 @@ def model_run_with_model_run_data_rows(client, configured_project,
 
 
 @pytest.fixture
-def model_run_with_model_run_all_project_data_rows(client, configured_project,
-                                                   model_run_predictions,
-                                                   model_run):
+def model_run_with_all_project_labels(client, configured_project,
+                                      model_run_predictions, model_run):
     configured_project.enable_model_assisted_labeling()
 
     upload_task = LabelImport.create_from_objects(

tests/integration/annotation_import/test_mea_prediction_import.py

@@ -39,20 +39,43 @@ def test_create_from_objects(model_run_with_model_run_data_rows,
 
 
 def test_create_from_objects_all_project_labels(
-        model_run_with_model_run_all_project_data_rows, object_predictions,
+        model_run_with_all_project_labels, object_predictions,
         annotation_import_test_helpers):
     name = str(uuid.uuid4())
 
-    annotation_import = model_run_with_model_run_all_project_data_rows.add_predictions(
+    annotation_import = model_run_with_all_project_labels.add_predictions(
         name=name, predictions=object_predictions)
 
-    assert annotation_import.model_run_id == model_run_with_model_run_all_project_data_rows.uid
+    assert annotation_import.model_run_id == model_run_with_all_project_labels.uid
     annotation_import_test_helpers.check_running_state(annotation_import, name)
     annotation_import_test_helpers.assert_file_content(
         annotation_import.input_file_url, object_predictions)
     annotation_import.wait_until_done()
 
 
+def test_model_run_project_labels(model_run_with_all_project_labels,
+                                  model_run_predictions):
+    model_run = model_run_with_all_project_labels
+    model_run_exported_labels = model_run.export_labels(download=True)
+    labels_indexed_by_schema_id = {}
+    for label in model_run_exported_labels:
+        # assuming exported array of label 'objects' has only one label per data row... as usually is when there are no label revisions
+        schema_id = label['Label']['objects'][0]['schemaId']
+        labels_indexed_by_schema_id[schema_id] = label
+
+    assert (len(
+        labels_indexed_by_schema_id.keys())) == len(model_run_predictions)
+
+    # making sure the labels are in this model run are all labels uploaded to the project
+    # by comparing some 'immutable' attributes
+    for expected_label in model_run_predictions:
+        schema_id = expected_label['schemaId']
+        actual_label = labels_indexed_by_schema_id[schema_id]
+        assert actual_label['Label']['objects'][0]['title'] == expected_label[
+            'name']
+        assert actual_label['DataRow ID'] == expected_label['dataRow']['id']
+
+
 def test_create_from_label_objects(model_run_with_model_run_data_rows,
                                    object_predictions,
                                    annotation_import_test_helpers):