codekingpro/portable-devtools

1from chromadb.api.client import Client
2from chromadb.config import System
3from chromadb.test.property import invariants
4
5
6def test_log_purge(sqlite_persistent: System) -> None:
7    client = Client.from_system(sqlite_persistent)
8
9    first_collection = client.create_collection(
10        "first_collection", metadata={"hnsw:sync_threshold": 10, "hnsw:batch_size": 10}
11    )
12    second_collection = client.create_collection(
13        "second_collection", metadata={"hnsw:sync_threshold": 10, "hnsw:batch_size": 10}
14    )
15    collections = [first_collection, second_collection]
16
17    # (Does not trigger a purge)
18    for i in range(5):
19        first_collection.add(ids=str(i), embeddings=[i, i])
20
21    # (Should trigger a purge)
22    for i in range(100):
23        second_collection.add(ids=str(i), embeddings=[i, i])
24
25    # The purge of the second collection should not be blocked by the first
26    invariants.log_size_below_max(client._system, collections, True)
27
28
29def test_log_purge_with_multiple_collections(sqlite_persistent: System) -> None:
30    client = Client.from_system(sqlite_persistent)
31
32    first_collection = client.create_collection(
33        "first_collection", metadata={"hnsw:sync_threshold": 10, "hnsw:batch_size": 10}
34    )
35    second_collection = client.create_collection(
36        "second_collection", metadata={"hnsw:sync_threshold": 10, "hnsw:batch_size": 10}
37    )
38    collections = [first_collection, second_collection]
39
40    # (Does not trigger a purge)
41    for i in range(15):
42        first_collection.add(ids=str(i), embeddings=[i, i])
43
44    # (Should trigger a purge)
45    for i in range(25):
46        second_collection.add(ids=str(i), embeddings=[i, i])
47
48    invariants.log_size_for_collections_match_expected(
49        client._system, collections, True
50    )
51