codekingpro/portable-devtools

1# This tests a very minimal of test_add in test_add.py as a example based test
2# instead of a property based test. We can use the delta to get the property
3# test working and then enable
4import random
5import time
6from chromadb.api import ClientAPI
7from chromadb.test.conftest import (
8    multi_region_test,
9    reset,
10    skip_if_not_cluster,
11)
12from chromadb.test.property import invariants
13from chromadb.test.utils.wait_for_version_increase import (
14    wait_for_version_increase,
15    get_collection_version,
16)
17import numpy as np
18
19
20@skip_if_not_cluster()
21@multi_region_test
22def test_add(
23    client: ClientAPI,
24) -> None:
25    seed = time.time()
26    random.seed(seed)
27    print("Generating data with seed ", seed)
28    reset(client)
29    collection = client.create_collection(
30        name="test",
31        metadata={"hnsw:construction_ef": 128, "hnsw:search_ef": 128, "hnsw:M": 128},
32    )
33
34    # Add 1000 records, where each embedding has 3 dimensions randomly generated
35    # between 0 and 1
36    ids = []
37    embeddings = []
38    for i in range(1000):
39        ids.append(str(i))
40        embeddings.append(np.random.rand(1, 3)[0])
41        collection.add(
42            ids=[str(i)],
43            embeddings=[embeddings[-1]],
44        )
45
46    random_query = np.random.rand(1, 3)[0]
47    print("Generated data with seed ", seed)
48
49    invariants.ann_accuracy(
50        collection,
51        {
52            "ids": ids,
53            "embeddings": embeddings,
54            "metadatas": None,
55            "documents": None,
56        },
57        10,
58        query_embeddings=[random_query],
59    )
60
61
62@skip_if_not_cluster()
63@multi_region_test
64def test_add_include_all_with_compaction_delay(client: ClientAPI) -> None:
65    seed = time.time()
66    random.seed(seed)
67    print("Generating data with seed ", seed)
68    reset(client)
69    collection = client.create_collection(
70        name="test_add_include_all_with_compaction_delay",
71        metadata={"hnsw:construction_ef": 128, "hnsw:search_ef": 128, "hnsw:M": 128},
72    )
73    initial_version = get_collection_version(client, collection.name)
74
75    ids = []
76    embeddings = []
77    documents = []
78    for i in range(1000):
79        ids.append(str(i))
80        embeddings.append(np.random.rand(1, 3)[0])
81        documents.append(f"document_{i}")
82        collection.add(
83            ids=[str(i)],
84            embeddings=[embeddings[-1]],
85            documents=[documents[-1]],
86        )
87
88    wait_for_version_increase(client, collection.name, initial_version, 120)
89
90    random_query_1 = np.random.rand(1, 3)[0]
91    random_query_2 = np.random.rand(1, 3)[0]
92    print("Generated data with seed ", seed)
93
94    # Query the collection with a random query
95    invariants.ann_accuracy(
96        collection,
97        {
98            "ids": ids,
99            "embeddings": embeddings,
100            "metadatas": None,
101            "documents": documents,
102        },
103        10,
104        query_embeddings=[random_query_1, random_query_2],
105    )
106