refactor scene integration tests to pass with new payload

Author: Drew Kaul

nucleus/constants.py

@@ -57,6 +57,7 @@
 REFERENCE_IDS_KEY = "reference_ids"
 REFERENCE_ID_KEY = "reference_id"
 REQUEST_ID_KEY = "requestId"
+SCENES = "scenes"
 SEGMENTATIONS_KEY = "segmentations"
 SLICE_ID_KEY = "slice_id"
 STATUS_CODE_KEY = "status_code"
@@ -65,6 +66,7 @@
 TYPE_KEY = "type"
 UPDATED_ITEMS = "updated_items"
 UPDATE_KEY = "update"
+URL = "url"
 VERTICES_KEY = "vertices"
 WIDTH_KEY = "width"
 YAW_KEY = "yaw"

nucleus/dataset.py

@@ -28,7 +28,9 @@
     NAME_KEY,
     REFERENCE_IDS_KEY,
     REQUEST_ID_KEY,
+    SCENES,
     UPDATE_KEY,
+    URL,
 )
 from .dataset_item import (
     DatasetItem,
@@ -255,22 +257,22 @@ def upload_scenes(
         asynchronous: bool = False,
     ) -> Union[dict, AsyncJob]:
         """
-        Uploads scene with given frames to the dataset
+        Uploads scenes with given frames to the dataset
 
         Parameters:
-        :param payload: dictionary containing frames to be uploaded as a scene
+        :param payload: dictionary containing scenes to be uploaded
         :param update: if True, overwrite scene on collision
         :param aynchronous: if True, return a job object representing asynchronous ingestion job
         :return:
         {
             'dataset_id': str,
-            'new_items': int,
-            'updated_items': int,
-            'ignored_items': int,
+            'new_scenes': int,
         }
         """
         if asynchronous:
-            check_all_frame_paths_remote(payload[FRAMES])
+            for scene in payload[SCENES]:
+                for frame in scene[FRAMES]:
+                    check_all_frame_paths_remote(frame[URL])
             request_id = serialize_and_write_to_presigned_url(
                 [payload], self.id, self._client
             )

tests/helpers.py

@@ -26,7 +26,24 @@
     "https://scaleapi-cust-lidar.s3.us-west-1.amazonaws.com/test-scale/frame-0.json",
 ]
 
-TEST_LIDAR_FRAMES = {"frames": TEST_POINTCLOUD_URLS}
+TEST_LIDAR_SCENES = {
+    "scenes": [
+        {
+            "referenceId": "scene_1",
+            "frames": [
+                {
+                    "lidar": {
+                        "url": TEST_POINTCLOUD_URLS[0],
+                        "type": "pointcloud",
+                        "referenceId": "lidar_frame_1",
+                        "metadata": {},
+                    },
+                },
+            ],
+            "metadata": {},
+        },
+    ],
+}
 
 TEST_DATASET_ITEMS = [
     DatasetItem(TEST_IMG_URLS[0], "1"),
@@ -81,21 +98,21 @@ def reference_id_from_url(url):
 
 TEST_CUBOID_ANNOTATIONS = [
     {
-        "label": f"[Pytest] Cuboid Annotation ${i}",
+        "label": f"[Pytest] Cuboid Annotation #{i}",
         "geometry": {
             "position": {
-                "x": 50 * i + 10,
-                "y": 60 * i + 10,
-                "z": 70 * i + 10,
+                "x": 50 * i + 5,
+                "y": 60 * i + 5,
+                "z": 70 * i + 5,
             },
             "dimensions": {
-                "x": 10 * i,
-                "y": 20 * i,
-                "z": 30 * i,
+                "x": 10 * i + 10,
+                "y": 20 * i + 10,
+                "z": 30 * i + 10,
             },
             "yaw": 5 * i,
         },
-        "annotation_id": f"[Pytest] Cuboid Annotation Annotation Id{i}",
+        "annotation_id": f"[Pytest] Cuboid Annotation Annotation Id {i}",
     }
     for i in range(len(TEST_POINTCLOUD_URLS))
 ]

tests/test_scene.py

@@ -1,9 +1,10 @@
+from nucleus.constants import SCENES
 import pytest
 
 from .helpers import (
     TEST_DATASET_3D_NAME,
     TEST_CUBOID_ANNOTATIONS,
-    TEST_LIDAR_FRAMES,
+    TEST_LIDAR_SCENES,
     assert_cuboid_annotation_matches_dict,
 )
 
@@ -22,34 +23,30 @@ def dataset(CLIENT):
 
 
 def test_scene_upload_sync(dataset):
-    payload = TEST_LIDAR_FRAMES
+    payload = TEST_LIDAR_SCENES
     response = dataset.upload_scenes(payload)
 
     assert response["dataset_id"] == dataset.id
-    assert response["new_items"] == 1
-    assert response["updated_items"] == 0
-    assert response["ignored_items"] == 0
+    assert response["new_scenes"] == len(TEST_LIDAR_SCENES[SCENES])
 
 
 @pytest.mark.integration
 def test_scene_and_cuboid_upload_sync(dataset):
-    payload = TEST_LIDAR_FRAMES
+    payload = TEST_LIDAR_SCENES
     response = dataset.upload_scenes(payload)
 
     assert response["dataset_id"] == dataset.id
-    assert response["new_items"] == 1
-    assert response["updated_items"] == 0
-    assert response["ignored_items"] == 0
+    assert response["new_scenes"] == len(TEST_LIDAR_SCENES[SCENES])
 
     TEST_CUBOID_ANNOTATIONS[0]["dataset_item_id"] = dataset.items[0].item_id
-    annotation = CuboidAnnotation.from_json(TEST_CUBOID_ANNOTATIONS[0])
-    response = dataset.annotate(annotations=[annotation])
+    annotations = [CuboidAnnotation.from_json(TEST_CUBOID_ANNOTATIONS[0])]
+    response = dataset.annotate(annotations)
 
     assert response["dataset_id"] == dataset.id
-    assert response["annotations_processed"] == 1
+    assert response["annotations_processed"] == len(annotations)
     assert response["annotations_ignored"] == 0
 
-    response = dataset.loc(annotation.item_id)["annotations"]["cuboid"]
+    response = dataset.loc(annotations[0].item_id)["annotations"]["cuboid"]
     assert len(response) == 1
     response_annotation = response[0]
     assert_cuboid_annotation_matches_dict(