fix: Container for Milvus database (#606)

I use this wonderful package for writing tests, but I did not find a
container for [Milvus vector database](https://milvus.io/docs)

Please check, I'm ready to correct comments

---------

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

Author: 3 people

index.rst

@@ -28,6 +28,7 @@ testcontainers-python facilitates the use of Docker containers for functional an
     modules/keycloak/README
     modules/localstack/README
     modules/memcached/README
+    modules/milvus/README
     modules/minio/README
     modules/mongodb/README
     modules/mssql/README

modules/milvus/README.rst

@@ -0,0 +1,2 @@
+.. autoclass:: testcontainers.milvus.MilvusContainer
+.. title:: testcontainers.milvus.MilvusContainer

modules/milvus/testcontainers/milvus/__init__.py

@@ -0,0 +1,85 @@
+#
+#    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 requests
+
+from testcontainers.core.config import testcontainers_config as c
+from testcontainers.core.generic import DockerContainer
+from testcontainers.core.waiting_utils import wait_container_is_ready, wait_for_logs
+
+
+class MilvusContainer(DockerContainer):
+    """
+    Milvus database container.
+
+    Read mode about Milvus: https://milvus.io/docs
+
+    Example:
+
+        The example spins up a Milvus database and connects to it client using MilvisClient.
+
+        .. doctest::
+
+            >>> from testcontainers.milvus import MilvusContainer
+            >>> with MilvusContainer("milvusdb/milvus:v2.4.4") as milvus_container:
+            ...     milvus_container.get_exposed_port(milvus_container.port) in milvus_container.get_connection_url()
+            True
+    """
+
+    def __init__(
+        self,
+        image: str = "milvusdb/milvus:latest",
+        port: int = 19530,
+        **kwargs,
+    ) -> None:
+        super().__init__(image=image, **kwargs)
+        self.port = port
+        self.healthcheck_port = 9091
+        self.with_exposed_ports(self.port, self.healthcheck_port)
+        self.cmd = "milvus run standalone"
+
+        envs = {"ETCD_USE_EMBED": "true", "ETCD_DATA_DIR": "/var/lib/milvus/etcd", "COMMON_STORAGETYPE": "local"}
+
+        for env, value in envs.items():
+            self.with_env(env, value)
+
+    def get_connection_url(self) -> str:
+        ip = self.get_container_host_ip()
+        port = self.get_exposed_port(self.port)
+        return f"http://{ip}:{port}"
+
+    @wait_container_is_ready()
+    def _connect(self) -> None:
+        msg = "Welcome to use Milvus!"
+        wait_for_logs(self, f".*{msg}.*", c.max_tries, c.sleep_time)
+        self._healthcheck()
+
+    def _get_healthcheck_url(self) -> str:
+        ip = self.get_container_host_ip()
+        port = self.get_exposed_port(self.healthcheck_port)
+        return f"http://{ip}:{port}"
+
+    @wait_container_is_ready(requests.exceptions.HTTPError)
+    def _healthcheck(self) -> None:
+        healthcheck_url = self._get_healthcheck_url()
+        response = requests.get(f"{healthcheck_url}/healthz", timeout=1)
+        response.raise_for_status()
+
+    def start(self) -> "MilvusContainer":
+        """This method starts the Milvus container and runs the healthcheck
+        to verify that the container is ready to use."""
+        self.with_command(self.cmd)
+        super().start()
+        self._connect()
+        self._healthcheck()
+        return self

modules/milvus/tests/test_milvus.py

@@ -0,0 +1,39 @@
+import pytest
+from pymilvus import MilvusClient
+
+from testcontainers.milvus import MilvusContainer
+
+VERSIONS = ["v2.4.0", "v2.4.4"]
+
+
+class ClientMilvusContainer(MilvusContainer):
+    def get_client(self, *, dbname: str = "default", token: str = "root:Milvus") -> MilvusClient:
+        connection_url = self.get_connection_url()
+        client = MilvusClient(uri=connection_url, dbname=dbname, token=token)
+        return client
+
+
+@pytest.mark.parametrize("version", VERSIONS)
+def test_run_milvus_success(version: str):
+    image = f"milvusdb/milvus:{version}"
+
+    with MilvusContainer(image=image) as milvus_container:
+        exposed_port = milvus_container.get_exposed_port(milvus_container.port)
+        url = milvus_container.get_connection_url()
+
+    assert url and exposed_port in url
+
+
+@pytest.mark.parametrize("version", VERSIONS)
+def test_milvus_client_success(version: str):
+    image = f"milvusdb/milvus:{version}"
+    test_collection = "test_collection"
+
+    with ClientMilvusContainer(image=image) as milvus_container:
+        client = milvus_container.get_client()
+        client.create_collection(test_collection, dimension=2)
+        collections = client.list_collections()
+        assert test_collection in collections
+
+        client.drop_collection(test_collection)
+        assert not client.has_collection(test_collection)