feat: Qdrant module

INDEX.rst

@@ -32,6 +32,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,68 @@
+#
+#    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 pathlib import Path
+from typing import Optional
+
+from qdrant_client import QdrantClient
+
+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_CONFIG_FILE_PATH = "/qdrant/config/config.yaml"
+
+    def __init__(
+        self,
+        image: str = "qdrant/qdrant:latest",
+        rest_port: int = 6333,
+        grpc_port: int = 6334,
+        container_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.container_api_key = container_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.container_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) -> "QdrantClient":
+        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.container_api_key,
+            **kwargs,
+        )
+
+    @property
+    def http_host_address(self) -> str:
+        return f"{self.get_container_host_ip()}:{self.get_exposed_port(self.rest_port)}"
+
+    @property
+    def grpc_host_address(self) -> str:
+        return f"{self.get_container_host_ip()}:{self.get_exposed_port(self.grpc_port)}"

modules/qdrant/tests/test_config.yaml

@@ -0,0 +1,4 @@
+log_level: INFO
+
+service:
+  api_key: "SOME_TEST_KEY"

modules/qdrant/tests/test_qdrant.py

@@ -0,0 +1,71 @@
+import pytest
+from testcontainers.qdrant import QdrantContainer
+import uuid
+from qdrant_client import QdrantClient
+from qdrant_client.http.exceptions import UnexpectedResponse
+from grpc import RpcError
+from pathlib import Path
+
+
+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(container_api_key=api_key) as qdrant:
+        with pytest.raises(UnexpectedResponse) as e:
+            QdrantClient(location=f"http://{qdrant.http_host_address}").get_collections()
+
+        assert "Invalid api-key" in str(e.value)
+
+        collections = (
+            QdrantClient(location=f"http://{qdrant.http_host_address}", api_key=api_key).get_collections().collections
+        )
+
+        assert len(collections) == 0
+
+
+def test_qdrant_with_api_key_grpc():
+    api_key = uuid.uuid4().hex
+
+    with QdrantContainer(container_api_key=api_key) as qdrant:
+        with pytest.raises(RpcError) as e:
+            QdrantClient(
+                url=f"http://{qdrant.grpc_host_address}",
+                grpc_port=qdrant.get_exposed_port(qdrant.grpc_port),
+                prefer_grpc=True,
+            ).get_collections()
+
+        assert "Invalid api-key" in str(e.value)
+
+        collections = (
+            QdrantClient(
+                url=f"http://{qdrant.grpc_host_address}",
+                grpc_port=qdrant.get_exposed_port(qdrant.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(UnexpectedResponse) as e:
+            QdrantClient(location=f"http://{qdrant.http_host_address}").get_collections()
+
+        assert "Invalid api-key" in str(e.value)