fix pytest

Author: sasha-scale

nucleus/__init__.py

@@ -723,6 +723,7 @@ def predict(
 
         for s_batch in s_batches:
             payload = construct_segmentation_payload(s_batch, update)
+            print("PAYLOAD: ", payload)
             response = self._make_request(
                 payload, f"modelRun/{model_run_id}/predict_segmentation"
             )

nucleus/prediction.py

@@ -37,6 +37,7 @@ def from_json(cls, payload: dict):
             ],
             reference_id=payload.get(REFERENCE_ID_KEY, None),
             item_id=payload.get(ITEM_ID_KEY, None),
+            annotation_id=payload.get(ANNOTATION_ID_KEY, None),
         )
 
 

tests/test_dataset.py

@@ -1,6 +1,6 @@
 import pytest
 
-from helpers import (
+from tests.helpers import (
     TEST_SLICE_NAME,
     TEST_DATASET_NAME,
     TEST_IMG_URLS,
@@ -30,6 +30,7 @@ def dataset(CLIENT):
 def test_dataset_create_and_delete(CLIENT):
     # Creation
     ds = CLIENT.create_dataset(TEST_DATASET_NAME)
+    print(ds)
     assert isinstance(ds, Dataset)
     assert ds.name == TEST_DATASET_NAME
     assert ds.model_runs == []
@@ -119,7 +120,7 @@ def test_slice_create_and_delete_and_list(dataset):
     assert response["dataset_id"] == dataset.id
     assert len(response["dataset_items"]) == 2
     for item in ds_items[:2]:
-        assert item.reference_id == response["dataset_items"][0]["reference_id"] or item.reference_id == response["dataset_items"][1]["reference_id"]
+        assert item.reference_id == response["dataset_items"][0]["ref_id"] or item.reference_id == response["dataset_items"][1]["ref_id"]
 
 
 def test_slice_append(dataset):
@@ -145,6 +146,7 @@ def test_slice_append(dataset):
     slc.append(reference_ids=[item.reference_id for item in ds_items[:3]])
 
     response = slc.info()
+    print(response)
     assert len(response["dataset_items"]) == 3
     for item in ds_items[:3]:
-        assert item.reference_id == response["dataset_items"][0]["reference_id"] or item.reference_id == response["dataset_items"][1]["reference_id"] or item.reference_id == response["dataset_items"][2]["reference_id"]
+        assert item.reference_id == response["dataset_items"][0]["ref_id"] or item.reference_id == response["dataset_items"][1]["ref_id"] or item.reference_id == response["dataset_items"][2]["ref_id"]