Await healthcheck for testing

Author: ejscribner

api/core/rag/datasource/vdb/couchbase/couchbase_vector.py

@@ -16,10 +16,10 @@
 from flask import current_app
 from pydantic import BaseModel, model_validator
 
-from core.rag.datasource.entity.embedding import Embeddings
 from core.rag.datasource.vdb.vector_base import BaseVector
 from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory
 from core.rag.datasource.vdb.vector_type import VectorType
+from core.rag.embedding.embedding_base import Embeddings
 from core.rag.models.document import Document
 from extensions.ext_redis import redis_client
 from models.dataset import Dataset

api/tests/integration_tests/vdb/couchbase/test_couchbase.py

@@ -1,3 +1,4 @@
+import subprocess
 import time
 
 from core.rag.datasource.vdb.couchbase.couchbase_vector import CouchbaseConfig, CouchbaseVector
@@ -8,21 +9,42 @@
 )
 
 
+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):
-        time.sleep(20)
         super().__init__()
         self.vector = CouchbaseVector(
             collection_name=self.collection_name,
             config=CouchbaseConfig(
-                connection_string="127.0.0.1",
+                connection_string="couchbase://127.0.0.1",
                 user="Administrator",
                 password="password",
                 bucket_name="Embeddings",
                 scope_name="_default",
             ),
         )
 
+    def search_by_vector(self):
+        # brief sleep to ensure document is indexed
+        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()

dev/pytest/pytest_vdb.sh

@@ -11,5 +11,5 @@ pytest api/tests/integration_tests/vdb/chroma \
   api/tests/integration_tests/vdb/vikingdb \
   api/tests/integration_tests/vdb/baidu \
   api/tests/integration_tests/vdb/tcvectordb \
-  api/tests/integration_tests/vdb/upstash
+  api/tests/integration_tests/vdb/upstash \
   api/tests/integration_tests/vdb/couchbase \

docker/couchbase-server/init-cbserver.sh

@@ -11,9 +11,9 @@ if ! [ -f "$FILE" ]; then
   # used to automatically create the cluster based on environment variables
   # https://docs.couchbase.com/server/current/cli/cbcli/couchbase-cli-cluster-init.html
 
-  echo $COUCHBASE_ADMINISTRATOR_USERNAME ":"  $COUCHBASE_ADMINISTRATOR_PASSWORD  
+  echo $COUCHBASE_ADMINISTRATOR_USERNAME ":"  $COUCHBASE_ADMINISTRATOR_PASSWORD
 
-  sleep 10s 
+  sleep 20s
   /opt/couchbase/bin/couchbase-cli cluster-init -c 127.0.0.1 \
   --cluster-username $COUCHBASE_ADMINISTRATOR_USERNAME \
   --cluster-password $COUCHBASE_ADMINISTRATOR_PASSWORD \
@@ -24,7 +24,7 @@ if ! [ -f "$FILE" ]; then
   --cluster-fts-ramsize $COUCHBASE_FTS_RAM_SIZE \
   --index-storage-setting default
 
-  sleep 2s 
+  sleep 2s
 
   # used to auto create the bucket based on environment variables
   # https://docs.couchbase.com/server/current/cli/cbcli/couchbase-cli-bucket-create.html
@@ -41,4 +41,4 @@ if ! [ -f "$FILE" ]; then
 fi 
   # docker compose will stop the container from running unless we do this
   # known issue and workaround
-  tail -f /dev/null
+  tail -f /dev/null

docker/docker-compose.yaml

@@ -480,7 +480,7 @@ services:
     environment:
       QDRANT_API_KEY: ${QDRANT_API_KEY:-difyai123456}
 
-    # The Couchbase vector store.
+  # The Couchbase vector store.
   couchbase-server:
     build: ./couchbase-server
     profiles:
@@ -505,6 +505,13 @@ services:
     command: sh -c "/opt/couchbase/init/init-cbserver.sh"
     volumes:
       - ./volumes/couchbase/data:/opt/couchbase/var/lib/couchbase/data
+    healthcheck:
+      # ensure bucket was created before proceeding
+      test: [ "CMD-SHELL", "curl -s -f -u Administrator:password http://localhost:8091/pools/default/buckets | grep -q '\\[{' || exit 1" ]
+      interval: 10s
+      retries: 10
+      start_period: 30s
+      timeout: 10s
 
   # The pgvector vector database.
   pgvector: