Skip to content

Commit

Permalink
fix: Container for Milvus database (#606)
Browse files Browse the repository at this point in the history
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 <ivan.belyaev@ailet.com>
Co-authored-by: David Ankin <daveankin@gmail.com>
  • Loading branch information
3 people committed Jun 13, 2024
1 parent f5a019b commit ec76df2
Show file tree
Hide file tree
Showing 6 changed files with 388 additions and 6 deletions.
1 change: 1 addition & 0 deletions index.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 2 additions & 0 deletions modules/milvus/README.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
.. autoclass:: testcontainers.milvus.MilvusContainer
.. title:: testcontainers.milvus.MilvusContainer
85 changes: 85 additions & 0 deletions modules/milvus/testcontainers/milvus/__init__.py
Original file line number Diff line number Diff line change
@@ -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
39 changes: 39 additions & 0 deletions modules/milvus/tests/test_milvus.py
Original file line number Diff line number Diff line change
@@ -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)
Loading

0 comments on commit ec76df2

Please sign in to comment.