Coverage for mlos_bench/mlos_bench/tests/docker_fixtures_util.py: 95%
37 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-10-30 00:51 +0000
1#
2# Copyright (c) Microsoft Corporation.
3# Licensed under the MIT License.
4#
5"""
6Helper functions for various docker test fixtures.
8Test functions needing to use these should import them and then add them to their
9namespace in a conftest.py file.
11The intent of keeping these separate from conftest.py is to allow individual test to
12setup their own docker-compose configurations that are independent.
14As such, each conftest.py should set their own docker_compose_file fixture pointing to
15the appropriate docker-compose.yml file(s) and set a unique docker_compose_project_name.
16"""
17# pylint: disable=redefined-outer-name
19import os
20import sys
21from collections.abc import Generator
22from typing import Any
24import pytest
25from fasteners import InterProcessLock, InterProcessReaderWriterLock
26from pytest_docker.plugin import Services as DockerServices
27from pytest_docker.plugin import get_docker_services
30# Fixtures to configure the pytest-docker plugin.
31@pytest.fixture(scope="session")
32def docker_setup() -> list[str] | str:
33 """Setup for docker services."""
34 if sys.platform == "darwin" or os.environ.get("HOST_OSTYPE", "").lower().startswith("darwin"):
35 # Workaround an oddity on macOS where the "docker-compose up"
36 # command always recreates the containers.
37 # That leads to races when multiple workers are trying to
38 # start and use the same services.
39 return ["up --build -d --no-recreate"]
40 else:
41 return ["up --build -d"]
44@pytest.fixture(scope="session")
45def docker_compose_file(pytestconfig: pytest.Config) -> list[str]:
46 """
47 Fixture for the path to the docker-compose file.
49 Parameters
50 ----------
51 pytestconfig : pytest.Config
53 Returns
54 -------
55 list[str]
56 List of paths to the docker-compose file(s).
57 """
58 _ = pytestconfig # unused
59 # Add additional configs as necessary here.
60 # return []
61 raise NotImplementedError("Please implement docker_compose_file in your conftest.py")
64@pytest.fixture(scope="session")
65def docker_compose_project_name(short_testrun_uid: str) -> str:
66 """
67 Fixture for the name of the docker-compose project.
69 Returns
70 -------
71 str
72 Name of the docker-compose project.
73 """
74 # Use the xdist testrun UID to ensure that the docker-compose project name
75 # is unique across sessions, but shared amongst workers.
76 # return f"""mlos_bench-test-{short_testrun_uid}-{__name__.replace(".", "-")}"""
77 raise NotImplementedError("Please implement docker_compose_project_name in your conftest.py")
80@pytest.fixture(scope="session")
81def docker_services_lock(
82 shared_temp_dir: str,
83 short_testrun_uid: str,
84) -> InterProcessReaderWriterLock:
85 """
86 Gets a pytest session lock for xdist workers to mark when they're using the docker
87 services.
89 Yields
90 ------
91 A lock to ensure that setup/teardown operations don't happen while a
92 worker is using the docker services.
93 """
94 return InterProcessReaderWriterLock(
95 f"{shared_temp_dir}/pytest_docker_services-{short_testrun_uid}.lock"
96 )
99@pytest.fixture(scope="session")
100def docker_setup_teardown_lock(shared_temp_dir: str, short_testrun_uid: str) -> InterProcessLock:
101 """
102 Gets a pytest session lock between xdist workers for the docker setup/teardown
103 operations.
105 Yields
106 ------
107 A lock to ensure that only one worker is doing setup/teardown at a time.
108 """
109 return InterProcessLock(
110 f"{shared_temp_dir}/pytest_docker_services-setup-teardown-{short_testrun_uid}.lock"
111 )
114@pytest.fixture(scope="session")
115def locked_docker_services(
116 docker_compose_command: Any,
117 docker_compose_file: Any,
118 docker_compose_project_name: Any,
119 docker_setup: Any,
120 docker_cleanup: Any,
121 docker_setup_teardown_lock: InterProcessLock,
122 docker_services_lock: InterProcessReaderWriterLock,
123) -> Generator[DockerServices, Any, None]:
124 """A locked version of the docker_services fixture to implement xdist single
125 instance locking.
126 """
127 # pylint: disable=too-many-arguments,too-many-positional-arguments
128 # Mark the services as in use with the reader lock.
129 docker_services_lock.acquire_read_lock()
130 # Acquire the setup lock to prevent multiple setup operations at once.
131 docker_setup_teardown_lock.acquire()
132 # This "with get_docker_services(...)"" pattern is in the default fixture.
133 # We call it instead of docker_services() to avoid pytest complaints about
134 # calling fixtures directly.
135 with get_docker_services(
136 docker_compose_command,
137 docker_compose_file,
138 docker_compose_project_name,
139 docker_setup,
140 docker_cleanup,
141 ) as docker_services:
142 # Release the setup/tear down lock in order to let the setup operation
143 # continue for other workers (should be a no-op at this point).
144 docker_setup_teardown_lock.release()
145 # Yield the services so that tests within this worker can use them.
146 yield docker_services
147 # Now tests that use those services get to run on this worker...
148 # Once the tests are done, release the read lock that marks the services as in use.
149 docker_services_lock.release_read_lock()
150 # Now as we prepare to execute the cleanup code on context exit we need
151 # to acquire the setup/teardown lock again.
152 # First we attempt to get the write lock so that we wait for other
153 # readers to finish and guard against a lock inversion possibility.
154 docker_services_lock.acquire_write_lock()
155 # Next, acquire the setup/teardown lock
156 # First one here is the one to do actual work, everyone else is basically a no-op.
157 # Upon context exit, we should execute the docker_cleanup code.
158 # And try to get the setup/tear down lock again.
159 docker_setup_teardown_lock.acquire()
160 # Finally, after the docker_cleanup code has finished, remove both locks.
161 docker_setup_teardown_lock.release()
162 docker_services_lock.release_write_lock()
165__all__ = [
166 # These two should be implemented in the conftest.py of the local test suite
167 # "docker_compose_file",
168 # "docker_compose_project_name",
169 "docker_setup",
170 "docker_services_lock",
171 "docker_setup_teardown_lock",
172 "locked_docker_services",
173]