-
Notifications
You must be signed in to change notification settings - Fork 11
Expand file tree
/
Copy pathtest_indexing.py
More file actions
59 lines (49 loc) · 1.59 KB
/
Copy pathtest_indexing.py
File metadata and controls
59 lines (49 loc) · 1.59 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
from nucleus.job import AsyncJob
import pytest
from .helpers import (
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 == {"message": "Beginning dataset deletion..."}
def test_index_integration(dataset):
signed_embeddings_url = TEST_INDEX_EMBEDDINGS_FILE
create_response = dataset.create_custom_index(
[signed_embeddings_url], embedding_dim=3
)
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]
# 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