Skip to content

Commit

Permalink
linting
Browse files Browse the repository at this point in the history
  • Loading branch information
Mehdi BEN ABDALLAH committed May 26, 2024
1 parent cae372b commit 67f5220
Show file tree
Hide file tree
Showing 7 changed files with 44 additions and 33 deletions.
30 changes: 16 additions & 14 deletions modules/cosmosdb/testcontainers/cosmosdb/_emulator.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,18 +2,22 @@
import socket
import ssl
from collections.abc import Iterable
from distutils.util import strtobool
from urllib.error import HTTPError, URLError
from urllib.request import urlopen

from typing_extensions import Self

from testcontainers.core.container import DockerContainer
from testcontainers.core.waiting_utils import wait_container_is_ready, wait_for_logs

from . import _grab as grab
from distutils.util import strtobool
from urllib.error import HTTPError, URLError
from urllib.request import urlopen

__all__ = ["CosmosDBEmulatorContainer"]

EMULATOR_PORT = 8081


class CosmosDBEmulatorContainer(DockerContainer):
"""
Abstract class for CosmosDB Emulator endpoints.
Expand All @@ -28,9 +32,7 @@ def __init__(
"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 = strtobool(
os.getenv("AZURE_COSMOS_EMULATOR_ENABLE_DATA_PERSISTENCE", "false")
),
enable_data_persistence: bool = strtobool(os.getenv("AZURE_COSMOS_EMULATOR_ENABLE_DATA_PERSISTENCE", "false")),
key: str = os.getenv(
"AZURE_COSMOS_EMULATOR_KEY",
"C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==",
Expand All @@ -52,7 +54,7 @@ def host(self) -> str:
Emulator host
"""
return self.get_container_host_ip()

@property
def server_certificate_pem(self) -> bytes:
"""
Expand All @@ -66,18 +68,17 @@ def start(self) -> Self:
self._wait_until_ready()
self._cert_pem_bytes = self._download_cert()
return self

def _configure(self) -> None:
all_ports = set([EMULATOR_PORT] + self.endpoint_ports)
all_ports = {EMULATOR_PORT, *self.endpoint_ports}
if self.bind_ports:
for port in all_ports:
self.with_bind_ports(port, port)
else:
self.with_exposed_ports(*all_ports)

(
self
.with_env("AZURE_COSMOS_EMULATOR_PARTITION_COUNT", str(self.partition_count))
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))
Expand All @@ -89,12 +90,13 @@ def _wait_until_ready(self) -> Self:
if self.bind_ports:
self._wait_for_url(f"https://{self.host}:{EMULATOR_PORT}/_explorer/index.html")
self._wait_for_query_success()

return self

def _download_cert(self) -> bytes:
with grab.file(
self.get_wrapped_container(), "/tmp/cosmos/appdata/.system/profiles/Client/AppData/Local/CosmosDBEmulator/emulator.pem"
self.get_wrapped_container(),
"/tmp/cosmos/appdata/.system/profiles/Client/AppData/Local/CosmosDBEmulator/emulator.pem",
) as cert:
return cert.read()

Expand All @@ -103,6 +105,6 @@ 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_query_success(self) -> None:
pass
32 changes: 16 additions & 16 deletions modules/cosmosdb/testcontainers/cosmosdb/_grab.py
Original file line number Diff line number Diff line change
@@ -1,26 +1,26 @@

from pathlib import Path
from os import path
import tarfile
import tempfile
from contextlib import contextmanager
from os import path
from pathlib import Path

from docker.models.containers import Container


@contextmanager
def file(container: Container, target: str):
target_path = Path(target)
assert target_path.is_absolute(), "target must be an absolute path"

with tempfile.TemporaryDirectory() as tmp:
archive = Path(tmp) / 'grabbed.tar'
target_path = Path(target)
assert target_path.is_absolute(), "target must be an absolute path"

# download from container as tar archive
with open(archive, 'wb') as f:
tar_bits, _ = container.get_archive(target)
for chunk in tar_bits:
f.write(chunk)
with tempfile.TemporaryDirectory() as tmp:
archive = Path(tmp) / "grabbed.tar"

# extract target file from tar archive
with tarfile.TarFile(archive) as tar:
yield tar.extractfile(path.basename(target))
# download from container as tar archive
with open(archive, "wb") as f:
tar_bits, _ = container.get_archive(target)
for chunk in tar_bits:
f.write(chunk)

# extract target file from tar archive
with tarfile.TarFile(archive) as tar:
yield tar.extractfile(path.basename(target))
6 changes: 4 additions & 2 deletions modules/cosmosdb/testcontainers/cosmosdb/mongodb.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,12 @@
import os

from ._emulator import CosmosDBEmulatorContainer

__all__ = ["CosmosDBMongoEndpointContainer"]

ENDPOINT_PORT = 10255


class CosmosDBMongoEndpointContainer(CosmosDBEmulatorContainer):
"""
CosmosDB MongoDB enpoint Emulator.
Expand All @@ -21,7 +23,7 @@ class CosmosDBMongoEndpointContainer(CosmosDBEmulatorContainer):

def __init__(
self,
mongodb_version: str = None,
mongodb_version: str,
image: str = os.getenv(
"AZURE_COSMOS_EMULATOR_IMAGE", "mcr.microsoft.com/cosmosdb/linux/azure-cosmos-emulator:mongodb"
),
Expand All @@ -37,7 +39,7 @@ def port(self) -> str:
The exposed port to the MongoDB endpoint
"""
return self.get_exposed_port(ENDPOINT_PORT)

def _configure(self) -> None:
super()._configure()
self.with_env("AZURE_COSMOS_EMULATOR_ENABLE_MONGODB_ENDPOINT", self.mongodb_version)
3 changes: 3 additions & 0 deletions modules/cosmosdb/testcontainers/cosmosdb/nosql.py
Original file line number Diff line number Diff line change
@@ -1,13 +1,16 @@
from azure.core.exceptions import ServiceRequestError
from azure.cosmos import CosmosClient as SyncCosmosClient
from azure.cosmos.aio import CosmosClient as AsyncCosmosClient

from testcontainers.core.waiting_utils import wait_container_is_ready

from ._emulator import CosmosDBEmulatorContainer

__all__ = ["CosmosDBNoSQLEndpointContainer"]

NOSQL_PORT = 8081


class CosmosDBNoSQLEndpointContainer(CosmosDBEmulatorContainer):
"""
CosmosDB NoSQL enpoint Emulator.
Expand Down
1 change: 1 addition & 0 deletions modules/cosmosdb/tests/test_emulator.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import pytest
from testcontainers.cosmosdb._emulator import CosmosDBEmulatorContainer


def test_runs():
with CosmosDBEmulatorContainer(partition_count=1, bind_ports=False) as emulator:
assert emulator.server_certificate_pem is not None
Expand Down
4 changes: 3 additions & 1 deletion modules/cosmosdb/tests/test_mongodb.py
Original file line number Diff line number Diff line change
@@ -1,13 +1,15 @@
import pytest
from testcontainers.cosmosdb import CosmosDBMongoEndpointContainer


def test_requires_a_version():
with pytest.raises(AssertionError, match="A MongoDB version is required"):
CosmosDBMongoEndpointContainer()
CosmosDBMongoEndpointContainer(mongodb_version=None)

# instanciates
CosmosDBMongoEndpointContainer(mongodb_version="4.0")


def test_runs():
with CosmosDBMongoEndpointContainer(mongodb_version="4.0", partition_count=1, bind_ports=False) as emulator:
assert emulator.env["AZURE_COSMOS_EMULATOR_ENABLE_MONGODB_ENDPOINT"] == "4.0"
Expand Down
1 change: 1 addition & 0 deletions modules/cosmosdb/tests/test_nosql.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import pytest
from testcontainers.cosmosdb import CosmosDBNoSQLEndpointContainer


def test_runs():
with CosmosDBNoSQLEndpointContainer(partition_count=1, bind_ports=False) as emulator:
assert emulator.get_exposed_port(8081) is not None, "The NoSQL endpoint's port should be exposed"

0 comments on commit 67f5220

Please sign in to comment.