fix: Qdrant module (#463)

This PR adds a module to spawn a [Qdrant](https://qdrant.tech) test
container.

---------

Co-authored-by: David Ankin <[email protected]>

Author: Anush008

conf.py

@@ -31,6 +31,7 @@
 extensions = [
     "sphinx.ext.autodoc",
     "sphinx.ext.doctest",
+    "sphinx.ext.intersphinx",
     "sphinx.ext.napoleon",
 ]
 
@@ -156,3 +157,7 @@
         "Miscellaneous",
     ),
 ]
+
+intersphinx_mapping = {
+    "python": ("https://docs.python.org/3", None),
+}

index.rst

@@ -36,6 +36,7 @@ testcontainers-python facilitates the use of Docker containers for functional an
     modules/opensearch/README
     modules/oracle/README
     modules/postgres/README
+    modules/qdrant/README
     modules/rabbitmq/README
     modules/redis/README
     modules/selenium/README

modules/qdrant/README.rst

@@ -0,0 +1,2 @@
+.. autoclass:: testcontainers.qdrant.QdrantContainer
+.. title:: testcontainers.qdrant.QdrantContainer

modules/qdrant/testcontainers/qdrant/__init__.py

@@ -0,0 +1,156 @@
+#
+#    Licensed under the Apache License, Version 2.0 (the "License"); you may
+#    not use this file except in compliance with the License. You may obtain
+#    a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+#    License for the specific language governing permissions and limitations
+#    under the License.
+import os
+from functools import cached_property
+from pathlib import Path
+from typing import Optional
+
+from testcontainers.core.config import TIMEOUT
+from testcontainers.core.generic import DbContainer
+from testcontainers.core.waiting_utils import wait_container_is_ready, wait_for_logs
+
+
+class QdrantContainer(DbContainer):
+    """
+    Qdrant vector database container.
+
+    Example:
+        .. doctest::
+
+            >>> from testcontainers.qdrant import QdrantContainer
+
+            >>> with QdrantContainer() as qdrant:
+            ...     client = qdrant.get_client()
+            ...     client.get_collections()
+            CollectionsResponse(collections=[])
+    """
+
+    QDRANT_CONFIG_FILE_PATH = "/qdrant/config/config.yaml"
+
+    def __init__(
+        self,
+        image: str = "qdrant/qdrant:v1.8.3",
+        rest_port: int = 6333,
+        grpc_port: int = 6334,
+        api_key: Optional[str] = None,
+        config_file_path: Optional[Path] = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(image, **kwargs)
+        self._rest_port = rest_port
+        self._grpc_port = grpc_port
+        self._api_key = api_key or os.getenv("QDRANT_CONTAINER_API_KEY")
+
+        if config_file_path:
+            self.with_volume_mapping(host=str(config_file_path), container=QdrantContainer.QDRANT_CONFIG_FILE_PATH)
+
+        self.with_exposed_ports(self._rest_port, self._grpc_port)
+
+    def _configure(self) -> None:
+        self.with_env("QDRANT__SERVICE__API_KEY", self._api_key)
+
+    @wait_container_is_ready()
+    def _connect(self) -> None:
+        wait_for_logs(self, ".*Actix runtime found; starting in Actix runtime.*", TIMEOUT)
+
+    def get_client(self, **kwargs):
+        """
+        Get a `qdrant_client.QdrantClient` instance associated with the container.
+
+        Args:
+            **kwargs: Additional keyword arguments to be passed to the `qdrant_client.QdrantClient` constructor.
+
+        Returns:
+            QdrantClient: An instance of the `qdrant_client.QdrantClient` class.
+
+        """
+
+        try:
+            from qdrant_client import QdrantClient
+        except ImportError as e:
+            raise ImportError("To use the `get_client` method, you must install the `qdrant_client` package.") from e
+        return QdrantClient(
+            host=self.get_container_host_ip(),
+            port=self.get_exposed_port(self._rest_port),
+            grpc_port=self.get_exposed_port(self._grpc_port),
+            api_key=self._api_key,
+            https=False,
+            **kwargs,
+        )
+
+    def get_async_client(self, **kwargs):
+        """
+        Get a `qdrant_client.AsyncQdrantClient` instance associated with the container.
+
+        Args:
+            **kwargs: Additional keyword arguments to be passed to the `qdrant_client.AsyncQdrantClient` constructor.
+
+        Returns:
+            QdrantClient: An instance of the `qdrant_client.AsyncQdrantClient` class.
+
+        """
+
+        try:
+            from qdrant_client import AsyncQdrantClient
+        except ImportError as e:
+            raise ImportError(
+                "To use the `get_async_client` method, you must install the `qdrant_client` package."
+            ) from e
+        return AsyncQdrantClient(
+            host=self.get_container_host_ip(),
+            port=self.get_exposed_port(self._rest_port),
+            grpc_port=self.get_exposed_port(self._grpc_port),
+            api_key=self._api_key,
+            https=False,
+            **kwargs,
+        )
+
+    @cached_property
+    def rest_host_address(self) -> str:
+        """
+        Get the REST host address of the Qdrant container.
+
+        Returns:
+            str: The REST host address of the Qdrant container.
+        """
+        return f"{self.get_container_host_ip()}:{self.exposed_rest_port}"
+
+    @cached_property
+    def grpc_host_address(self) -> str:
+        """
+        Get the GRPC host address of the Qdrant container.
+
+        Returns:
+            str: The GRPC host address of the Qdrant container.
+        """
+        return f"{self.get_container_host_ip()}:{self.exposed_grpc_port}"
+
+    @cached_property
+    def exposed_rest_port(self) -> int:
+        """
+        Get the exposed REST port of the Qdrant container.
+
+        Returns:
+            int: The REST port of the Qdrant container.
+        """
+        return self.get_exposed_port(self._rest_port)
+
+    @cached_property
+    def exposed_grpc_port(self) -> int:
+        """
+        Get the exposed GRPC port of the Qdrant container.
+
+        Returns:
+            int: The GRPC port of the Qdrant container.
+        """
+        return self.get_exposed_port(self._grpc_port)

modules/qdrant/tests/test_config.yaml

@@ -0,0 +1,6 @@
+# Qdrant image configuration file for testing
+# Reference: https://qdrant.tech/documentation/guides/configuration/#configuration-file-example
+log_level: INFO
+
+service:
+  api_key: "SOME_TEST_KEY"

modules/qdrant/tests/test_qdrant.py

@@ -0,0 +1,80 @@
+import pytest
+from testcontainers.qdrant import QdrantContainer
+import uuid
+from grpc import RpcError
+from pathlib import Path
+
+import qdrant_client
+
+
+def test_docker_run_qdrant():
+    with QdrantContainer() as qdrant:
+        client = qdrant.get_client()
+        collections = client.get_collections().collections
+        assert len(collections) == 0
+
+        client = qdrant.get_client(prefer_grpc=True)
+        collections = client.get_collections().collections
+        assert len(collections) == 0
+
+
+def test_qdrant_with_api_key_http():
+    api_key = uuid.uuid4().hex
+
+    with QdrantContainer(api_key=api_key) as qdrant:
+        with pytest.raises(qdrant_client.http.exceptions.UnexpectedResponse) as e:
+            # Construct a client without an API key
+            qdrant_client.QdrantClient(location=f"http://{qdrant.rest_host_address}").get_collections()
+
+        assert "Must provide an API key" in str(e.value)
+
+        # Construct a client with an API key
+        collections = (
+            qdrant_client.QdrantClient(location=f"http://{qdrant.rest_host_address}", api_key=api_key)
+            .get_collections()
+            .collections
+        )
+
+        assert len(collections) == 0
+
+        # Get an automatically configured client instance
+        collections = qdrant.get_client().get_collections().collections
+
+        assert len(collections) == 0
+
+
+def test_qdrant_with_api_key_grpc():
+    api_key = uuid.uuid4().hex
+
+    with QdrantContainer(api_key=api_key) as qdrant:
+        with pytest.raises(RpcError) as e:
+            qdrant_client.QdrantClient(
+                url=f"http://{qdrant.grpc_host_address}",
+                grpc_port=qdrant.exposed_grpc_port,
+                prefer_grpc=True,
+            ).get_collections()
+
+        assert "Must provide an API key" in str(e.value)
+
+        collections = (
+            qdrant_client.QdrantClient(
+                url=f"http://{qdrant.grpc_host_address}",
+                grpc_port=qdrant.exposed_grpc_port,
+                prefer_grpc=True,
+                api_key=api_key,
+            )
+            .get_collections()
+            .collections
+        )
+
+        assert len(collections) == 0
+
+
+def test_qdrant_with_config_file():
+    config_file_path = Path(__file__).with_name("test_config.yaml")
+
+    with QdrantContainer(config_file_path=config_file_path) as qdrant:
+        with pytest.raises(qdrant_client.http.exceptions.UnexpectedResponse) as e:
+            qdrant_client.QdrantClient(location=f"http://{qdrant.rest_host_address}").get_collections()
+
+        assert "Must provide an API key" in str(e.value)