fix: add chroma container (#515)

based on #477

---------

Co-authored-by: Trayan Azarov <[email protected]>

Author: alexanderankin

index.rst

@@ -17,6 +17,7 @@ testcontainers-python facilitates the use of Docker containers for functional an
     core/README
     modules/arangodb/README
     modules/azurite/README
+    modules/chroma/README
     modules/clickhouse/README
     modules/elasticsearch/README
     modules/google/README

modules/chroma/README.rst

@@ -0,0 +1,2 @@
+.. autoclass:: testcontainers.chroma.ChromaContainer
+.. title:: testcontainers.minio.ChromaContainer

modules/chroma/testcontainers/chroma/__init__.py

@@ -0,0 +1,81 @@
+from typing import TYPE_CHECKING
+
+from requests import ConnectionError, get
+
+from testcontainers.core.container import DockerContainer
+from testcontainers.core.utils import raise_for_deprecated_parameter
+from testcontainers.core.waiting_utils import wait_container_is_ready
+
+if TYPE_CHECKING:
+    from requests import Response
+
+
+class ChromaContainer(DockerContainer):
+    """
+    The example below spins up a ChromaDB container, performs a healthcheck and creates a collection.
+    The method :code:`get_client` can be used to create a client for the Chroma Python Client.
+
+    Example:
+
+        .. doctest::
+
+            >>> import chromadb
+            >>> from testcontainers.chroma import ChromaContainer
+
+            >>> with ChromaContainer() as chroma:
+            ...   config = chroma.get_config()
+            ...   client = chromadb.HttpClient(host=config["host"], port=config["port"])
+            ...   col = client.get_or_create_collection("test")
+            ...   col.name
+            'test'
+    """
+
+    def __init__(
+        self,
+        image: str = "chromadb/chroma:latest",
+        port: int = 8000,
+        **kwargs,
+    ) -> None:
+        """
+        Args:
+            image: Docker image to use for the MinIO container.
+            port: Port to expose on the container.
+            access_key: Access key for client connections.
+            secret_key: Secret key for client connections.
+        """
+        raise_for_deprecated_parameter(kwargs, "port_to_expose", "port")
+        super().__init__(image, **kwargs)
+        self.port = port
+
+        self.with_exposed_ports(self.port)
+        # self.with_command(f"server /data --address :{self.port}")
+
+    def get_config(self) -> dict:
+        """This method returns the configuration of the Chroma container,
+        including the endpoint.
+
+        Returns:
+            dict: {`endpoint`: str}
+        """
+        host_ip = self.get_container_host_ip()
+        exposed_port = self.get_exposed_port(self.port)
+        return {
+            "endpoint": f"{host_ip}:{exposed_port}",
+            "host": host_ip,
+            "port": exposed_port,
+        }
+
+    @wait_container_is_ready(ConnectionError)
+    def _healthcheck(self) -> None:
+        """This is an internal method used to check if the Chroma container
+        is healthy and ready to receive requests."""
+        url = f"http://{self.get_config()['endpoint']}/api/v1/heartbeat"
+        response: Response = get(url)
+        response.raise_for_status()
+
+    def start(self) -> "ChromaContainer":
+        """This method starts the Chroma container and runs the healthcheck
+        to verify that the container is ready to use."""
+        super().start()
+        self._healthcheck()
+        return self

modules/chroma/tests/test_chroma.py

@@ -0,0 +1,9 @@
+from testcontainers.chroma import ChromaContainer
+import chromadb
+
+
+def test_docker_run_chroma():
+    with ChromaContainer(image="chromadb/chroma:0.4.24") as chroma:
+        client = chromadb.HttpClient(host=chroma.get_config()["host"], port=chroma.get_config()["port"])
+        col = client.get_or_create_collection("test")
+        assert col.name == "test"