Skip to content

Commit cce5a04

Browse files
alexanderankinalexanderankin
committed
linting
1 parent d11370f commit cce5a04

File tree

2 files changed

+150
-139
lines changed

2 files changed

+150
-139
lines changed

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

Lines changed: 147 additions & 137 deletions
Original file line numberDiff line numberDiff line change
@@ -1,159 +1,169 @@
1-
from testcontainers.core.container import DockerContainer
2-
from testcontainers.core.waiting_utils import wait_for_logs, wait_container_is_ready
31
import os
4-
import ssl
52
import socket
6-
from typing import Iterable, Callable
7-
from typing_extensions import Self
3+
import ssl
4+
from collections.abc import Iterable
5+
from enum import Enum, auto
6+
from typing import Callable
7+
from urllib.error import HTTPError, URLError
8+
from urllib.request import urlopen
9+
10+
from azure.core.exceptions import ServiceRequestError
811
from azure.cosmos import CosmosClient as SyncCosmosClient
912
from azure.cosmos.aio import CosmosClient as AsyncCosmosClient
10-
from azure.core.exceptions import ServiceRequestError
11-
12-
from urllib.request import urlopen
13-
from urllib.error import HTTPError, URLError
13+
from typing_extensions import Self
1414

15-
from enum import Enum, auto
15+
from testcontainers.core.container import DockerContainer
16+
from testcontainers.core.waiting_utils import wait_container_is_ready, wait_for_logs
1617

1718
__all__ = ["CosmosDBEmulatorContainer", "Endpoints"]
1819

20+
1921
class Endpoints(Enum):
20-
Direct = auto()
21-
Gremlin = auto()
22-
Table = auto()
23-
MongoDB = auto()
24-
Cassandra = auto()
22+
Direct = auto()
23+
Gremlin = auto()
24+
Table = auto()
25+
MongoDB = auto()
26+
Cassandra = auto()
27+
2528

26-
ALL_ENDPOINTS = { e for e in Endpoints }
29+
ALL_ENDPOINTS = set(Endpoints)
2730

2831
# Ports mostly derived from https://docs.microsoft.com/en-us/azure/cosmos-db/emulator-command-line-parameters
2932
EMULATOR_PORT = 8081
3033
endpoint_ports = {
31-
Endpoints.Direct : frozenset([10251, 10252, 10253, 10254]),
32-
Endpoints.Gremlin : frozenset([8901]),
33-
Endpoints.Table : frozenset([8902]),
34-
Endpoints.MongoDB : frozenset([10255]),
35-
Endpoints.Cassandra: frozenset([10350]),
34+
Endpoints.Direct: frozenset([10251, 10252, 10253, 10254]),
35+
Endpoints.Gremlin: frozenset([8901]),
36+
Endpoints.Table: frozenset([8902]),
37+
Endpoints.MongoDB: frozenset([10255]),
38+
Endpoints.Cassandra: frozenset([10350]),
3639
}
3740

41+
3842
def is_truthy_string(s: str):
39-
return s.lower().strip() in {"true", "yes", "y", "1"}
43+
return s.lower().strip() in {"true", "yes", "y", "1"}
44+
4045

4146
class CosmosDBEmulatorContainer(DockerContainer):
42-
"""
47+
"""
4348
CosmosDB Emulator container.
4449
4550
Example:
46-
47-
.. doctest::
48-
>>> from testcontainers.cosmosdb import CosmosDBEmulatorContainer
49-
>>> with CosmosDBEmulatorContainer() as cosmosdb:
50-
... db = cosmosdb.sync_client().create_database_if_not_exists("test")
51-
52-
.. doctest::
53-
>>> from testcontainers.cosmosdb import CosmosDBEmulatorContainer
54-
>>> with CosmosDBEmulatorContainer() as emulator:
55-
... cosmosdb = CosmosClient(url=emulator.url, credential=emulator.key, connection_verify=False)
56-
... db = cosmosdb.create_database_if_not_exists("test")
57-
58-
.. doctest::
59-
>>> from testcontainers.cosmosdb import CosmosDBEmulatorContainer, Endpoints
60-
>>> with CosmosDBEmulatorContainer(endpoints=[Endpoints.MongoDB]) as emulator:
61-
... print(f"Point yout MongoDB client to {emulator.host}:{emulator.ports(Endpoints.MongoDB)[0]}")
62-
"""
63-
def __init__(
64-
self,
65-
image: str = os.getenv("AZURE_COSMOS_EMULATOR_IMAGE", "mcr.microsoft.com/cosmosdb/linux/azure-cosmos-emulator:latest"),
66-
partition_count: int = os.getenv("AZURE_COSMOS_EMULATOR_PARTITION_COUNT", None),
67-
enable_data_persistence: bool = is_truthy_string(os.getenv("AZURE_COSMOS_EMULATOR_ENABLE_DATA_PERSISTENCE", "false")),
68-
bind_ports: bool = is_truthy_string(os.getenv("AZURE_COSMOS_EMULATOR_BIND_PORTS", "true")),
69-
key: str = os.getenv("AZURE_COSMOS_EMULATOR_KEY", "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw=="),
70-
endpoints: Iterable[Endpoints] = [], # the emulator image does not support host-container port mapping
71-
**docker_client_kw,
72-
):
73-
super().__init__(image=image, **docker_client_kw)
74-
self.partition_count = partition_count
75-
self.key = key
76-
self.enable_data_persistence = enable_data_persistence
77-
self.endpoints = frozenset(endpoints)
78-
self.bind_ports = bind_ports
79-
80-
def start(self) -> Self:
81-
self._configure()
82-
super().start()
83-
self._wait_until_ready()
84-
return self
85-
86-
@property
87-
def url(self) -> str:
88-
"""
89-
Returns the url to interact with the emulator
90-
"""
91-
return f"https://{self.host}:{self.get_exposed_port(EMULATOR_PORT)}"
92-
93-
@property
94-
def host(self) -> str:
95-
return self.get_container_host_ip()
96-
97-
def ports(self, endpoint: Endpoints) -> Iterable[int]:
98-
assert endpoint in self.endpoints, f"Endpoint {endpoint} is not exposed"
99-
return { self.get_exposed_port(p) for p in endpoint_ports[endpoint] }
100-
101-
def async_client(self) -> AsyncCosmosClient:
102-
"""
103-
Returns an asynchronous CosmosClient instance to interact with the CosmosDB server
104-
"""
105-
return AsyncCosmosClient(url=self.url, credential=self.key, connection_verify=False)
106-
107-
def sync_client(self) -> SyncCosmosClient:
108-
"""
109-
Returns a synchronous CosmosClient instance to interact with the CosmosDB server
110-
"""
111-
return SyncCosmosClient(url=self.url, credential=self.key, connection_verify=False)
112-
113-
def _configure(self) -> None:
114-
self.with_bind_ports(EMULATOR_PORT, EMULATOR_PORT)
115-
116-
endpoints_ports = []
117-
for endpoint in self.endpoints:
118-
endpoints_ports.extend(endpoint_ports[endpoint])
119-
120-
if self.bind_ports:
121-
[ self.with_bind_ports(port, port) for port in endpoints_ports ]
122-
else:
123-
self.with_exposed_ports(*endpoints_ports)
124-
125-
(
126-
self
127-
.with_env("AZURE_COSMOS_EMULATOR_PARTITION_COUNT", str(self.partition_count))
128-
.with_env("AZURE_COSMOS_EMULATOR_IP_ADDRESS_OVERRIDE", socket.gethostbyname(socket.gethostname()))
129-
.with_env("AZURE_COSMOS_EMULATOR_ENABLE_DATA_PERSISTENCE", str(self.enable_data_persistence))
130-
.with_env("AZURE_COSMOS_EMULATOR_KEY", str(self.key))
131-
)
132-
133-
def _wait_until_ready(self) -> Self:
134-
"""
135-
Waits until the CosmosDB Emulator image is ready to be used.
136-
"""
137-
(
138-
self
139-
._wait_for_logs(container=self, predicate="Started\\s*$")
140-
._wait_for_url(f"{self.url}/_explorer/index.html")
141-
._wait_for_query_success(lambda sync_client: list(sync_client.list_databases()))
142-
)
143-
return self
144-
145-
@wait_container_is_ready(HTTPError, URLError)
146-
def _wait_for_url(self, url: str) -> Self:
147-
with urlopen(url, context=ssl._create_unverified_context()) as response:
148-
response.read()
149-
return self
150-
151-
def _wait_for_logs(self, *args, **kwargs) -> Self:
152-
wait_for_logs(*args, **kwargs)
153-
return self
154-
155-
@wait_container_is_ready(ServiceRequestError)
156-
def _wait_for_query_success(self, query: Callable[[SyncCosmosClient], None]) -> Self:
157-
with self.sync_client() as c:
158-
query(c)
159-
return self
51+
.. doctest::
52+
>>> from testcontainers.cosmosdb import CosmosDBEmulatorContainer
53+
>>> with CosmosDBEmulatorContainer() as cosmosdb:
54+
... db = cosmosdb.sync_client().create_database_if_not_exists("test")
55+
56+
.. doctest::
57+
>>> from testcontainers.cosmosdb import CosmosDBEmulatorContainer
58+
>>> with CosmosDBEmulatorContainer() as emulator:
59+
... cosmosdb = CosmosClient(url=emulator.url, credential=emulator.key, connection_verify=False)
60+
... db = cosmosdb.create_database_if_not_exists("test")
61+
62+
.. doctest::
63+
>>> from testcontainers.cosmosdb import CosmosDBEmulatorContainer, Endpoints
64+
>>> with CosmosDBEmulatorContainer(endpoints=[Endpoints.MongoDB]) as emulator:
65+
... print(f"Point yout MongoDB client to {emulator.host}:{emulator.ports(Endpoints.MongoDB)[0]}")
66+
"""
67+
68+
def __init__(
69+
self,
70+
image: str = os.getenv(
71+
"AZURE_COSMOS_EMULATOR_IMAGE", "mcr.microsoft.com/cosmosdb/linux/azure-cosmos-emulator:latest"
72+
),
73+
partition_count: int = os.getenv("AZURE_COSMOS_EMULATOR_PARTITION_COUNT", None),
74+
enable_data_persistence: bool = is_truthy_string(
75+
os.getenv("AZURE_COSMOS_EMULATOR_ENABLE_DATA_PERSISTENCE", "false")
76+
),
77+
bind_ports: bool = is_truthy_string(os.getenv("AZURE_COSMOS_EMULATOR_BIND_PORTS", "true")),
78+
key: str = os.getenv(
79+
"AZURE_COSMOS_EMULATOR_KEY",
80+
"C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==",
81+
),
82+
endpoints: Iterable[Endpoints] = [], # the emulator image does not support host-container port mapping
83+
**docker_client_kw,
84+
):
85+
super().__init__(image=image, **docker_client_kw)
86+
self.partition_count = partition_count
87+
self.key = key
88+
self.enable_data_persistence = enable_data_persistence
89+
self.endpoints = frozenset(endpoints)
90+
self.bind_ports = bind_ports
91+
92+
def start(self) -> Self:
93+
self._configure()
94+
super().start()
95+
self._wait_until_ready()
96+
return self
97+
98+
@property
99+
def url(self) -> str:
100+
"""
101+
Returns the url to interact with the emulator
102+
"""
103+
return f"https://{self.host}:{self.get_exposed_port(EMULATOR_PORT)}"
104+
105+
@property
106+
def host(self) -> str:
107+
return self.get_container_host_ip()
108+
109+
def ports(self, endpoint: Endpoints) -> Iterable[int]:
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 async_client(self) -> AsyncCosmosClient:
114+
"""
115+
Returns an asynchronous CosmosClient instance to interact with the CosmosDB server
116+
"""
117+
return AsyncCosmosClient(url=self.url, credential=self.key, connection_verify=False)
118+
119+
def sync_client(self) -> SyncCosmosClient:
120+
"""
121+
Returns a synchronous CosmosClient instance to interact with the CosmosDB server
122+
"""
123+
return SyncCosmosClient(url=self.url, credential=self.key, connection_verify=False)
124+
125+
def _configure(self) -> None:
126+
self.with_bind_ports(EMULATOR_PORT, EMULATOR_PORT)
127+
128+
endpoints_ports = []
129+
for endpoint in self.endpoints:
130+
endpoints_ports.extend(endpoint_ports[endpoint])
131+
132+
if self.bind_ports:
133+
[self.with_bind_ports(port, port) for port in endpoints_ports]
134+
else:
135+
self.with_exposed_ports(*endpoints_ports)
136+
137+
(
138+
self.with_env("AZURE_COSMOS_EMULATOR_PARTITION_COUNT", str(self.partition_count))
139+
.with_env("AZURE_COSMOS_EMULATOR_IP_ADDRESS_OVERRIDE", socket.gethostbyname(socket.gethostname()))
140+
.with_env("AZURE_COSMOS_EMULATOR_ENABLE_DATA_PERSISTENCE", str(self.enable_data_persistence))
141+
.with_env("AZURE_COSMOS_EMULATOR_KEY", str(self.key))
142+
)
143+
144+
def _wait_until_ready(self) -> Self:
145+
"""
146+
Waits until the CosmosDB Emulator image is ready to be used.
147+
"""
148+
(
149+
self._wait_for_logs(container=self, predicate="Started\\s*$")
150+
._wait_for_url(f"{self.url}/_explorer/index.html")
151+
._wait_for_query_success(lambda sync_client: list(sync_client.list_databases()))
152+
)
153+
return self
154+
155+
@wait_container_is_ready(HTTPError, URLError)
156+
def _wait_for_url(self, url: str) -> Self:
157+
with urlopen(url, context=ssl._create_unverified_context()) as response:
158+
response.read()
159+
return self
160+
161+
def _wait_for_logs(self, *args, **kwargs) -> Self:
162+
wait_for_logs(*args, **kwargs)
163+
return self
164+
165+
@wait_container_is_ready(ServiceRequestError)
166+
def _wait_for_query_success(self, query: Callable[[SyncCosmosClient], None]) -> Self:
167+
with self.sync_client() as c:
168+
query(c)
169+
return self

‎modules/cosmosdb/tests/test_cosmosdb.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import pytest
22
from testcontainers.cosmosdb import CosmosDBEmulatorContainer
33

4+
45
def test_docker_run():
5-
with CosmosDBEmulatorContainer(partition_count=1) as cosmosdb:
6-
list(cosmosdb.sync_client().list_databases())
6+
with CosmosDBEmulatorContainer(partition_count=1) as cosmosdb:
7+
list(cosmosdb.sync_client().list_databases())

0 commit comments

Comments
 (0)