import multiprocessing import os import socket import subprocess import tempfile import time from typing import ( Any, Generator, Iterator, List, Optional, Sequence, Tuple, Callable, cast, ) from uuid import UUID import hypothesis import pytest import uvicorn from httpx import ConnectError from typing_extensions import Protocol from chromadb.api.async_fastapi import AsyncFastAPI from chromadb.api.fastapi import FastAPI import chromadb.server.fastapi from chromadb.api import ClientAPI, ServerAPI, BaseAPI from chromadb.config import Settings, System from chromadb.db.mixins import embeddings_queue from chromadb.ingest import Producer from chromadb.types import SeqId, OperationRecord from chromadb.api.client import Client as ClientCreator, AdminClient from chromadb.api.async_client import ( AsyncAdminClient, AsyncClient as AsyncClientCreator, ) from chromadb.utils.async_to_sync import async_class_to_sync import logging logger = logging.getLogger(__name__) VALID_PRESETS = ["fast", "normal", "slow"] CURRENT_PRESET = os.getenv("PROPERTY_TESTING_PRESET", "fast") if CURRENT_PRESET not in VALID_PRESETS: raise ValueError( f"Invalid property testing preset: {CURRENT_PRESET}. Must be one of {VALID_PRESETS}." ) hypothesis.settings.register_profile( "base", deadline=45000, suppress_health_check=[ hypothesis.HealthCheck.data_too_large, hypothesis.HealthCheck.large_base_example, hypothesis.HealthCheck.function_scoped_fixture, ], ) hypothesis.settings.register_profile( "fast", hypothesis.settings.get_profile("base"), max_examples=50 ) # Hypothesis's default max_examples is 100 hypothesis.settings.register_profile( "normal", hypothesis.settings.get_profile("base"), max_examples=100 ) hypothesis.settings.register_profile( "slow", hypothesis.settings.get_profile("base"), max_examples=500, stateful_step_count=100, ) hypothesis.settings.load_profile(CURRENT_PRESET) def reset(api: BaseAPI) -> None: api.reset() def override_hypothesis_profile( fast: Optional[hypothesis.settings] = None, normal: Optional[hypothesis.settings] = None, slow: Optional[hypothesis.settings] = None, ) -> Optional[hypothesis.settings]: """Override Hypothesis settings for specific profiles. For example, to override max_examples only when the current profile is 'fast': override_hypothesis_profile( fast=hypothesis.settings(max_examples=50), ) Settings will be merged with the default/active profile. """ allowable_override_keys = [ "deadline", "max_examples", "stateful_step_count", "suppress_health_check", ] override_profiles = { "fast": fast, "normal": normal, "slow": slow, } overriding_profile = override_profiles.get(CURRENT_PRESET) if overriding_profile is not None: overridden_settings = { key: value for key, value in overriding_profile.__dict__.items() if key in allowable_override_keys } return hypothesis.settings(hypothesis.settings.default, **overridden_settings) return cast(hypothesis.settings, hypothesis.settings.default) NOT_CLUSTER_ONLY = os.getenv("CHROMA_CLUSTER_TEST_ONLY") != "1" COMPACTION_SLEEP = 120 def skip_if_not_cluster() -> pytest.MarkDecorator: return pytest.mark.skipif( NOT_CLUSTER_ONLY, reason="Requires Kubernetes to be running with a valid config", ) def generate_self_signed_certificate() -> None: config_path = os.path.join( os.path.dirname(os.path.abspath(__file__)), "openssl.cnf" ) print(f"Config path: {config_path}") # Debug print to verify path if not os.path.exists(config_path): raise FileNotFoundError(f"Config file not found at {config_path}") subprocess.run( [ "openssl", "req", "-x509", "-newkey", "rsa:4096", "-keyout", "serverkey.pem", "-out", "servercert.pem", "-days", "365", "-nodes", "-subj", "/CN=localhost", "-config", config_path, ] ) def find_free_port() -> int: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: s.bind(("", 0)) s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) return s.getsockname()[1] # type: ignore def _run_server( port: int, is_persistent: bool = False, persist_directory: Optional[str] = None, chroma_server_authn_provider: Optional[str] = None, chroma_server_authn_credentials_file: Optional[str] = None, chroma_server_authn_credentials: Optional[str] = None, chroma_auth_token_transport_header: Optional[str] = None, chroma_server_authz_provider: Optional[str] = None, chroma_server_authz_config_file: Optional[str] = None, chroma_server_ssl_certfile: Optional[str] = None, chroma_server_ssl_keyfile: Optional[str] = None, chroma_overwrite_singleton_tenant_database_access_from_auth: Optional[bool] = False, ) -> None: """Run a Chroma server locally""" if is_persistent and persist_directory: settings = Settings( chroma_api_impl="chromadb.api.segment.SegmentAPI", chroma_sysdb_impl="chromadb.db.impl.sqlite.SqliteDB", chroma_producer_impl="chromadb.db.impl.sqlite.SqliteDB", chroma_consumer_impl="chromadb.db.impl.sqlite.SqliteDB", chroma_segment_manager_impl="chromadb.segment.impl.manager.local.LocalSegmentManager", is_persistent=is_persistent, persist_directory=persist_directory, allow_reset=True, chroma_server_authn_provider=chroma_server_authn_provider, chroma_server_authn_credentials_file=chroma_server_authn_credentials_file, chroma_server_authn_credentials=chroma_server_authn_credentials, chroma_auth_token_transport_header=chroma_auth_token_transport_header, chroma_server_authz_provider=chroma_server_authz_provider, chroma_server_authz_config_file=chroma_server_authz_config_file, chroma_overwrite_singleton_tenant_database_access_from_auth=chroma_overwrite_singleton_tenant_database_access_from_auth, ) else: settings = Settings( chroma_api_impl="chromadb.api.segment.SegmentAPI", chroma_sysdb_impl="chromadb.db.impl.sqlite.SqliteDB", chroma_producer_impl="chromadb.db.impl.sqlite.SqliteDB", chroma_consumer_impl="chromadb.db.impl.sqlite.SqliteDB", chroma_segment_manager_impl="chromadb.segment.impl.manager.local.LocalSegmentManager", is_persistent=False, allow_reset=True, chroma_server_authn_provider=chroma_server_authn_provider, chroma_server_authn_credentials_file=chroma_server_authn_credentials_file, chroma_server_authn_credentials=chroma_server_authn_credentials, chroma_auth_token_transport_header=chroma_auth_token_transport_header, chroma_server_authz_provider=chroma_server_authz_provider, chroma_server_authz_config_file=chroma_server_authz_config_file, chroma_overwrite_singleton_tenant_database_access_from_auth=chroma_overwrite_singleton_tenant_database_access_from_auth, ) server = chromadb.server.fastapi.FastAPI(settings) uvicorn.run( server.app(), host="0.0.0.0", port=port, log_level="error", timeout_keep_alive=30, ssl_keyfile=chroma_server_ssl_keyfile, ssl_certfile=chroma_server_ssl_certfile, ) def _await_server(api: ServerAPI, attempts: int = 0) -> None: try: api.heartbeat() except ConnectError as e: if attempts > 15: raise e else: time.sleep(4) _await_server(api, attempts + 1) def _fastapi_fixture( is_persistent: bool = False, chroma_api_impl: str = "chromadb.api.fastapi.FastAPI", chroma_server_authn_provider: Optional[str] = None, chroma_client_auth_provider: Optional[str] = None, chroma_server_authn_credentials_file: Optional[str] = None, chroma_server_authn_credentials: Optional[str] = None, chroma_client_auth_credentials: Optional[str] = None, chroma_auth_token_transport_header: Optional[str] = None, chroma_server_authz_provider: Optional[str] = None, chroma_server_authz_config_file: Optional[str] = None, chroma_server_ssl_certfile: Optional[str] = None, chroma_server_ssl_keyfile: Optional[str] = None, chroma_overwrite_singleton_tenant_database_access_from_auth: Optional[bool] = False, ) -> Generator[System, None, None]: """Fixture generator that launches a server in a separate process, and yields a fastapi client connect to it""" port = find_free_port() ctx = multiprocessing.get_context("spawn") args: Tuple[ int, bool, Optional[str], Optional[str], Optional[str], Optional[str], Optional[str], Optional[str], Optional[str], Optional[str], Optional[str], Optional[bool], ] = ( port, False, None, chroma_server_authn_provider, chroma_server_authn_credentials_file, chroma_server_authn_credentials, chroma_auth_token_transport_header, chroma_server_authz_provider, chroma_server_authz_config_file, chroma_server_ssl_certfile, chroma_server_ssl_keyfile, chroma_overwrite_singleton_tenant_database_access_from_auth, ) def run(args: Any) -> Generator[System, None, None]: proc = ctx.Process(target=_run_server, args=args, daemon=True) proc.start() settings = Settings( chroma_api_impl=chroma_api_impl, chroma_server_host="localhost", chroma_server_http_port=port, allow_reset=True, chroma_client_auth_provider=chroma_client_auth_provider, chroma_client_auth_credentials=chroma_client_auth_credentials, chroma_auth_token_transport_header=chroma_auth_token_transport_header, chroma_server_ssl_verify=chroma_server_ssl_certfile, chroma_server_ssl_enabled=True if chroma_server_ssl_certfile else False, chroma_overwrite_singleton_tenant_database_access_from_auth=chroma_overwrite_singleton_tenant_database_access_from_auth, ) system = System(settings) api = system.instance(ServerAPI) system.start() _await_server(api if isinstance(api, FastAPI) else async_class_to_sync(api)) yield system system.stop() proc.kill() proc.join() if is_persistent: persist_directory = tempfile.TemporaryDirectory() args = ( port, is_persistent, persist_directory.name, chroma_server_authn_provider, chroma_server_authn_credentials_file, chroma_server_authn_credentials, chroma_auth_token_transport_header, chroma_server_authz_provider, chroma_server_authz_config_file, chroma_server_ssl_certfile, chroma_server_ssl_keyfile, chroma_overwrite_singleton_tenant_database_access_from_auth, ) yield from run(args) try: persist_directory.cleanup() # (Older versions of Python throw NotADirectoryError sometimes instead of PermissionError) # (when we drop support for Python < 3.10, we should use ignore_cleanup_errors=True with the context manager instead) except (PermissionError, NotADirectoryError) as e: # todo: what's holding onto directory contents on Windows? if os.name == "nt": pass else: raise e else: yield from run(args) def fastapi() -> Generator[System, None, None]: return _fastapi_fixture(is_persistent=False) def async_fastapi() -> Generator[System, None, None]: return _fastapi_fixture( is_persistent=False, chroma_api_impl="chromadb.api.async_fastapi.AsyncFastAPI", ) def fastapi_persistent() -> Generator[System, None, None]: return _fastapi_fixture(is_persistent=True) def fastapi_ssl() -> Generator[System, None, None]: generate_self_signed_certificate() return _fastapi_fixture( is_persistent=False, chroma_server_ssl_certfile="./servercert.pem", chroma_server_ssl_keyfile="./serverkey.pem", ) def basic_http_client() -> Generator[System, None, None]: settings = Settings( chroma_api_impl="chromadb.api.fastapi.FastAPI", chroma_server_http_port=8000, chroma_server_host="localhost", allow_reset=True, ) system = System(settings) api = system.instance(ServerAPI) _await_server(api) system.start() yield system system.stop() def fastapi_server_basic_auth_valid_cred_single_user() -> Generator[System, None, None]: # This (and similar usage below) should use the delete_on_close parameter # instead of delete=False, but it's only available in Python 3.12 and later. # We must explicitly close the file before spawning a subprocess to avoid # file locking issues on Windows. with tempfile.NamedTemporaryFile("w", suffix=".htpasswd", delete=False) as f: f.write("admin:$2y$05$e5sRb6NCcSH3YfbIxe1AGu2h5K7OOd982OXKmd8WyQ3DRQ4MvpnZS\n") f.close() for item in _fastapi_fixture( is_persistent=False, chroma_server_authn_provider="chromadb.auth.basic_authn.BasicAuthenticationServerProvider", chroma_server_authn_credentials_file=f.name, chroma_client_auth_provider="chromadb.auth.basic_authn.BasicAuthClientProvider", chroma_client_auth_credentials="admin:admin", ): yield item def fastapi_server_basic_auth_valid_cred_multiple_users() -> ( Generator[System, None, None] ): creds = { "user": "$2y$10$kY9hn.Wlfcj7n1Cnjmy1kuIhEFIVBsfbNWLQ5ahoKmdc2HLA4oP6i", "user2": "$2y$10$CymQ63tic/DRj8dD82915eoM4ke3d6RaNKU4dj4IVJlHyea0yeGDS", "admin": "$2y$05$e5sRb6NCcSH3YfbIxe1AGu2h5K7OOd982OXKmd8WyQ3DRQ4MvpnZS", } with tempfile.NamedTemporaryFile("w", suffix=".htpasswd", delete=False) as f: for user, cred in creds.items(): f.write(f"{user}:{cred}\n") f.close() for item in _fastapi_fixture( is_persistent=False, chroma_server_authn_provider="chromadb.auth.basic_authn.BasicAuthenticationServerProvider", chroma_server_authn_credentials_file=f.name, chroma_client_auth_provider="chromadb.auth.basic_authn.BasicAuthClientProvider", chroma_client_auth_credentials="admin:admin", ): yield item def fastapi_server_basic_auth_invalid_cred() -> Generator[System, None, None]: with tempfile.NamedTemporaryFile("w", suffix=".htpasswd", delete=False) as f: f.write("admin:$2y$05$e5sRb6NCcSH3YfbIxe1AGu2h5K7OOd982OXKmd8WyQ3DRQ4MvpnZS\n") f.close() for item in _fastapi_fixture( is_persistent=False, chroma_server_authn_provider="chromadb.auth.basic_authn.BasicAuthenticationServerProvider", chroma_server_authn_credentials_file=f.name, chroma_client_auth_provider="chromadb.auth.basic_authn.BasicAuthClientProvider", chroma_client_auth_credentials="admin:admin1", ): yield item def fastapi_server_basic_authn_rbac_authz() -> Generator[System, None, None]: with tempfile.NamedTemporaryFile( "w", suffix=".htpasswd", delete=False ) as server_authn_file: server_authn_file.write( "admin:$2y$05$e5sRb6NCcSH3YfbIxe1AGu2h5K7OOd982OXKmd8WyQ3DRQ4MvpnZS\n" ) server_authn_file.close() with tempfile.NamedTemporaryFile( "w", suffix=".authz", delete=False ) as server_authz_file: server_authz_file.write( """ roles_mapping: admin: actions: [ "system:reset", "tenant:create_tenant", "tenant:get_tenant", "db:create_database", "db:get_database", "db:list_collections", "db:create_collection", "db:get_or_create_collection", "collection:get_collection", "collection:delete_collection", "collection:update_collection", "collection:add", "collection:delete", "collection:get", "collection:query", "collection:peek", "collection:update", "collection:upsert", "collection:count", ] users: - id: admin role: admin """ ) server_authz_file.close() for item in _fastapi_fixture( is_persistent=False, chroma_client_auth_provider="chromadb.auth.basic_authn.BasicAuthClientProvider", chroma_client_auth_credentials="admin:admin", chroma_server_authn_provider="chromadb.auth.basic_authn.BasicAuthenticationServerProvider", chroma_server_authn_credentials_file=server_authn_file.name, chroma_server_authz_provider="chromadb.auth.simple_rbac_authz.SimpleRBACAuthorizationProvider", chroma_server_authz_config_file=server_authz_file.name, ): yield item def fastapi_fixture_admin_and_singleton_tenant_db_user() -> ( Generator[System, None, None] ): with tempfile.NamedTemporaryFile("w", suffix=".authn", delete=False) as f: f.write( """ users: - id: admin tokens: - admin-token - id: singleton_user tenant: singleton_tenant databases: - singleton_database tokens: - singleton-token """ ) f.close() for item in _fastapi_fixture( is_persistent=False, chroma_overwrite_singleton_tenant_database_access_from_auth=True, chroma_client_auth_provider="chromadb.auth.token_authn.TokenAuthClientProvider", chroma_client_auth_credentials="admin-token", chroma_server_authn_provider="chromadb.auth.token_authn.TokenAuthenticationServerProvider", chroma_server_authn_credentials_file=f.name, ): yield item def integration() -> Generator[System, None, None]: """Fixture generator for returning a client configured via environmenet variables, intended for externally configured integration tests """ settings = Settings(allow_reset=True) system = System(settings) system.start() yield system system.stop() def sqlite_fixture() -> Generator[System, None, None]: """Fixture generator for segment-based API using in-memory Sqlite""" settings = Settings( chroma_api_impl="chromadb.api.segment.SegmentAPI", chroma_sysdb_impl="chromadb.db.impl.sqlite.SqliteDB", chroma_producer_impl="chromadb.db.impl.sqlite.SqliteDB", chroma_consumer_impl="chromadb.db.impl.sqlite.SqliteDB", chroma_segment_manager_impl="chromadb.segment.impl.manager.local.LocalSegmentManager", is_persistent=False, allow_reset=True, ) system = System(settings) system.start() yield system system.stop() @pytest.fixture def sqlite() -> Generator[System, None, None]: yield from sqlite_fixture() def sqlite_persistent_fixture() -> Generator[System, None, None]: """Fixture generator for segment-based API using persistent Sqlite""" save_path = tempfile.TemporaryDirectory() settings = Settings( chroma_api_impl="chromadb.api.segment.SegmentAPI", chroma_sysdb_impl="chromadb.db.impl.sqlite.SqliteDB", chroma_producer_impl="chromadb.db.impl.sqlite.SqliteDB", chroma_consumer_impl="chromadb.db.impl.sqlite.SqliteDB", chroma_segment_manager_impl="chromadb.segment.impl.manager.local.LocalSegmentManager", allow_reset=True, is_persistent=True, persist_directory=save_path.name, ) system = System(settings) system.start() yield system system.stop() try: save_path.cleanup() # (Older versions of Python throw NotADirectoryError sometimes instead of PermissionError) # (when we drop support for Python < 3.10, we should use ignore_cleanup_errors=True with the context manager instead) except (PermissionError, NotADirectoryError) as e: # todo: what's holding onto directory contents on Windows? if os.name == "nt": pass else: raise e @pytest.fixture def sqlite_persistent() -> Generator[System, None, None]: yield from sqlite_persistent_fixture() def system_fixtures() -> List[Callable[[], Generator[System, None, None]]]: fixtures = [ fastapi, async_fastapi, fastapi_persistent, sqlite_fixture, sqlite_persistent_fixture, ] if "CHROMA_INTEGRATION_TEST" in os.environ: fixtures.append(integration) if "CHROMA_INTEGRATION_TEST_ONLY" in os.environ: fixtures = [integration] if "CHROMA_CLUSTER_TEST_ONLY" in os.environ: fixtures = [basic_http_client] return fixtures def system_http_server_fixtures() -> List[Callable[[], Generator[System, None, None]]]: fixtures = [ fixture for fixture in system_fixtures() if fixture != sqlite_fixture and fixture != sqlite_persistent_fixture ] return fixtures def system_fixtures_auth() -> List[Callable[[], Generator[System, None, None]]]: fixtures = [ fastapi_server_basic_auth_valid_cred_single_user, fastapi_server_basic_auth_valid_cred_multiple_users, ] return fixtures def system_fixtures_authn_rbac_authz() -> ( List[Callable[[], Generator[System, None, None]]] ): fixtures = [fastapi_server_basic_authn_rbac_authz] return fixtures def system_fixtures_root_and_singleton_tenant_db_user() -> ( List[Callable[[], Generator[System, None, None]]] ): fixtures = [fastapi_fixture_admin_and_singleton_tenant_db_user] return fixtures def system_fixtures_wrong_auth() -> List[Callable[[], Generator[System, None, None]]]: fixtures = [fastapi_server_basic_auth_invalid_cred] return fixtures def system_fixtures_ssl() -> List[Callable[[], Generator[System, None, None]]]: fixtures = [fastapi_ssl] return fixtures @pytest.fixture(scope="module", params=system_fixtures_wrong_auth()) def system_wrong_auth( request: pytest.FixtureRequest, ) -> Generator[ServerAPI, None, None]: yield from request.param() @pytest.fixture(scope="module", params=system_fixtures_authn_rbac_authz()) def system_authn_rbac_authz( request: pytest.FixtureRequest, ) -> Generator[ServerAPI, None, None]: yield from request.param() @pytest.fixture(scope="module", params=system_http_server_fixtures()) def system_http_server( request: pytest.FixtureRequest, ) -> Generator[ServerAPI, None, None]: yield from request.param() @pytest.fixture(scope="module", params=system_fixtures()) def system(request: pytest.FixtureRequest) -> Generator[ServerAPI, None, None]: yield from request.param() @pytest.fixture(scope="module", params=system_fixtures_ssl()) def system_ssl(request: pytest.FixtureRequest) -> Generator[ServerAPI, None, None]: yield from request.param() @pytest.fixture(scope="module", params=system_fixtures_auth()) def system_auth(request: pytest.FixtureRequest) -> Generator[ServerAPI, None, None]: yield from request.param() @async_class_to_sync class AsyncClientCreatorSync(AsyncClientCreator): pass @async_class_to_sync class AsyncAdminClientSync(AsyncAdminClient): pass @pytest.fixture(scope="function") def api(system: System) -> Generator[ServerAPI, None, None]: system.reset_state() api = system.instance(ServerAPI) if isinstance(api, AsyncFastAPI): transformed = async_class_to_sync(api) yield transformed else: yield api class ClientFactories: """This allows consuming tests to be parameterized by async/sync versions of the client and papers over the async implementation. If you don't need to manually construct clients, use the `client` fixture instead. """ _system: System # Need to track created clients so we can call .clear_system_cache() during teardown _created_clients: List[ClientAPI] = [] def __init__(self, system: System): self._system = system def create_client(self, *args: Any, **kwargs: Any) -> ClientCreator: if kwargs.get("settings") is None: kwargs["settings"] = self._system.settings if ( self._system.settings.chroma_api_impl == "chromadb.api.async_fastapi.AsyncFastAPI" ): client = cast(ClientCreator, AsyncClientCreatorSync.create(*args, **kwargs)) self._created_clients.append(client) return client client = ClientCreator(*args, **kwargs) self._created_clients.append(client) return client def create_admin_client(self, *args: Any, **kwargs: Any) -> AdminClient: if ( self._system.settings.chroma_api_impl == "chromadb.api.async_fastapi.AsyncFastAPI" ): return cast(AdminClient, AsyncAdminClientSync(*args, **kwargs)) return AdminClient(*args, **kwargs) def create_admin_client_from_system(self) -> AdminClient: if ( self._system.settings.chroma_api_impl == "chromadb.api.async_fastapi.AsyncFastAPI" ): return cast(AdminClient, AsyncAdminClientSync.from_system(self._system)) return AdminClient.from_system(self._system) @pytest.fixture(scope="function") def client_factories(system: System) -> Generator[ClientFactories, None, None]: system.reset_state() factories = ClientFactories(system) yield factories while len(factories._created_clients) > 0: client = factories._created_clients.pop() client.clear_system_cache() del client @pytest.fixture(scope="function") def client(system: System) -> Generator[ClientAPI, None, None]: system.reset_state() if system.settings.chroma_api_impl == "chromadb.api.async_fastapi.AsyncFastAPI": client = cast(Any, AsyncClientCreatorSync.from_system_async(system)) yield client client.clear_system_cache() else: client = ClientCreator.from_system(system) yield client client.clear_system_cache() @pytest.fixture(scope="function") def http_client(system_http_server: System) -> Generator[ClientAPI, None, None]: system_http_server.reset_state() if ( system_http_server.settings.chroma_api_impl == "chromadb.api.async_fastapi.AsyncFastAPI" ): client = cast(Any, AsyncClientCreatorSync.from_system_async(system_http_server)) yield client client.clear_system_cache() else: client = ClientCreator.from_system(system_http_server) yield client client.clear_system_cache() @pytest.fixture(scope="function") def client_ssl(system_ssl: System) -> Generator[ClientAPI, None, None]: system_ssl.reset_state() client = ClientCreator.from_system(system_ssl) yield client client.clear_system_cache() @pytest.fixture(scope="function") def api_wrong_cred( system_wrong_auth: System, ) -> Generator[ServerAPI, None, None]: system_wrong_auth.reset_state() api = system_wrong_auth.instance(ServerAPI) yield api @pytest.fixture(scope="function") def api_with_authn_rbac_authz( system_authn_rbac_authz: System, ) -> Generator[ServerAPI, None, None]: system_authn_rbac_authz.reset_state() api = system_authn_rbac_authz.instance(ServerAPI) yield api @pytest.fixture(scope="function") def api_with_server_auth(system_auth: System) -> Generator[ServerAPI, None, None]: _sys = system_auth _sys.reset_state() api = _sys.instance(ServerAPI) yield api # Producer / Consumer fixtures # class ProducerFn(Protocol): def __call__( self, producer: Producer, collection_id: UUID, embeddings: Iterator[OperationRecord], n: int, ) -> Tuple[Sequence[OperationRecord], Sequence[SeqId]]: ... def produce_n_single( producer: Producer, collection_id: UUID, embeddings: Iterator[OperationRecord], n: int, ) -> Tuple[Sequence[OperationRecord], Sequence[SeqId]]: submitted_embeddings = [] seq_ids = [] for _ in range(n): e = next(embeddings) seq_id = producer.submit_embedding(collection_id, e) submitted_embeddings.append(e) seq_ids.append(seq_id) return submitted_embeddings, seq_ids def produce_n_batch( producer: Producer, collection_id: UUID, embeddings: Iterator[OperationRecord], n: int, ) -> Tuple[Sequence[OperationRecord], Sequence[SeqId]]: submitted_embeddings = [] seq_ids: Sequence[SeqId] = [] for _ in range(n): e = next(embeddings) submitted_embeddings.append(e) seq_ids = producer.submit_embeddings(collection_id, submitted_embeddings) return submitted_embeddings, seq_ids def produce_fn_fixtures() -> List[ProducerFn]: return [produce_n_single, produce_n_batch] @pytest.fixture(scope="module", params=produce_fn_fixtures()) def produce_fns( request: pytest.FixtureRequest, ) -> Generator[ProducerFn, None, None]: yield request.param def pytest_configure(config): # type: ignore embeddings_queue._called_from_test = True def is_client_in_process(client: ClientAPI) -> bool: """Returns True if the client is in-process (a SQLite client), False if it's out-of-process (a HTTP client).""" return client.get_settings().chroma_server_http_port is None @pytest.fixture(autouse=True) def log_tests(request): """Automatically logs the start and end of each test.""" test_name = request.node.name logger.debug(f"Starting test: {test_name}") # Yield control back to the test, allowing it to execute yield logger.debug(f"Finished test: {test_name}")