run integration tests

.dockerignore

@@ -0,0 +1,3 @@
+venv
+.git
+examples

.gitignore

@@ -15,4 +15,5 @@ index_data
 /index_data
 
 venv
-.env
+.env
+.chroma

Dockerfile

@@ -1,4 +1,4 @@
-FROM --platform=linux/arm64 python:3.10 AS chroma_server
+FROM --platform=linux/arm64 python:3.10
 
 #RUN apt-get update -qq
 #RUN apt-get install python3.10 python3-pip -y --no-install-recommends && rm -rf /var/lib/apt/lists_/*
@@ -13,4 +13,4 @@ COPY ./ /
 
 EXPOSE 8000
 
-CMD ["uvicorn", "chroma.app:app", "--host", "0.0.0.0", "--port", "8000", "--proxy-headers"]
+CMD ["uvicorn", "chroma.app:app", "--host", "0.0.0.0", "--port", "8000", "--workers", "1", "--proxy-headers"]

bin/integration-test

@@ -0,0 +1,19 @@
+#!/usr/bin/env bash
+
+set -e
+
+function cleanup {
+  docker-compose down --rmi local --volumes
+}
+
+trap cleanup EXIT
+
+docker-compose up --build -d
+
+export CHROMA_INTEGRATION_TEST=1
+export CHROMA_API_IMPL=rest
+export CHROMA_SERVER_HOST=localhost
+export CHROMA_SERVER_HTTP_PORT=8000
+
+python -m pytest
+

chroma/app.py

@@ -1,8 +1,5 @@
 import chroma
 from chroma.server.fastapi import FastAPI
-settings = chroma.config.Settings(
-                                chroma_db_impl="clickhouse",
-                                clickhouse_host="clickhouse",
-                                clickhouse_port="9000",)
+settings = chroma.config.Settings()
 server = FastAPI(settings)
 app = server.app()

chroma/test/test_api.py

@@ -5,6 +5,7 @@
 import time
 import tempfile
 import copy
+import os
 from multiprocessing import Process
 import uvicorn
 from requests.exceptions import ConnectionError
@@ -15,6 +16,11 @@ def local_api():
                                    chroma_db_impl="duckdb",
                                    chroma_cache_dir=tempfile.gettempdir()))
 
+@pytest.fixture
+def fastapi_integration_api():
+    return chroma.get_api() # configured by environment variables
+
+
 def _build_fastapi_api():
     return chroma.get_api(Settings(chroma_api_impl="rest",
                                    chroma_server_host="localhost",
@@ -55,9 +61,12 @@ def fastapi_server():
     yield
     proc.kill()
 
-
 test_apis = [local_api, fastapi_api]
 
+if 'CHROMA_INTEGRATION_TEST' in os.environ:
+    print("Including integration tests")
+    test_apis.append(fastapi_integration_api)
+
 
 @pytest.mark.parametrize('api_fixture', test_apis)
 def test_heartbeat(api_fixture, request):

chroma/test/test_chroma.py

@@ -1,5 +1,6 @@
 import pytest
 import unittest
+import os
 from unittest.mock import patch
 
 import chroma
@@ -32,13 +33,15 @@ class GetAPITest(unittest.TestCase):
 
     @patch('chroma.db.duckdb.DuckDB', autospec=True)
     @patch('chroma.api.local.LocalAPI', autospec=True)
+    @patch.dict(os.environ, {}, clear=True)
     def test_local(self, mock_api, mock_db):
         api = chroma.get_api(chroma.config.Settings(chroma_cache_dir="./foo"))
         assert mock_api.called
         assert mock_db.called
 
     @patch('chroma.db.duckdb.DuckDB', autospec=True)
     @patch('chroma.api.celery.CeleryAPI', autospec=True)
+    @patch.dict(os.environ, {}, clear=True)
     def test_celery(self, mock_api, mock_db):
         api = chroma.get_api(chroma.config.Settings(chroma_api_impl="celery",
                                                     chroma_cache_dir="./foo",
@@ -49,6 +52,7 @@ def test_celery(self, mock_api, mock_db):
 
 
     @patch('chroma.api.fastapi.FastAPI', autospec=True)
+    @patch.dict(os.environ, {}, clear=True)
     def test_fastapi(self, mock):
         api = chroma.get_api(chroma.config.Settings(chroma_api_impl="rest",
                                                     chroma_cache_dir="./foo",
@@ -58,6 +62,7 @@ def test_fastapi(self, mock):
 
 
     @patch('chroma.api.arrowflight.ArrowFlightAPI', autospec=True)
+    @patch.dict(os.environ, {}, clear=True)
     def test_arrowflight(self, mock):
         api = chroma.get_api(chroma.config.Settings(chroma_api_impl="arrowflight",
                                                     chroma_cache_dir="./foo",

docker-compose.yml

@@ -9,15 +9,16 @@ services:
     build:
       context: .
       dockerfile: Dockerfile
-      target: chroma_server
     volumes:
       - ./:/chroma
       - index_data:/index_data
     command: uvicorn chroma.app:app --reload --workers 1 --host 0.0.0.0 --port 8000
     environment:
+      - CHROMA_DB_IMPL=clickhouse
+      - CLICKHOUSE_HOST=clickhouse
+      - CLICKHOUSE_PORT=9000
       - CELERY_BROKER_URL=redis://redis:6379/0
       - CELERY_RESULT_BACKEND=redis://redis:6379/0
-      - CLICKHOUSE_TCP_PORT=9000
     ports:
       - 8000:8000
     depends_on: