| import subprocess |
| import time |
|
|
| from core.rag.datasource.vdb.couchbase.couchbase_vector import CouchbaseConfig, CouchbaseVector |
| from tests.integration_tests.vdb.test_vector_store import ( |
| AbstractVectorTest, |
| get_example_text, |
| setup_mock_redis, |
| ) |
|
|
|
|
| def wait_for_healthy_container(service_name="couchbase-server", timeout=300): |
| start_time = time.time() |
| while time.time() - start_time < timeout: |
| result = subprocess.run( |
| ["docker", "inspect", "--format", "{{.State.Health.Status}}", service_name], capture_output=True, text=True |
| ) |
| if result.stdout.strip() == "healthy": |
| print(f"{service_name} is healthy!") |
| return True |
| else: |
| print(f"Waiting for {service_name} to be healthy...") |
| time.sleep(10) |
| raise TimeoutError(f"{service_name} did not become healthy in time") |
|
|
|
|
| class CouchbaseTest(AbstractVectorTest): |
| def __init__(self): |
| super().__init__() |
| self.vector = CouchbaseVector( |
| collection_name=self.collection_name, |
| config=CouchbaseConfig( |
| connection_string="couchbase://127.0.0.1", |
| user="Administrator", |
| password="password", |
| bucket_name="Embeddings", |
| scope_name="_default", |
| ), |
| ) |
|
|
| def search_by_vector(self): |
| |
| time.sleep(5) |
| hits_by_vector = self.vector.search_by_vector(query_vector=self.example_embedding) |
| assert len(hits_by_vector) == 1 |
|
|
|
|
| def test_couchbase(setup_mock_redis): |
| wait_for_healthy_container("couchbase-server", timeout=60) |
| CouchbaseTest().run_all_tests() |
|
|
