From 139eb716b744c23ec0afa1a4017e86d2ad9f0657 Mon Sep 17 00:00:00 2001 From: Mehdi BEN ABDALLAH <@mbenabda> Date: Fri, 24 May 2024 21:57:04 +0200 Subject: [PATCH] feat(cosmosdb) : add support for the CosmosDB emulator --- index.rst | 1 + modules/cosmosdb/README.rst | 2 + .../testcontainers/cosmosdb/__init__.py | 158 ++++++++++++++++++ modules/cosmosdb/tests/test_cosmosdb.py | 6 + pyproject.toml | 4 + 5 files changed, 171 insertions(+) create mode 100644 modules/cosmosdb/README.rst create mode 100644 modules/cosmosdb/testcontainers/cosmosdb/__init__.py create mode 100644 modules/cosmosdb/tests/test_cosmosdb.py diff --git a/index.rst b/index.rst index 3c7fcc140..9091738b8 100644 --- a/index.rst +++ b/index.rst @@ -20,6 +20,7 @@ testcontainers-python facilitates the use of Docker containers for functional an modules/cassandra/README modules/chroma/README modules/clickhouse/README + modules/cosmosdb/README modules/elasticsearch/README modules/google/README modules/influxdb/README diff --git a/modules/cosmosdb/README.rst b/modules/cosmosdb/README.rst new file mode 100644 index 000000000..bf97ac105 --- /dev/null +++ b/modules/cosmosdb/README.rst @@ -0,0 +1,2 @@ +.. autoclass:: testcontainers.cosmosdb.CosmosDBEmulatorContainer +.. title:: testcontainers.cosmosdb.CosmosDBEmulatorContainer diff --git a/modules/cosmosdb/testcontainers/cosmosdb/__init__.py b/modules/cosmosdb/testcontainers/cosmosdb/__init__.py new file mode 100644 index 000000000..259ea68de --- /dev/null +++ b/modules/cosmosdb/testcontainers/cosmosdb/__init__.py @@ -0,0 +1,158 @@ +from testcontainers.core.container import DockerContainer +from testcontainers.core.waiting_utils import wait_for_logs, wait_container_is_ready +import os +import ssl +import socket +from typing import Iterable, Callable +from typing_extensions import Self +from azure.cosmos import CosmosClient as SyncCosmosClient +from azure.cosmos.aio import CosmosClient as AsyncCosmosClient +from azure.core.exceptions import ServiceRequestError + +from urllib.request import urlopen +from urllib.error import HTTPError, URLError + +from enum import Enum, auto + +__all__ = ["CosmosDBEmulatorContainer", "Endpoints"] + +class Endpoints(Enum): + Direct = auto() + Gremlin = auto() + Table = auto() + MongoDB = auto() + Cassandra = auto() + +ALL_ENDPOINTS = { e for e in Endpoints } + +# Ports mostly derived from https://docs.microsoft.com/en-us/azure/cosmos-db/emulator-command-line-parameters +EMULATOR_PORT = 8081 +endpoint_ports = { + Endpoints.Direct : frozenset([10251, 10252, 10253, 10254]), + Endpoints.Gremlin : frozenset([8901]), + Endpoints.Table : frozenset([8902]), + Endpoints.MongoDB : frozenset([10255]), + Endpoints.Cassandra: frozenset([10350]), +} + +def is_truthy_string(s: str): + return s.lower().strip() in {"true", "yes", "y", "1"} + +class CosmosDBEmulatorContainer(DockerContainer): + """ + CosmosDB Emulator container. + + Example: + + .. doctest:: + >>> from testcontainers.cosmosdb import CosmosDBEmulatorContainer + >>> with CosmosDBEmulatorContainer() as cosmosdb: + ... db = cosmosdb.sync_client().create_database_if_not_exists("test") + + .. doctest:: + >>> from testcontainers.cosmosdb import CosmosDBEmulatorContainer + >>> with CosmosDBEmulatorContainer() as emulator: + ... cosmosdb = CosmosClient(url=emulator.url, credential=emulator.key, connection_verify=False) + ... db = cosmosdb.create_database_if_not_exists("test") + + .. doctest:: + >>> from testcontainers.cosmosdb import CosmosDBEmulatorContainer, Endpoints + >>> with CosmosDBEmulatorContainer(endpoints=[Endpoints.MongoDB]) as emulator: + ... print(f"Point yout MongoDB client to {emulator.host}:{emulator.ports(Endpoints.MongoDB)[0]}") + """ + def __init__( + self, + image: str = os.getenv("AZURE_COSMOS_EMULATOR_IMAGE", "mcr.microsoft.com/cosmosdb/linux/azure-cosmos-emulator:latest"), + partition_count: int = os.getenv("AZURE_COSMOS_EMULATOR_PARTITION_COUNT", None), + enable_data_persistence: bool = is_truthy_string(os.getenv("AZURE_COSMOS_EMULATOR_ENABLE_DATA_PERSISTENCE", "false")), + bind_ports: bool = is_truthy_string(os.getenv("AZURE_COSMOS_EMULATOR_BIND_PORTS", "true")), + key: str = os.getenv("AZURE_COSMOS_EMULATOR_KEY", "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw=="), + endpoints: Iterable[Endpoints] = ALL_ENDPOINTS, # the emulator image does not support host-container port mapping + **docker_client_kw, + ): + super().__init__(image=image, **docker_client_kw) + + self.partition_count = partition_count + self.key = key + self.enable_data_persistence = enable_data_persistence + self.endpoints = frozenset(endpoints) + + self.with_bind_ports(EMULATOR_PORT, EMULATOR_PORT) + + endpoints_ports = [] + for endpoint in self.endpoints: + endpoints_ports.extend(endpoint_ports[endpoint]) + + if bind_ports: + [ self.with_bind_ports(port, port) for port in endpoints_ports ] + else: + self.with_exposed_ports(*endpoints_ports) + + def start(self) -> Self: + self._configure() + super().start() + self._wait_until_ready() + return self + + @property + def url(self) -> str: + """ + Returns the url to interact with the emulator + """ + return f"https://{self.host}:{self.get_exposed_port(EMULATOR_PORT)}" + + @property + def host(self) -> str: + return self.get_container_host_ip() + + def ports(self, endpoint: Endpoints) -> Iterable[int]: + assert endpoint in self.endpoints, f"Endpoint {endpoint} is not exposed" + return { self.get_exposed_port(p) for p in endpoint_ports[endpoint] } + + def async_client(self) -> AsyncCosmosClient: + """ + Returns an asynchronous CosmosClient instance to interact with the CosmosDB server + """ + return AsyncCosmosClient(url=self.url, credential=self.key, connection_verify=False) + + def sync_client(self) -> SyncCosmosClient: + """ + Returns a synchronous CosmosClient instance to interact with the CosmosDB server + """ + return SyncCosmosClient(url=self.url, credential=self.key, connection_verify=False) + + def _configure(self) -> None: + ( + self + .with_env("AZURE_COSMOS_EMULATOR_PARTITION_COUNT", str(self.partition_count)) + .with_env("AZURE_COSMOS_EMULATOR_IP_ADDRESS_OVERRIDE", socket.gethostbyname(socket.gethostname())) + .with_env("AZURE_COSMOS_EMULATOR_ENABLE_DATA_PERSISTENCE", str(self.enable_data_persistence)) + .with_env("AZURE_COSMOS_EMULATOR_KEY", str(self.key)) + ) + + @wait_container_is_ready(HTTPError, URLError, ServiceRequestError) + def _wait_until_ready(self) -> Self: + """ + Waits until the CosmosDB Emulator image is ready to be used. + """ + ( + self + ._wait_for_logs(container=self, predicate="Started\\s*$") + ._wait_for_url(f"{self.url}/_explorer/index.html") + ._wait_for_query_success(lambda sync_client: list(sync_client.list_databases())) + ) + return self + + def _wait_for_url(self, url: str) -> Self: + with urlopen(url, context=ssl._create_unverified_context()) as response: + response.read() + return self + + def _wait_for_logs(self, *args, **kwargs) -> Self: + wait_for_logs(*args, **kwargs) + return self + + def _wait_for_query_success(self, query: Callable[[SyncCosmosClient], None]) -> Self: + with self.sync_client() as c: + query(c) + return self diff --git a/modules/cosmosdb/tests/test_cosmosdb.py b/modules/cosmosdb/tests/test_cosmosdb.py new file mode 100644 index 000000000..17cfe257a --- /dev/null +++ b/modules/cosmosdb/tests/test_cosmosdb.py @@ -0,0 +1,6 @@ +import pytest +from testcontainers.cosmosdb import CosmosDBEmulatorContainer + +def test_docker_run(): + with CosmosDBEmulatorContainer(partition_count=1) as cosmosdb: + list(cosmosdb.sync_client().list_databases()) diff --git a/pyproject.toml b/pyproject.toml index 2594f3e44..76cd6992d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,6 +33,7 @@ packages = [ { include = "testcontainers", from = "modules/cassandra" }, { include = "testcontainers", from = "modules/chroma" }, { include = "testcontainers", from = "modules/clickhouse" }, + { include = "testcontainers", from = "modules/cosmosdb" }, { include = "testcontainers", from = "modules/elasticsearch" }, { include = "testcontainers", from = "modules/google" }, { include = "testcontainers", from = "modules/influxdb" }, @@ -75,6 +76,7 @@ typing-extensions = "*" python-arango = { version = "^7.8", optional = true } azure-storage-blob = { version = "^12.19", optional = true } clickhouse-driver = { version = "*", optional = true } +azure-cosmos = { version = "*", optional = true } google-cloud-pubsub = { version = ">=2", optional = true } google-cloud-datastore = { version = ">=2", optional = true } influxdb = { version = "*", optional = true } @@ -105,6 +107,7 @@ arangodb = ["python-arango"] azurite = ["azure-storage-blob"] cassandra = [] clickhouse = ["clickhouse-driver"] +cosmosdb = ["azure-cosmos"] elasticsearch = [] google = ["google-cloud-pubsub", "google-cloud-datastore"] influxdb = ["influxdb", "influxdb-client"] @@ -250,6 +253,7 @@ mypy_path = [ # "modules/azurite", # "modules/cassandra", # "modules/clickhouse", +# "modules/cosmosdb", # "modules/elasticsearch", # "modules/google", # "modules/k3s",