integration tests for public dataset (#162)

* integration tests for public dataset

* add sanitize to fix test

Co-authored-by: Ubuntu <diego.ardila@scale.com>

Author: ardila

nucleus/dataset.py

@@ -564,9 +564,12 @@ def iloc(self, i: int) -> dict:
                     }
                 }
         """
-        response = self._client.dataitem_iloc(self.id, i)
+        response = self._client.make_request(
+            {}, f"dataset/{self.id}/iloc/{i}", requests.get
+        )
         return format_dataset_item_response(response)
 
+    @sanitize_string_args
     def refloc(self, reference_id: str) -> dict:
         """Retrieves a dataset item by reference ID.
 
@@ -587,7 +590,9 @@ def refloc(self, reference_id: str) -> dict:
                     }
                 }
         """
-        response = self._client.dataitem_ref_id(self.id, reference_id)
+        response = self._client.make_request(
+            {}, f"dataset/{self.id}/refloc/{reference_id}", requests.get
+        )
         return format_dataset_item_response(response)
 
     def loc(self, dataset_item_id: str) -> dict:
@@ -611,7 +616,9 @@ def loc(self, dataset_item_id: str) -> dict:
                     }
                 }
         """
-        response = self._client.dataitem_loc(self.id, dataset_item_id)
+        response = self._client.make_request(
+            {}, f"dataset/{self.id}/loc/{dataset_item_id}", requests.get
+        )
         return format_dataset_item_response(response)
 
     def ground_truth_loc(self, reference_id: str, annotation_id: str):

tests/test_public_datasets.py

@@ -0,0 +1,18 @@
+from nucleus.dataset import Dataset
+
+PANDASET_ID = "ds_bwhjbyfb8mjj0ykagxf0"
+
+
+def test_get_pandaset_items(CLIENT):
+    dataset: Dataset = CLIENT.get_dataset(PANDASET_ID)
+    items = dataset.items
+    items_and_annotations = dataset.items_and_annotations()
+
+    target_item = items[0]
+    assert {_["item"].reference_id for _ in items_and_annotations} == set(
+        [i.reference_id for i in items]
+    )
+    ref_item = dataset.refloc(target_item.reference_id)
+    assert ref_item["item"] == target_item
+    index_item = dataset.iloc(0)
+    assert index_item["item"] in items