Skip to content

Commit db8ecd0

Browse files
ardilaUbuntu
andauthored
integration tests for public dataset (#162)
* integration tests for public dataset * add sanitize to fix test Co-authored-by: Ubuntu <diego.ardila@scale.com>
1 parent e2cfc7a commit db8ecd0

File tree

2 files changed

+28
-3
lines changed

2 files changed

+28
-3
lines changed

nucleus/dataset.py

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -564,9 +564,12 @@ def iloc(self, i: int) -> dict:
564564
}
565565
}
566566
"""
567-
response = self._client.dataitem_iloc(self.id, i)
567+
response = self._client.make_request(
568+
{}, f"dataset/{self.id}/iloc/{i}", requests.get
569+
)
568570
return format_dataset_item_response(response)
569571

572+
@sanitize_string_args
570573
def refloc(self, reference_id: str) -> dict:
571574
"""Retrieves a dataset item by reference ID.
572575
@@ -587,7 +590,9 @@ def refloc(self, reference_id: str) -> dict:
587590
}
588591
}
589592
"""
590-
response = self._client.dataitem_ref_id(self.id, reference_id)
593+
response = self._client.make_request(
594+
{}, f"dataset/{self.id}/refloc/{reference_id}", requests.get
595+
)
591596
return format_dataset_item_response(response)
592597

593598
def loc(self, dataset_item_id: str) -> dict:
@@ -611,7 +616,9 @@ def loc(self, dataset_item_id: str) -> dict:
611616
}
612617
}
613618
"""
614-
response = self._client.dataitem_loc(self.id, dataset_item_id)
619+
response = self._client.make_request(
620+
{}, f"dataset/{self.id}/loc/{dataset_item_id}", requests.get
621+
)
615622
return format_dataset_item_response(response)
616623

617624
def ground_truth_loc(self, reference_id: str, annotation_id: str):

tests/test_public_datasets.py

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
from nucleus.dataset import Dataset
2+
3+
PANDASET_ID = "ds_bwhjbyfb8mjj0ykagxf0"
4+
5+
6+
def test_get_pandaset_items(CLIENT):
7+
dataset: Dataset = CLIENT.get_dataset(PANDASET_ID)
8+
items = dataset.items
9+
items_and_annotations = dataset.items_and_annotations()
10+
11+
target_item = items[0]
12+
assert {_["item"].reference_id for _ in items_and_annotations} == set(
13+
[i.reference_id for i in items]
14+
)
15+
ref_item = dataset.refloc(target_item.reference_id)
16+
assert ref_item["item"] == target_item
17+
index_item = dataset.iloc(0)
18+
assert index_item["item"] in items

0 commit comments

Comments
 (0)