Skip to content

Commit

Permalink
azure isolate
Browse files Browse the repository at this point in the history
  • Loading branch information
provinzkraut committed Dec 22, 2024
1 parent cf00c53 commit a366b46
Show file tree
Hide file tree
Showing 2 changed files with 54 additions and 38 deletions.
84 changes: 46 additions & 38 deletions src/pytest_databases/docker/azure_blob.py
Original file line number Diff line number Diff line change
@@ -1,75 +1,81 @@
from __future__ import annotations

import json
import subprocess # noqa: S404
from dataclasses import dataclass
from typing import TYPE_CHECKING, AsyncGenerator, Generator

import pytest
from azure.storage.blob import ContainerClient
from azure.storage.blob.aio import ContainerClient as AsyncContainerClient

from pytest_databases.helpers import get_xdist_worker_num
from pytest_databases.types import ServiceContainer
from pytest_databases.helpers import get_xdist_worker_count, get_xdist_worker_num
from pytest_databases.types import ServiceContainer, XdistIsolationLevel

if TYPE_CHECKING:
from pytest_databases._service import DockerService


DEFAULT_ACCOUNT_KEY = "Eby8vdM02xNOcqFlqUwJPLlmEtlCDXJ1OUzFT50uSRZ6IFsuFq2UVErCz4I6tq/K1SZFPTOtr/KBHBeksoGMGw=="
DEFAULT_ACCOUNT_NAME = "devstoreaccount1"


@dataclass
class AzureBlobService(ServiceContainer):
connection_string: str
account_url: str
account_key: str = "Eby8vdM02xNOcqFlqUwJPLlmEtlCDXJ1OUzFT50uSRZ6IFsuFq2UVErCz4I6tq/K1SZFPTOtr/KBHBeksoGMGw=="
account_name: str = "devstoreaccount1"


def _get_container_ids(compose_file_name: str) -> list[str]:
proc = subprocess.run(
[ # noqa: S607
"docker",
"container",
"ls",
f"--filter=label=com.docker.compose.project.config_files={compose_file_name}",
"--format=json",
],
capture_output=True,
text=True,
check=True,
)
return [json.loads(line)["ID"] for line in proc.stdout.splitlines()]


def _get_container_logs(container_id: str) -> str:
return subprocess.run(
["docker", "logs", container_id], # noqa: S607
capture_output=True,
text=True,
check=True,
).stdout
account_key: str
account_name: str


@pytest.fixture(scope="session")
def azure_blob_xdist_isolation_level() -> XdistIsolationLevel:
return "database"


@pytest.fixture(scope="session")
def azure_blob_service_startup_delay() -> int:
return 1
def azurite_in_memory() -> bool:
return True


def _create_account_options(number: int) -> list[tuple[str, str]]:
return [(f"test_account_{i}", DEFAULT_ACCOUNT_KEY) for i in range(number)]


@pytest.fixture(scope="session")
def azure_blob_service(
docker_service: DockerService,
azurite_in_memory: bool,
azure_blob_xdist_isolation_level: XdistIsolationLevel,
) -> Generator[ServiceContainer, None, None]:
command = "azurite-blob --blobHost 0.0.0.0 --blobPort 10000"
if azurite_in_memory:
command += " --inMemoryPersistence"

name = "azurite-blob"
env = {}
account_name = DEFAULT_ACCOUNT_NAME
account_key = DEFAULT_ACCOUNT_KEY

worker_num = get_xdist_worker_num()
if azure_blob_xdist_isolation_level == "server":
name = f"{name}_{worker_num}"
else:
accounts = _create_account_options(get_xdist_worker_count())
env["AZURITE_ACCOUNTS"] = ",".join(f"{name}:{key}" for name, key in accounts)
account_name, account_key = accounts[worker_num - 1]

with docker_service.run(
image="mcr.microsoft.com/azure-storage/azurite",
name="azurite-blob",
command="azurite-blob --blobHost 0.0.0.0 --blobPort 10000",
name=name,
command=command,
wait_for_log="Azurite Blob service successfully listens on",
container_port=10000,
env=env,
) as service:
account_url = f"http://127.0.0.1:{service.port}/devstoreaccount1"
account_url = f"http://127.0.0.1:{service.port}/{account_name}"
connection_string = (
"DefaultEndpointsProtocol=http;"
"AccountName=devstoreaccount1;"
"AccountKey=Eby8vdM02xNOcqFlqUwJPLlmEtlCDXJ1OUzFT50uSRZ6IFsuFq2UVErCz4I6tq/K1SZFPTOtr/KBHBeksoGMGw==;"
f"AccountName={account_name};"
f"AccountKey={account_key};"
f"BlobEndpoint={account_url};"
)

Expand All @@ -78,6 +84,8 @@ def azure_blob_service(
port=service.port,
connection_string=connection_string,
account_url=account_url,
account_key=account_key,
account_name=account_name,
)


Expand Down
8 changes: 8 additions & 0 deletions src/pytest_databases/helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,3 +31,11 @@ def get_xdist_worker_num() -> int:
if worker_id is None or worker_id == "master":
return 0
return int(worker_id.replace("gw", ""))


def get_xdist_worker_count() -> int:
return int(os.getenv("PYTEST_XDIST_WORKER_COUNT", "0"))


def is_xdist() -> bool:
return get_xdist_worker_count() > 0

0 comments on commit a366b46

Please sign in to comment.