[SDK-35] Improve mea tests (#1243)

vbrodsky · web-flow · commit 40f4dbcc1fa9 · 2023-10-02T16:03:19.000-07:00
diff --git a/labelbox/schema/data_row_metadata.py b/labelbox/schema/data_row_metadata.py
@@ -255,6 +255,7 @@ def _parse_ontology(raw_ontology) -> List[DataRowMetadataSchema]:
                 options = []
                 for option in schema["options"]:
                     option["uid"] = option["id"]
+
                     options.append(
                         DataRowMetadataSchema(**{
                             **option,
diff --git a/tests/integration/annotation_import/conftest.py b/tests/integration/annotation_import/conftest.py
@@ -517,7 +517,7 @@ def configured_project(client, initial_dataset, ontology, rand_gen, image_url):
     project._wait_until_data_rows_are_processed(data_row_ids=data_row_ids,
                                                 sleep_interval=3)
 
-    batch = project.create_batch(
+    project.create_batch(
         rand_gen(str),
         data_row_ids,  # sample of data row objects
         5  # priority between 1(Highest) - 5(lowest)
@@ -894,9 +894,9 @@ def segmentation_inference(prediction_id_mapping):
     segmentation = prediction_id_mapping['superpixel'].copy()
     segmentation.update({
         'mask': {
-            # TODO: Use a real URI
-            'instanceURI': "sampleuri",
-            'colorRGB': [0, 0, 0]
+            "instanceURI":
+                "https://storage.googleapis.com/labelbox-datasets/image_sample_data/raster_seg.png",
+            "colorRGB": (255, 255, 255)
         }
     })
     del segmentation['tool']
@@ -1006,7 +1006,6 @@ def model_run_predictions(polygon_inference, rectangle_inference,
     return [polygon_inference, rectangle_inference, line_inference]
 
 
-# also used for label imports
 @pytest.fixture
 def object_predictions(polygon_inference, rectangle_inference, line_inference,
                        entity_inference, segmentation_inference):
@@ -1016,6 +1015,17 @@ def object_predictions(polygon_inference, rectangle_inference, line_inference,
     ]
 
 
+@pytest.fixture
+def object_predictions_for_annotation_import(polygon_inference,
+                                             rectangle_inference,
+                                             line_inference,
+                                             segmentation_inference):
+    return [
+        polygon_inference, rectangle_inference, line_inference,
+        segmentation_inference
+    ]
+
+
 @pytest.fixture
 def classification_predictions(checklist_inference, text_inference):
     return [checklist_inference, text_inference]
@@ -1074,6 +1084,8 @@ def model_run_with_training_metadata(rand_gen, model):
 def model_run_with_data_rows(client, configured_project, model_run_predictions,
                              model_run, wait_for_label_processing):
     configured_project.enable_model_assisted_labeling()
+    use_data_row_ids = [p['dataRow']['id'] for p in model_run_predictions]
+    model_run.upsert_data_rows(use_data_row_ids)
 
     upload_task = LabelImport.create_from_objects(
         client, configured_project.uid, f"label-import-{uuid.uuid4()}",
@@ -1096,6 +1108,8 @@ def model_run_with_all_project_labels(client, configured_project,
                                       model_run_predictions, model_run,
                                       wait_for_label_processing):
     configured_project.enable_model_assisted_labeling()
+    use_data_row_ids = [p['dataRow']['id'] for p in model_run_predictions]
+    model_run.upsert_data_rows(use_data_row_ids)
 
     upload_task = LabelImport.create_from_objects(
         client, configured_project.uid, f"label-import-{uuid.uuid4()}",
diff --git a/tests/integration/annotation_import/test_mea_prediction_import.py b/tests/integration/annotation_import/test_mea_prediction_import.py
@@ -11,21 +11,13 @@
 """
 
 
-def test_create_from_url(model_run_with_data_rows,
-                         annotation_import_test_helpers):
-    name = str(uuid.uuid4())
-    url = "https://storage.googleapis.com/labelbox-public-bucket/predictions_test_v2.ndjson"
-    annotation_import = model_run_with_data_rows.add_predictions(
-        name=name, predictions=url)
-    assert annotation_import.model_run_id == model_run_with_data_rows.uid
-    annotation_import_test_helpers.check_running_state(annotation_import, name,
-                                                       url)
-    annotation_import.wait_until_done()
-
-
-def test_create_from_objects(model_run_with_data_rows, object_predictions,
+def test_create_from_objects(model_run_with_data_rows,
+                             object_predictions_for_annotation_import,
                              annotation_import_test_helpers):
     name = str(uuid.uuid4())
+    object_predictions = object_predictions_for_annotation_import
+    use_data_row_ids = [p['dataRow']['id'] for p in object_predictions]
+    model_run_with_data_rows.upsert_data_rows(use_data_row_ids)
 
     annotation_import = model_run_with_data_rows.add_predictions(
         name=name, predictions=object_predictions)
@@ -36,15 +28,19 @@ def test_create_from_objects(model_run_with_data_rows, object_predictions,
         annotation_import.input_file_url, object_predictions)
     annotation_import.wait_until_done()
 
+    assert annotation_import.state == AnnotationImportState.FINISHED
+    annotation_import_test_helpers.download_and_assert_status(
+        annotation_import.status_file_url)
+
 
 def test_create_from_objects_global_key(client, model_run_with_data_rows,
-                                        entity_inference,
+                                        polygon_inference,
                                         annotation_import_test_helpers):
     name = str(uuid.uuid4())
-    dr = client.get_data_row(entity_inference['dataRow']['id'])
-    del entity_inference['dataRow']['id']
-    entity_inference['dataRow']['globalKey'] = dr.global_key
-    object_predictions = [entity_inference]
+    dr = client.get_data_row(polygon_inference['dataRow']['id'])
+    del polygon_inference['dataRow']['id']
+    polygon_inference['dataRow']['globalKey'] = dr.global_key
+    object_predictions = [polygon_inference]
 
     annotation_import = model_run_with_data_rows.add_predictions(
         name=name, predictions=object_predictions)
@@ -55,6 +51,10 @@ def test_create_from_objects_global_key(client, model_run_with_data_rows,
         annotation_import.input_file_url, object_predictions)
     annotation_import.wait_until_done()
 
+    assert annotation_import.state == AnnotationImportState.FINISHED
+    annotation_import_test_helpers.download_and_assert_status(
+        annotation_import.status_file_url)
+
 
 def test_create_from_objects_with_confidence(predictions_with_confidence,
                                              model_run_with_data_rows,
@@ -77,15 +77,20 @@ def test_create_from_objects_with_confidence(predictions_with_confidence,
     annotation_import_test_helpers.assert_file_content(
         annotation_import.input_file_url, predictions_with_confidence)
     annotation_import.wait_until_done()
+
     assert annotation_import.state == AnnotationImportState.FINISHED
     annotation_import_test_helpers.download_and_assert_status(
         annotation_import.status_file_url)
 
 
 def test_create_from_objects_all_project_labels(
-        model_run_with_all_project_labels, object_predictions,
+        model_run_with_all_project_labels,
+        object_predictions_for_annotation_import,
         annotation_import_test_helpers):
     name = str(uuid.uuid4())
+    object_predictions = object_predictions_for_annotation_import
+    use_data_row_ids = [p['dataRow']['id'] for p in object_predictions]
+    model_run_with_all_project_labels.upsert_data_rows(use_data_row_ids)
 
     annotation_import = model_run_with_all_project_labels.add_predictions(
         name=name, predictions=object_predictions)
@@ -96,13 +101,18 @@ def test_create_from_objects_all_project_labels(
         annotation_import.input_file_url, object_predictions)
     annotation_import.wait_until_done()
 
+    assert annotation_import.state == AnnotationImportState.FINISHED
+    annotation_import_test_helpers.download_and_assert_status(
+        annotation_import.status_file_url)
+
 
 def test_model_run_project_labels(model_run_with_all_project_labels,
                                   model_run_predictions):
     model_run = model_run_with_all_project_labels
     # TODO: Move to export_v2
     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']
@@ -121,11 +131,17 @@ def test_model_run_project_labels(model_run_with_all_project_labels,
         assert actual_label['DataRow ID'] == expected_label['dataRow']['id']
 
 
-def test_create_from_label_objects(model_run_with_data_rows, object_predictions,
+def test_create_from_label_objects(model_run_with_data_rows,
+                                   object_predictions_for_annotation_import,
                                    annotation_import_test_helpers):
     name = str(uuid.uuid4())
+    use_data_row_ids = [
+        p['dataRow']['id'] for p in object_predictions_for_annotation_import
+    ]
+    model_run_with_data_rows.upsert_data_rows(use_data_row_ids)
 
-    predictions = list(NDJsonConverter.deserialize(object_predictions))
+    predictions = list(
+        NDJsonConverter.deserialize(object_predictions_for_annotation_import))
 
     annotation_import = model_run_with_data_rows.add_predictions(
         name=name, predictions=predictions)
@@ -137,25 +153,39 @@ def test_create_from_label_objects(model_run_with_data_rows, object_predictions,
         annotation_import.input_file_url, normalized_predictions)
     annotation_import.wait_until_done()
 
+    assert annotation_import.state == AnnotationImportState.FINISHED
+    annotation_import_test_helpers.download_and_assert_status(
+        annotation_import.status_file_url)
+
 
 def test_create_from_local_file(tmp_path, model_run_with_data_rows,
-                                object_predictions,
+                                object_predictions_for_annotation_import,
                                 annotation_import_test_helpers):
+    use_data_row_ids = [
+        p['dataRow']['id'] for p in object_predictions_for_annotation_import
+    ]
+    model_run_with_data_rows.upsert_data_rows(use_data_row_ids)
+
     name = str(uuid.uuid4())
     file_name = f"{name}.ndjson"
     file_path = tmp_path / file_name
     with file_path.open("w") as f:
-        parser.dump(object_predictions, f)
+        parser.dump(object_predictions_for_annotation_import, f)
 
     annotation_import = model_run_with_data_rows.add_predictions(
         name=name, predictions=str(file_path))
 
     assert annotation_import.model_run_id == model_run_with_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.input_file_url,
+        object_predictions_for_annotation_import)
     annotation_import.wait_until_done()
 
+    assert annotation_import.state == AnnotationImportState.FINISHED
+    annotation_import_test_helpers.download_and_assert_status(
+        annotation_import.status_file_url)
+
 
 def test_get(client, model_run_with_data_rows, annotation_import_test_helpers):
     name = str(uuid.uuid4())
diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py
@@ -25,6 +25,7 @@
 from support.integration_client import Environ, IntegrationClient, EphemeralClient, AdminClient
 
 IMG_URL = "https://picsum.photos/200/300.jpg"
+MASKABLE_IMG_URL = "https://storage.googleapis.com/labelbox-datasets/image_sample_data/2560px-Kitano_Street_Kobe01s5s4110.jpeg"
 SMALL_DATASET_URL = "https://storage.googleapis.com/lb-artifacts-testing-public/sdk_integration_test/potato.jpeg"
 DATA_ROW_PROCESSING_WAIT_TIMEOUT_SECONDS = 30
 DATA_ROW_PROCESSING_WAIT_SLEEP_INTERNAL_SECONDS = 3
@@ -108,7 +109,7 @@ def client(environ: str):
 
 @pytest.fixture(scope="session")
 def image_url(client):
-    return client.upload_data(requests.get(IMG_URL).content,
+    return client.upload_data(requests.get(MASKABLE_IMG_URL).content,
                               content_type="image/jpeg",
                               filename="image.jpeg",
                               sign=True)
diff --git a/tests/integration/export_v2/test_export_video.py b/tests/integration/export_v2/test_export_video.py
@@ -1,17 +1,34 @@
 import time
+
+import pytest
 import labelbox as lb
 from labelbox.data.annotation_types.data.video import VideoData
 import labelbox.types as lb_types
 from labelbox.schema.annotation_import import AnnotationImportState
 
 
+@pytest.fixture
+def user_id(client):
+    return client.get_user().uid
+
+
+@pytest.fixture
+def org_id(client):
+    return client.get_organization().uid
+
+
 def test_export_v2_video(client, configured_project_without_data_rows,
-                         video_data, video_data_row,
-                         bbox_video_annotation_objects, rand_gen):
+                         video_data, video_data_row, is_adv_enabled,
+                         bbox_video_annotation_objects, rand_gen, user_id,
+                         org_id):
+
+    orgid = client.get_organization().uid
+    userid = client.get_user().uid
 
     project = configured_project_without_data_rows
     project_id = project.uid
     labels = []
+
     _, data_row_uids = video_data
     project.create_batch(
         rand_gen(str),
@@ -156,6 +173,7 @@ def test_export_v2_video(client, configured_project_without_data_rows,
            )  #note we create 1 label per data row, 1 data row so 1 label
     export_label = project_export_labels[0]
     assert (export_label['label_kind']) == 'Video'
+
     assert (export_label['label_details'].keys()
            ) == expected_export_label['label_details'].keys()
 
