FazBrowse GitHub Viewer | Trending |
URL:
| Home
Tools: [Download Repo ZIP]   [Original HTTPS Page]

fix: add chroma container by alexanderankin · Pull Request #515 · testcontainers/testcontainers-python · GitHub

Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension .lock  (1) .py  (2) .rst  (2) .toml  (1) All 4 file types selected
Only manifest files
Viewed files
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Unified
Split
Hide whitespace
Diff view
Unified
Split
Hide whitespace
1 change: 1 addition & 0 deletions index.rst
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 2 additions & 0 deletions modules/chroma/README.rst
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
.. autoclass:: testcontainers.chroma.ChromaContainer
.. title:: testcontainers.minio.ChromaContainer
81 changes: 81 additions & 0 deletions modules/chroma/testcontainers/chroma/__init__.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Original file line number Diff line number Diff line change
@@ -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
9 changes: 9 additions & 0 deletions modules/chroma/tests/test_chroma.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Original file line number Diff line number Diff line change
@@ -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"
Loading

Back | FazBrowse Home | New Git URL