add tests for Frame class methods

Author: Drew Kaul

nucleus/constants.py

@@ -67,6 +67,7 @@
 NUCLEUS_ENDPOINT = "https://api.scale.com/v1/nucleus"
 NUM_SENSORS_KEY = "num_sensors"
 ORIGINAL_IMAGE_URL_KEY = "original_image_url"
+POINTCLOUD_KEY = "pointcloud"
 POINTCLOUD_LOCATION_KEY = "pointcloud_location"
 POINTCLOUD_URL_KEY = "pointcloud_url"
 POSITION_KEY = "position"

nucleus/scene.py

@@ -38,10 +38,10 @@ def get_item(self, sensor_name: str):
         return self.items[sensor_name]
 
     def get_items(self):
-        return self.items.values()
+        return list(self.items.values())
 
     def get_sensors(self):
-        return self.items.keys()
+        return list(self.items.keys())
 
     def get_index(self):
         return self.index

tests/helpers.py

@@ -24,6 +24,10 @@
 
 TEST_POINTCLOUD_URLS = [
     "https://scaleapi-cust-lidar.s3.us-west-1.amazonaws.com/test-scale/frame-0.json",
+    "https://scaleapi-cust-lidar.s3.us-west-1.amazonaws.com/test-scale/frame-1.json",
+    "https://scaleapi-cust-lidar.s3.us-west-1.amazonaws.com/test-scale/frame-2.json",
+    "https://scaleapi-cust-lidar.s3.us-west-1.amazonaws.com/test-scale/frame-3.json",
+    "https://scaleapi-cust-lidar.s3.us-west-1.amazonaws.com/test-scale/frame-4.json",
 ]
 
 TEST_LIDAR_SCENES = {
@@ -56,7 +60,14 @@
                             }
                         },
                     },
-                }
+                },
+                {
+                    "lidar": {
+                        "pointcloud_url": TEST_POINTCLOUD_URLS[0],
+                        "reference_id": "lidar_frame_2",
+                        "metadata": {},
+                    },
+                },
             ],
             "metadata": {},
         },
@@ -71,6 +82,14 @@
     DatasetItem(TEST_IMG_URLS[3], "4"),
 ]
 
+TEST_LIDAR_ITEMS = [
+    DatasetItem(pointcloud_location=TEST_POINTCLOUD_URLS[0], reference_id="1"),
+    DatasetItem(pointcloud_location=TEST_POINTCLOUD_URLS[1], reference_id="2"),
+    DatasetItem(pointcloud_location=TEST_POINTCLOUD_URLS[2], reference_id="3"),
+    DatasetItem(pointcloud_location=TEST_POINTCLOUD_URLS[3], reference_id="4"),
+    DatasetItem(pointcloud_location=TEST_POINTCLOUD_URLS[4], reference_id="5"),
+]
+
 LOCAL_FILENAME = "tests/test_img.jpg"
 TEST_PREDS = [
     BoxPrediction("[Pytest Box Prediction 1]", 0, 0, 100, 100, "1"),

tests/test_scene.py

@@ -3,7 +3,11 @@
     ANNOTATIONS_KEY,
     DATASET_ITEM_ID_KEY,
     FRAMES_KEY,
+    IMAGE_KEY,
+    IMAGE_URL_KEY,
     ITEM_KEY,
+    POINTCLOUD_KEY,
+    POINTCLOUD_URL_KEY,
     REFERENCE_ID_KEY,
     SCENES_KEY,
     UPDATE_KEY,
@@ -12,11 +16,14 @@
 from nucleus import (
     CuboidAnnotation,
     LidarScene,
+    Frame,
 )
 
 from .helpers import (
     TEST_DATASET_3D_NAME,
     TEST_CUBOID_ANNOTATIONS,
+    TEST_DATASET_ITEMS,
+    TEST_LIDAR_ITEMS,
     TEST_LIDAR_SCENES,
     assert_cuboid_annotation_matches_dict,
 )
@@ -31,6 +38,32 @@ def dataset(CLIENT):
     assert response == {"message": "Beginning dataset deletion..."}
 
 
+def test_frame_add_item(dataset):
+    frame = Frame(index=0)
+    frame.add_item(TEST_DATASET_ITEMS[0], "camera")
+    frame.add_item(TEST_LIDAR_ITEMS[0], "lidar")
+
+    assert frame.get_index() == 0
+    assert frame.get_sensors() == ["camera", "lidar"]
+    for item in frame.get_items():
+        assert item in [TEST_DATASET_ITEMS[0], TEST_LIDAR_ITEMS[0]]
+    assert frame.get_item("lidar") == TEST_LIDAR_ITEMS[0]
+    assert frame.to_payload() == {
+        "camera": {
+            "url": TEST_DATASET_ITEMS[0].image_location,
+            "reference_id": TEST_DATASET_ITEMS[0].reference_id,
+            "type": IMAGE_KEY,
+            "metadata": TEST_DATASET_ITEMS[0].metadata or {},
+        },
+        "lidar": {
+            "url": TEST_LIDAR_ITEMS[0].pointcloud_location,
+            "reference_id": TEST_LIDAR_ITEMS[0].reference_id,
+            "type": POINTCLOUD_KEY,
+            "metadata": TEST_LIDAR_ITEMS[0].metadata or {},
+        },
+    }
+
+
 def test_scene_upload_sync(dataset):
     payload = TEST_LIDAR_SCENES
     scenes = [
@@ -44,33 +77,6 @@ def test_scene_upload_sync(dataset):
     assert response["new_scenes"] == len(scenes)
 
 
-@pytest.mark.integration
-def test_scene_upload_async(dataset):
-    payload = TEST_LIDAR_SCENES
-    scenes = [
-        LidarScene.from_json(scene_json) for scene_json in payload[SCENES_KEY]
-    ]
-    update = payload[UPDATE_KEY]
-
-    job = dataset.append(scenes, update=update, asynchronous=True)
-    job.sleep_until_complete()
-    status = job.status()
-
-    assert status == {
-        "job_id": job.job_id,
-        "status": "Completed",
-        "message": {
-            "SceneUploadResponse": {
-                "errors": [],
-                "dataset_id": dataset.id,
-                "new_scenes": len(scenes),
-                "ignored_scenes": 0,
-                "scenes_errored": 0,
-            }
-        },
-    }
-
-
 @pytest.mark.integration
 def test_scene_and_cuboid_upload_sync(dataset):
     payload = TEST_LIDAR_SCENES
@@ -103,3 +109,30 @@ def test_scene_and_cuboid_upload_sync(dataset):
     assert_cuboid_annotation_matches_dict(
         response_annotations[0], TEST_CUBOID_ANNOTATIONS[0]
     )
+
+
+@pytest.mark.integration
+def test_scene_upload_async(dataset):
+    payload = TEST_LIDAR_SCENES
+    scenes = [
+        LidarScene.from_json(scene_json) for scene_json in payload[SCENES_KEY]
+    ]
+    update = payload[UPDATE_KEY]
+
+    job = dataset.append(scenes, update=update, asynchronous=True)
+    job.sleep_until_complete()
+    status = job.status()
+
+    assert status == {
+        "job_id": job.job_id,
+        "status": "Completed",
+        "message": {
+            "SceneUploadResponse": {
+                "errors": [],
+                "dataset_id": dataset.id,
+                "new_scenes": len(scenes),
+                "ignored_scenes": 0,
+                "scenes_errored": 0,
+            }
+        },
+    }