Test&Mutation Fixes

Author: mnoszczak

labelbox/schema/annotation_import.py

@@ -473,14 +473,14 @@ def _get_url_mutation(cls) -> str:
 
     @classmethod
     def _get_model_run_data_rows_mutation(cls) -> str:
-        return """mutation createMalPredictionImportForModelRunDataRows($projectId : ID!, $dataRowIds: [ID!]!, $name: String!, $modelRunId: ID!) {
+        return """mutation createMalPredictionImportForModelRunDataRowsPyApi($dataRowIds: [ID!]!, $name: String!, $modelRunId: ID!, $projectId:ID!) {
             createMalPredictionImportForModelRunDataRows(data: {
                 name: $name
                 modelRunId: $modelRunId
                 dataRowIds: $dataRowIds
                 projectId: $projectId
-            }) {%s}
-        }""" % query.results_query_part(cls)
+            }) {id importType inputFileUrl errorFileUrl project { id name } name statusFileUrl state progress}
+        }"""
 
     @classmethod
     def _get_file_mutation(cls) -> str:

labelbox/schema/model_run.py

@@ -144,17 +144,22 @@ def upsert_predictions_and_send_to_project(
         kwargs = dict(client=self.client, model_run_id=self.uid, name=name)
         project = self.client.get_project(project_id)
         import_job = self.add_predictions(name, predictions)
-        prediction_statuses = import_job.statuses()
+        prediction_statuses = import_job.statuses
         mea_to_mal_data_rows_set = set([
             row['dataRow']['id']
             for row in prediction_statuses
             if row['status'] == 'SUCCESS'
         ])
         mea_to_mal_data_rows = list(
             mea_to_mal_data_rows_set)[:DATAROWS_IMPORT_LIMIT]
-        logger.warning(
-            f"Got {len(mea_to_mal_data_rows_set)} data rows to import, trimmed down to {DATAROWS_IMPORT_LIMIT} data rows"
-        )
+
+        if len(mea_to_mal_data_rows) >= DATAROWS_IMPORT_LIMIT:
+
+            logger.warning(
+                f"Got {len(mea_to_mal_data_rows_set)} data rows to import, trimmed down to {DATAROWS_IMPORT_LIMIT} data rows"
+            )
+        if len(mea_to_mal_data_rows) == 0:
+            return import_job, None, None
 
         batch = project.create_batch(name, mea_to_mal_data_rows, priority)
         mal_prediction_import = Entity.MALPredictionImport.create_for_model_run_data_rows(

tests/integration/annotation_import/conftest.py

@@ -121,6 +121,20 @@ def configured_project(client, ontology, rand_gen, image_url):
     dataset.delete()
 
 
+@pytest.fixture
+def configured_project_without_data_rows(client, ontology, rand_gen):
+    project = client.create_project(name=rand_gen(str))
+    dataset = client.create_dataset(name=rand_gen(str))
+    editor = list(
+        client.get_labeling_frontends(
+            where=LabelingFrontend.name == "editor"))[0]
+    project.setup(editor, ontology)
+    project.update(queue_mode=project.QueueMode.Batch)
+    yield project
+    project.delete()
+    dataset.delete()
+
+
 @pytest.fixture
 def prediction_id_mapping(configured_project):
     #Maps tool types to feature schema ids

tests/integration/annotation_import/test_upsert_prediction_import.py

@@ -0,0 +1,121 @@
+import uuid
+import ndjson
+import pytest
+
+from labelbox.schema.annotation_import import AnnotationImportState, MEAPredictionImport
+"""
+- Here we only want to check that the uploads are calling the validation
+- Then with unit tests we can check the types of errors raised
+
+"""
+
+
+def test_create_from_url(model_run_with_model_run_data_rows,
+                         configured_project_without_data_rows,
+                         annotation_import_test_helpers):
+    name = str(uuid.uuid4())
+    url = "https://storage.googleapis.com/labelbox-public-bucket/predictions_test_v2.ndjson"
+
+    annotation_import, batch, mal_prediction_import = model_run_with_model_run_data_rows.upsert_predictions_and_send_to_project(
+        name=name,
+        predictions=url,
+        project_id=configured_project_without_data_rows.uid,
+        priority=5)
+
+    assert annotation_import.model_run_id == model_run_with_model_run_data_rows.uid
+    annotation_import_test_helpers.check_running_state(annotation_import, name,
+                                                       url)
+    annotation_import.wait_until_done()
+
+    if batch:
+        assert batch.project().uid == configured_project_without_data_rows.uid
+    if mal_prediction_import:
+        mal_prediction_import.wait_until_done()
+
+
+def test_create_from_objects(model_run_with_model_run_data_rows,
+                             configured_project_without_data_rows,
+                             object_predictions,
+                             annotation_import_test_helpers):
+    name = str(uuid.uuid4())
+
+    annotation_import, batch, mal_prediction_import = model_run_with_model_run_data_rows.upsert_predictions_and_send_to_project(
+        name=name,
+        predictions=object_predictions,
+        project_id=configured_project_without_data_rows.uid,
+        priority=5)
+
+    assert annotation_import.model_run_id == model_run_with_model_run_data_rows.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()
+
+    if batch:
+        assert batch.project().uid == configured_project_without_data_rows.uid
+
+    if mal_prediction_import:
+        annotation_import_test_helpers.check_running_state(
+            mal_prediction_import, name)
+        mal_prediction_import.wait_until_done()
+
+
+def test_create_from_local_file(tmp_path, model_run_with_model_run_data_rows,
+                                configured_project_without_data_rows,
+                                object_predictions,
+                                annotation_import_test_helpers):
+    name = str(uuid.uuid4())
+    file_name = f"{name}.ndjson"
+    file_path = tmp_path / file_name
+    with file_path.open("w") as f:
+        ndjson.dump(object_predictions, f)
+
+    annotation_import = model_run_with_model_run_data_rows.add_predictions(
+        name=name, predictions=str(file_path))
+
+    annotation_import, batch, mal_prediction_import = model_run_with_model_run_data_rows.upsert_predictions_and_send_to_project(
+        name=name,
+        predictions=str(file_path),
+        project_id=configured_project_without_data_rows.uid,
+        priority=5)
+
+    assert annotation_import.model_run_id == model_run_with_model_run_data_rows.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()
+
+    if batch:
+        assert batch.project().uid == configured_project_without_data_rows.uid
+
+    if mal_prediction_import:
+        annotation_import_test_helpers.check_running_state(
+            mal_prediction_import, name)
+        mal_prediction_import.wait_until_done()
+
+
+@pytest.mark.slow
+def test_wait_till_done(model_run_predictions,
+                        model_run_with_model_run_data_rows):
+    name = str(uuid.uuid4())
+    annotation_import = model_run_with_model_run_data_rows.add_predictions(
+        name=name, predictions=model_run_predictions)
+
+    assert len(annotation_import.inputs) == len(model_run_predictions)
+    annotation_import.wait_until_done()
+    assert annotation_import.state == AnnotationImportState.FINISHED
+    # Check that the status files are being returned as expected
+    assert len(annotation_import.errors) == 0
+    assert len(annotation_import.inputs) == len(model_run_predictions)
+    input_uuids = [
+        input_annot['uuid'] for input_annot in annotation_import.inputs
+    ]
+    inference_uuids = [pred['uuid'] for pred in model_run_predictions]
+    assert set(input_uuids) == set(inference_uuids)
+    assert len(annotation_import.statuses) == len(model_run_predictions)
+    for status in annotation_import.statuses:
+        assert status['status'] == 'SUCCESS'
+    status_uuids = [
+        input_annot['uuid'] for input_annot in annotation_import.statuses
+    ]
+    assert set(input_uuids) == set(status_uuids)