Skip to content

Commit 6f1be19

Browse files
author
Mehdi BEN ABDALLAH
committed
wip
1 parent cfc12c5 commit 6f1be19

File tree

10 files changed

+258
-206
lines changed

10 files changed

+258
-206
lines changed

‎modules/cosmosdb/README.rst

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,3 @@
1-
.. autoclass:: testcontainers.cosmosdb.CosmosDBEmulatorContainer
2-
.. autoclass:: testcontainers.cosmosdb.Endpoints
1+
.. autoclass:: testcontainers.cosmosdb.MongoDBEmulatorContainer
2+
.. autoclass:: testcontainers.cosmosdb.NoSQLEmulatorContainer
33
.. title:: testcontainers.cosmosdb.CosmosDBEmulatorContainer

‎modules/cosmosdb/testcontainers/cosmosdb/__init__.py

Lines changed: 3 additions & 182 deletions
Original file line numberDiff line numberDiff line change
@@ -1,183 +1,4 @@
1-
import os
2-
import socket
3-
import ssl
4-
from collections.abc import Iterable
5-
from enum import Enum, auto
6-
from typing import Callable, Optional
7-
from urllib.error import HTTPError, URLError
8-
from urllib.request import urlopen
1+
from .mongodb import MongoDBEmulatorContainer
2+
from .nosql import NoSQLEmulatorContainer
93

10-
from azure.core.exceptions import ServiceRequestError
11-
from azure.cosmos import CosmosClient as SyncCosmosClient
12-
from azure.cosmos.aio import CosmosClient as AsyncCosmosClient
13-
from typing_extensions import Self
14-
15-
from testcontainers.core.container import DockerContainer
16-
from testcontainers.core.waiting_utils import wait_container_is_ready, wait_for_logs
17-
18-
__all__ = ["CosmosDBEmulatorContainer", "Endpoints"]
19-
20-
21-
class Endpoints(Enum):
22-
MongoDB = auto()
23-
24-
25-
# Ports mostly derived from https://docs.microsoft.com/en-us/azure/cosmos-db/emulator-command-line-parameters
26-
EMULATOR_PORT = 8081
27-
endpoint_ports = {
28-
Endpoints.MongoDB: frozenset([10255]),
29-
}
30-
31-
32-
def is_truthy_string(s: str):
33-
return s.lower().strip() in {"true", "yes", "y", "1"}
34-
35-
36-
class CosmosDBEmulatorContainer(DockerContainer):
37-
"""
38-
CosmosDB Emulator container.
39-
40-
Example:
41-
.. doctest::
42-
>>> from testcontainers.cosmosdb import CosmosDBEmulatorContainer
43-
>>> with CosmosDBEmulatorContainer() as cosmosdb:
44-
... db = cosmosdb.insecure_sync_client().create_database_if_not_exists("test")
45-
46-
.. doctest::
47-
>>> from testcontainers.cosmosdb import CosmosDBEmulatorContainer
48-
>>> with CosmosDBEmulatorContainer() as emulator:
49-
... cosmosdb = CosmosClient(url=emulator.url, credential=emulator.key, connection_verify=False)
50-
... db = cosmosdb.create_database_if_not_exists("test")
51-
52-
.. doctest::
53-
>>> from testcontainers.cosmosdb import CosmosDBEmulatorContainer, Endpoints
54-
>>> with CosmosDBEmulatorContainer(endpoints=[Endpoints.MongoDB], mongodb_version="4.0") as emulator:
55-
... print(f"Point yout MongoDB client to {emulator.host}:{next(iter(emulator.ports(Endpoints.MongoDB)))}")
56-
"""
57-
58-
def __init__(
59-
self,
60-
image: str = os.getenv(
61-
"AZURE_COSMOS_EMULATOR_IMAGE", "mcr.microsoft.com/cosmosdb/linux/azure-cosmos-emulator:latest"
62-
),
63-
partition_count: int = os.getenv("AZURE_COSMOS_EMULATOR_PARTITION_COUNT", None),
64-
enable_data_persistence: bool = is_truthy_string(
65-
os.getenv("AZURE_COSMOS_EMULATOR_ENABLE_DATA_PERSISTENCE", "false")
66-
),
67-
bind_ports: bool = is_truthy_string(os.getenv("AZURE_COSMOS_EMULATOR_BIND_PORTS", "true")),
68-
key: str = os.getenv(
69-
"AZURE_COSMOS_EMULATOR_KEY",
70-
"C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==",
71-
),
72-
endpoints: Iterable[Endpoints] = [], # the emulator image does not support host-container port mapping
73-
mongodb_version: Optional[str] = None,
74-
**docker_client_kw,
75-
):
76-
super().__init__(image=image, **docker_client_kw)
77-
self.partition_count = partition_count
78-
self.key = key
79-
self.enable_data_persistence = enable_data_persistence
80-
self.endpoints = frozenset(endpoints)
81-
self.bind_ports = bind_ports
82-
assert (Endpoints.MongoDB not in self.endpoints) or (
83-
mongodb_version is not None
84-
), "A MongoDB version is required to use the MongoDB Endpoint"
85-
self.mongodb_version = mongodb_version
86-
87-
@property
88-
def url(self) -> str:
89-
"""
90-
The url to the CosmosDB server
91-
"""
92-
return f"https://{self.host}:{self.get_exposed_port(EMULATOR_PORT)}"
93-
94-
@property
95-
def host(self) -> str:
96-
return self.get_container_host_ip()
97-
98-
@property
99-
def certificate_pem(self) -> bytes:
100-
"""
101-
PEM-encoded certificate of the CosmosDB server
102-
"""
103-
return self._cert_pem_bytes
104-
105-
def ports(self, endpoint: Endpoints) -> Iterable[int]:
106-
"""
107-
Returns the set of exposed ports for a given endpoint.
108-
If bind_ports is True, the returned ports will be the NAT-ed ports reachable from the host.
109-
"""
110-
assert endpoint in self.endpoints, f"Endpoint {endpoint} is not exposed"
111-
return {self.get_exposed_port(p) for p in endpoint_ports[endpoint]}
112-
113-
def insecure_async_client(self) -> AsyncCosmosClient:
114-
"""
115-
Returns an asynchronous CosmosClient instance
116-
"""
117-
return AsyncCosmosClient(url=self.url, credential=self.key, connection_verify=False)
118-
119-
def insecure_sync_client(self) -> SyncCosmosClient:
120-
"""
121-
Returns a synchronous CosmosClient instance
122-
"""
123-
return SyncCosmosClient(url=self.url, credential=self.key, connection_verify=False)
124-
125-
def start(self) -> Self:
126-
self._configure()
127-
super().start()
128-
self._wait_until_ready()
129-
self._cert_pem_bytes = self._download_cert()
130-
return self
131-
132-
def _configure(self) -> None:
133-
self.with_bind_ports(EMULATOR_PORT, EMULATOR_PORT)
134-
135-
endpoints_ports = []
136-
for endpoint in self.endpoints:
137-
endpoints_ports.extend(endpoint_ports[endpoint])
138-
139-
if self.bind_ports:
140-
[self.with_bind_ports(port, port) for port in endpoints_ports]
141-
else:
142-
self.with_exposed_ports(*endpoints_ports)
143-
144-
(
145-
self.with_env("AZURE_COSMOS_EMULATOR_PARTITION_COUNT", str(self.partition_count))
146-
.with_env("AZURE_COSMOS_EMULATOR_IP_ADDRESS_OVERRIDE", socket.gethostbyname(socket.gethostname()))
147-
.with_env("AZURE_COSMOS_EMULATOR_ENABLE_DATA_PERSISTENCE", str(self.enable_data_persistence))
148-
.with_env("AZURE_COSMOS_EMULATOR_KEY", str(self.key))
149-
)
150-
151-
if Endpoints.MongoDB in self.endpoints:
152-
self.with_env("AZURE_COSMOS_EMULATOR_ENABLE_MONGODB_ENDPOINT", self.mongodb_version)
153-
154-
def _wait_until_ready(self) -> Self:
155-
"""
156-
Waits until the CosmosDB Emulator image is ready to be used.
157-
"""
158-
(
159-
self._wait_for_logs(container=self, predicate="Started\\s*$")
160-
._wait_for_url(f"{self.url}/_explorer/index.html")
161-
._wait_for_query_success(lambda sync_client: list(sync_client.list_databases()))
162-
)
163-
return self
164-
165-
@wait_container_is_ready(HTTPError, URLError)
166-
def _wait_for_url(self, url: str) -> Self:
167-
with urlopen(url, context=ssl._create_unverified_context()) as response:
168-
response.read()
169-
return self
170-
171-
def _wait_for_logs(self, *args, **kwargs) -> Self:
172-
wait_for_logs(*args, **kwargs)
173-
return self
174-
175-
@wait_container_is_ready(ServiceRequestError)
176-
def _wait_for_query_success(self, query: Callable[[SyncCosmosClient], None]) -> Self:
177-
with self.insecure_sync_client() as c:
178-
query(c)
179-
return self
180-
181-
def _download_cert(self) -> bytes:
182-
with urlopen(f"{self.url}/_explorer/emulator.pem", context=ssl._create_unverified_context()) as response:
183-
return response.read()
4+
__all__ = ["MongoDBEmulatorContainer", "NoSQLEmulatorContainer"]

‎modules/cosmosdb/testcontainers/cosmosdb/_emulator.py

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
import os
2+
import socket
3+
import ssl
4+
from collections.abc import Iterable
5+
from typing_extensions import Self
6+
from testcontainers.core.container import DockerContainer
7+
from testcontainers.core.waiting_utils import wait_container_is_ready, wait_for_logs
8+
from . import _grab as grab
9+
from distutils.util import strtobool
10+
from urllib.error import HTTPError, URLError
11+
from urllib.request import urlopen
12+
13+
__all__ = ["CosmosDBEmulatorContainer"]
14+
15+
EMULATOR_PORT = 8081
16+
17+
class CosmosDBEmulatorContainer(DockerContainer):
18+
"""
19+
CosmosDB Emulator container.
20+
"""
21+
22+
def __init__(
23+
self,
24+
image: str = os.getenv(
25+
"AZURE_COSMOS_EMULATOR_IMAGE", "mcr.microsoft.com/cosmosdb/linux/azure-cosmos-emulator:latest"
26+
),
27+
partition_count: int = os.getenv("AZURE_COSMOS_EMULATOR_PARTITION_COUNT", None),
28+
enable_data_persistence: bool = strtobool(
29+
os.getenv("AZURE_COSMOS_EMULATOR_ENABLE_DATA_PERSISTENCE", "false")
30+
),
31+
key: str = os.getenv(
32+
"AZURE_COSMOS_EMULATOR_KEY",
33+
"C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==",
34+
),
35+
bind_ports: bool = strtobool(os.getenv("AZURE_COSMOS_EMULATOR_BIND_PORTS", "true")),
36+
endpoint_ports: Iterable[int] = [],
37+
**other_kwargs,
38+
):
39+
super().__init__(image=image, **other_kwargs)
40+
self.endpoint_ports = endpoint_ports
41+
self.partition_count = partition_count
42+
self.key = key
43+
self.enable_data_persistence = enable_data_persistence
44+
self.bind_ports = bind_ports
45+
46+
@property
47+
def host(self) -> str:
48+
return self.get_container_host_ip()
49+
50+
@property
51+
def server_certificate_pem(self) -> bytes:
52+
"""
53+
PEM-encoded server certificate
54+
"""
55+
return self._cert_pem_bytes
56+
57+
def start(self) -> Self:
58+
self._configure()
59+
super().start()
60+
self._wait_until_ready()
61+
self._cert_pem_bytes = self._download_cert()
62+
return self
63+
64+
def _configure(self) -> None:
65+
all_ports = set([EMULATOR_PORT] + self.endpoint_ports)
66+
if self.bind_ports:
67+
for port in all_ports:
68+
self.with_bind_ports(port, port)
69+
else:
70+
self.with_exposed_ports(*all_ports)
71+
72+
(
73+
self
74+
.with_env("AZURE_COSMOS_EMULATOR_PARTITION_COUNT", str(self.partition_count))
75+
.with_env("AZURE_COSMOS_EMULATOR_IP_ADDRESS_OVERRIDE", socket.gethostbyname(socket.gethostname()))
76+
.with_env("AZURE_COSMOS_EMULATOR_ENABLE_DATA_PERSISTENCE", str(self.enable_data_persistence))
77+
.with_env("AZURE_COSMOS_EMULATOR_KEY", str(self.key))
78+
)
79+
80+
def _wait_until_ready(self) -> Self:
81+
wait_for_logs(container=self, predicate="Started\\s*$")
82+
83+
if self.bind_ports:
84+
self._wait_for_url(f"https://{self.host}:{EMULATOR_PORT}/_explorer/index.html")
85+
self._wait_for_query_success()
86+
87+
return self
88+
89+
def _download_cert(self) -> bytes:
90+
with grab.file(
91+
self._container, "/tmp/cosmos/appdata/.system/profiles/Client/AppData/Local/CosmosDBEmulator/emulator.pem"
92+
) as cert:
93+
return cert.read()
94+
95+
@wait_container_is_ready(HTTPError, URLError)
96+
def _wait_for_url(self, url: str) -> Self:
97+
with urlopen(url, context=ssl._create_unverified_context()) as response:
98+
response.read()
99+
return self
100+
101+
def _wait_for_query_success(self) -> None:
102+
pass

‎modules/cosmosdb/testcontainers/cosmosdb/_grab.py

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
2+
from pathlib import Path
3+
from os import path
4+
import tarfile
5+
import tempfile
6+
from contextlib import contextmanager
7+
8+
@contextmanager
9+
def file(container, target):
10+
target_path = Path(target)
11+
assert target_path.is_absolute(), "target must be an absolute path"
12+
13+
with tempfile.TemporaryDirectory() as tmpdirname:
14+
archive = Path(tmpdirname) / 'grabbed.tar'
15+
16+
# download from container as tar archive
17+
with open(archive, 'wb') as f:
18+
tar_bits, _ = container.get_archive(target)
19+
for chunk in tar_bits:
20+
f.write(chunk)
21+
22+
# extract target file from tar archive
23+
with tarfile.TarFile(archive) as tar:
24+
yield tar.extractfile(path.basename(target))
25+

‎modules/cosmosdb/testcontainers/cosmosdb/mongodb.py

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
import os
2+
from ._emulator import CosmosDBEmulatorContainer
3+
4+
__all__ = ["MongoDBEmulatorContainer"]
5+
6+
ENDPOINT_PORT = 10255
7+
8+
class MongoDBEmulatorContainer(CosmosDBEmulatorContainer):
9+
"""
10+
CosmosDB MongoDB enpoint Emulator.
11+
12+
Example:
13+
14+
.. doctest::
15+
>>> from testcontainers.cosmosdb import MongoDBEmulatorContainer
16+
>>> with CosmosDBEmulatorContainer(mongodb_version="4.0") as emulator:
17+
... print(f"Point yout MongoDB client to {emulator.host}:{emulator.port}}")
18+
"""
19+
20+
def __init__(
21+
self,
22+
mongodb_version: str = None,
23+
image: str = os.getenv(
24+
"AZURE_COSMOS_EMULATOR_IMAGE", "mcr.microsoft.com/cosmosdb/linux/azure-cosmos-emulator:mongodb"
25+
),
26+
**other_kwargs,
27+
):
28+
super().__init__(image=image, endpoint_ports=[ENDPOINT_PORT], **other_kwargs)
29+
assert mongodb_version is not None, "A MongoDB version is required to use the MongoDB Endpoint"
30+
self.mongodb_version = mongodb_version
31+
32+
@property
33+
def port(self) -> str:
34+
return self.get_exposed_port(ENDPOINT_PORT)
35+
36+
def _configure(self) -> None:
37+
super()._configure()
38+
self.with_env("AZURE_COSMOS_EMULATOR_ENABLE_MONGODB_ENDPOINT", self.mongodb_version)

0 commit comments

Comments
 (0)