diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..9f25234 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,3 @@ +venv +.git +examples \ No newline at end of file diff --git a/.gitignore b/.gitignore index f69bafd..c2406ac 100644 --- a/.gitignore +++ b/.gitignore @@ -15,4 +15,5 @@ index_data /index_data venv -.env \ No newline at end of file +.env +.chroma diff --git a/Dockerfile b/Dockerfile index e0aab6d..006f1fe 100644 --- a/Dockerfile +++ b/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"] diff --git a/bin/integration-test b/bin/integration-test new file mode 100755 index 0000000..ded5854 --- /dev/null +++ b/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 + diff --git a/bin/test b/bin/test deleted file mode 100755 index 793da14..0000000 --- a/bin/test +++ /dev/null @@ -1,11 +0,0 @@ -#!/usr/bin/env bash - -set -e - -function cleanup { - docker-compose -f docker-compose.test.yml down -} - -trap cleanup EXIT - -docker-compose -f docker-compose.test.yml run --rm server_test diff --git a/chroma/app.py b/chroma/app.py index 80909f3..f064aae 100644 --- a/chroma/app.py +++ b/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() diff --git a/chroma/test/test_api.py b/chroma/test/test_api.py index 68d98aa..0a4064e 100644 --- a/chroma/test/test_api.py +++ b/chroma/test/test_api.py @@ -5,6 +5,7 @@ import pytest 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): diff --git a/chroma/test/test_chroma.py b/chroma/test/test_chroma.py index 064443b..e7d8e4d 100644 --- a/chroma/test/test_chroma.py +++ b/chroma/test/test_chroma.py @@ -1,5 +1,6 @@ import pytest import unittest +import os from unittest.mock import patch import chroma @@ -32,6 +33,7 @@ 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 @@ -39,6 +41,7 @@ class GetAPITest(unittest.TestCase): @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 @@ class GetAPITest(unittest.TestCase): @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 @@ class GetAPITest(unittest.TestCase): @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", diff --git a/docker-compose.yml b/docker-compose.yml index 704ca82..1ebdac1 100644 --- a/docker-compose.yml +++ b/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: