add testing for indexing endpoints

Author: sasha-scale

nucleus/constants.py

@@ -21,7 +21,8 @@
 ANNOTATION_UPDATE_KEY = "update"
 DEFAULT_ANNOTATION_UPDATE_MODE = False
 STATUS_CODE_KEY = "status_code"
-SUCCESS_STATUS_CODES = [200, 201]
+STATUS_KEY = "status"
+SUCCESS_STATUS_CODES = [200, 201, 202]
 ERRORS_KEY = "errors"
 MODEL_RUN_ID_KEY = "model_run_id"
 MODEL_ID_KEY = "model_id"
@@ -57,3 +58,5 @@
 INDEX_KEY = "index"
 SEGMENTATIONS_KEY = "segmentations"
 EMBEDDINGS_URL_KEY = "embeddings_url"
+JOB_ID_KEY = "job_id"
+MESSAGE_KEY = "message"

tests/helpers.py

@@ -120,6 +120,7 @@ def reference_id_from_url(url):
     for i in range(len(TEST_POLYGON_ANNOTATIONS))
 ]
 
+TEST_INDEX_EMBEDDINGS_FILE = "https://scale-ml.s3.amazonaws.com/tmp/sasha/pytest_embeddings_payload.json"
 
 # Asserts that a box annotation instance matches a dict representing its properties.
 # Useful to check annotation uploads/updates match.

tests/test_indexing.py

@@ -0,0 +1,53 @@
+import pytest
+
+from helpers import (
+    get_signed_url,
+    TEST_INDEX_EMBEDDINGS_FILE,
+    TEST_IMG_URLS,
+    TEST_DATASET_NAME,
+    reference_id_from_url,
+)
+
+from nucleus import DatasetItem
+
+from nucleus.constants import (
+    ERROR_PAYLOAD,
+    JOB_ID_KEY,
+    MESSAGE_KEY,
+    STATUS_KEY,
+)
+
+
+@pytest.fixture()
+def dataset(CLIENT):
+    ds = CLIENT.create_dataset(TEST_DATASET_NAME)
+    ds_items = []
+    for url in TEST_IMG_URLS:
+        ds_items.append(
+            DatasetItem(
+                image_location=url,
+                reference_id=reference_id_from_url(url),
+            )
+        )
+
+    response = ds.append(ds_items)
+    assert ERROR_PAYLOAD not in response.json()
+    yield ds
+
+    response = CLIENT.delete_dataset(ds.id)
+    assert response == {}
+
+
+def test_index_integration(dataset):
+    signed_embeddings_url = get_signed_url(TEST_INDEX_EMBEDDINGS_FILE)
+    create_response = dataset.create_custom_index(signed_embeddings_url)
+    assert JOB_ID_KEY in create_response
+    assert MESSAGE_KEY in create_response
+    job_id = create_response[JOB_ID_KEY]
+
+    # Job can error because pytest dataset fixture gets deleted
+    # As a workaround, we'll just check htat we got some response
+    job_status_response = dataset.check_index_status(job_id)
+    assert STATUS_KEY in job_status_response
+    assert JOB_ID_KEY in job_status_response
+    assert MESSAGE_KEY in job_status_response