| | 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() |
| |
|
