fix failing tests

Author: Bihan Jiang

nucleus/__init__.py

@@ -1163,19 +1163,14 @@ def create_custom_index(
         embeddings_urls: list of urls, each of which being a json mapping dataset_item_id -> embedding vector
         embedding_dim: the dimension of the embedding vectors, must be consistent for all embedding vectors in the index.
         """
-        response_objects = self.make_request(
+        return self.make_request(
             {
                 EMBEDDINGS_URL_KEY: embeddings_urls,
                 EMBEDDING_DIMENSION_KEY: embedding_dim,
             },
             f"indexing/{dataset_id}",
             requests_command=requests.post,
         )
-        job = AsyncJob.from_json(response_objects, self)
-        dataset_id = response_objects[DATASET_ID_KEY]
-        message = response_objects[MESSAGE_KEY]
-
-        return dataset_id, job, message
 
     def check_index_status(self, job_id: str):
         return self.make_request(

nucleus/job.py

@@ -5,9 +5,9 @@
 from nucleus.constants import (
     JOB_CREATION_TIME_KEY,
     JOB_ID_KEY,
-    JOB_STATUS_KEY,
     JOB_LAST_KNOWN_STATUS_KEY,
     JOB_TYPE_KEY,
+    STATUS_KEY,
 )
 
 JOB_POLLING_INTERVAL = 5
@@ -27,7 +27,7 @@ def status(self) -> Dict[str, str]:
             route=f"job/{self.job_id}",
             requests_command=requests.get,
         )
-        self.job_last_known_status = response[JOB_STATUS_KEY]
+        self.job_last_known_status = response[STATUS_KEY]
         return response
 
     def errors(self) -> List[str]:

tests/test_dataset.py

@@ -169,12 +169,12 @@ def test_dataset_append_async(dataset: Dataset):
     status = job.status()
     status["message"]["PayloadUrl"] = ""
     assert status == {
-        "job_id": job.id,
+        "job_id": job.job_id,
         "status": "Completed",
         "message": {
             "PayloadUrl": "",
             "image_upload_step": {"errored": 0, "pending": 0, "completed": 5},
-            "started_image_processing": f"Dataset: {dataset.id}, Job: {job.id}",
+            "started_image_processing": f"Dataset: {dataset.id}, Job: {job.job_id}",
             "ingest_to_reupload_queue": {
                 "epoch": 1,
                 "total": 5,
@@ -204,7 +204,7 @@ def test_dataset_append_async_with_1_bad_url(dataset: Dataset):
     status = job.status()
     status["message"]["PayloadUrl"] = ""
     assert status == {
-        "job_id": f"{job.id}",
+        "job_id": f"{job.job_id}",
         "status": "Errored",
         "message": {
             "PayloadUrl": "",
@@ -220,7 +220,7 @@ def test_dataset_append_async_with_1_bad_url(dataset: Dataset):
                 "datasetId": f"{dataset.id}",
                 "processed": 5,
             },
-            "started_image_processing": f"Dataset: {dataset.id}, Job: {job.id}",
+            "started_image_processing": f"Dataset: {dataset.id}, Job: {job.job_id}",
         },
     }
     # The error is fairly detailed and subject to change. What's important is we surface which URLs failed.
@@ -286,7 +286,7 @@ def test_annotate_async(dataset: Dataset):
     )
     job.sleep_until_complete()
     assert job.status() == {
-        "job_id": job.id,
+        "job_id": job.job_id,
         "status": "Completed",
         "message": {
             "annotation_upload": {
@@ -321,7 +321,7 @@ def test_annotate_async_with_error(dataset: Dataset):
     job.sleep_until_complete()
 
     assert job.status() == {
-        "job_id": job.id,
+        "job_id": job.job_id,
         "status": "Completed",
         "message": {
             "annotation_upload": {

tests/test_indexing.py

@@ -1,3 +1,4 @@
+from nucleus.job import AsyncJob
 import pytest
 
 from .helpers import (
@@ -42,7 +43,11 @@ def test_index_integration(dataset):
     create_response = dataset.create_custom_index(
         [signed_embeddings_url], embedding_dim=3
     )
-    assert JOB_ID_KEY in create_response
+    job = AsyncJob.from_json(create_response, client="Nucleus Client")
+    assert job.job_id
+    assert job.job_last_known_status
+    assert job.job_type
+    assert job.job_creation_time
     assert MESSAGE_KEY in create_response
     job_id = create_response[JOB_ID_KEY]
 

tests/test_jobs.py

@@ -28,4 +28,4 @@ def test_job_creation_and_listing(CLIENT):
     jobs = CLIENT.list_jobs()
 
     for job in jobs:
-        assert eval(print(job)) == job
+        assert eval(str(job)) == job

tests/test_prediction.py

@@ -290,7 +290,7 @@ def test_mixed_pred_upload_async(model_run: ModelRun):
     job.sleep_until_complete()
 
     assert job.status() == {
-        "job_id": job.id,
+        "job_id": job.job_id,
         "status": "Completed",
         "message": {
             "prediction_upload": {
@@ -328,7 +328,7 @@ def test_mixed_pred_upload_async_with_error(model_run: ModelRun):
     job.sleep_until_complete()
 
     assert job.status() == {
-        "job_id": job.id,
+        "job_id": job.job_id,
         "status": "Completed",
         "message": {
             "prediction_upload": {