codekingpro/portable-devtools

1from typing import Sequence
2from chromadb.test.conftest import (
3    reset,
4    skip_if_not_cluster,
5)
6from chromadb.api import ClientAPI
7from kubernetes import client as k8s_client, config
8import time
9
10
11@skip_if_not_cluster()
12def test_reroute(
13    client: ClientAPI,
14) -> None:
15    reset(client)
16    collection = client.create_collection(
17        name="test",
18        metadata={"hnsw:construction_ef": 128, "hnsw:search_ef": 128, "hnsw:M": 128},
19    )
20
21    ids = [str(i) for i in range(10)]
22    embeddings: list[Sequence[float]] = [
23        [float(i), float(i), float(i)] for i in range(10)
24    ]
25    collection.add(ids=ids, embeddings=embeddings)
26    collection.query(query_embeddings=[embeddings[0]])
27
28    # Restart the query service using k8s api, in order to trigger a reroute
29    # of the query service
30    config.load_kube_config()
31    v1 = k8s_client.CoreV1Api()
32    # Find all pods with the label "app=query"
33    res = v1.list_namespaced_pod("chroma", label_selector="app=query-service")
34    assert len(res.items) > 0
35    items = res.items
36    seen_ids = set()
37
38    # Restart all the pods by deleting them
39    for item in items:
40        seen_ids.add(item.metadata.uid)
41        name = item.metadata.name
42        namespace = item.metadata.namespace
43        v1.delete_namespaced_pod(name, namespace)
44
45    # Wait until we have len(seen_ids) pods running with new UIDs
46    timeout_secs = 10
47    start_time = time.time()
48    while True:
49        res = v1.list_namespaced_pod("chroma", label_selector="app=query-service")
50        items = res.items
51        new_ids = set([item.metadata.uid for item in items])
52        if len(new_ids) == len(seen_ids) and len(new_ids.intersection(seen_ids)) == 0:
53            break
54        if time.time() - start_time > timeout_secs:
55            assert False, "Timed out waiting for new pods to start"
56        time.sleep(1)
57
58    # Wait for the query service to be ready, or timeout
59    while True:
60        res = v1.list_namespaced_pod("chroma", label_selector="app=query-service")
61        items = res.items
62        ready = True
63        for item in items:
64            if item.status.phase != "Running":
65                ready = False
66                break
67        if ready:
68            break
69        if time.time() - start_time > timeout_secs:
70            assert False, "Timed out waiting for new pods to be ready"
71        time.sleep(1)
72
73    time.sleep(1)
74    collection.query(query_embeddings=[embeddings[0]])
75