Coverage for mlos_bench/mlos_bench/tests/conftest.py: 96%
57 statements
« prev ^ index » next coverage.py v7.9.2, created at 2025-07-14 00:55 +0000
1#
2# Copyright (c) Microsoft Corporation.
3# Licensed under the MIT License.
4#
5"""Common fixtures for mock TunableGroups and Environment objects."""
7import os
8import sys
9from collections.abc import Generator
10from typing import Any
12import pytest
13from fasteners import InterProcessLock, InterProcessReaderWriterLock
14from pytest_docker.plugin import Services as DockerServices
15from pytest_docker.plugin import get_docker_services
17from mlos_bench.environments.mock_env import MockEnv
18from mlos_bench.tests import SEED, resolve_host_name, tunable_groups_fixtures
19from mlos_bench.tunables.tunable_groups import TunableGroups
21# pylint: disable=redefined-outer-name
22# -- Ignore pylint complaints about pytest references to
23# `tunable_groups` fixture as both a function and a parameter.
25# Expose some of those as local names so they can be picked up as fixtures by pytest.
26tunable_groups_config = tunable_groups_fixtures.tunable_groups_config
27tunable_groups = tunable_groups_fixtures.tunable_groups
28mixed_numerics_tunable_groups = tunable_groups_fixtures.mixed_numerics_tunable_groups
29covariant_group = tunable_groups_fixtures.covariant_group
32HOST_DOCKER_NAME = "host.docker.internal"
35@pytest.fixture(scope="session")
36def docker_hostname() -> str:
37 """Returns the local hostname to use to connect to the test ssh server."""
38 if sys.platform != "win32" and resolve_host_name(HOST_DOCKER_NAME):
39 # On Linux, if we're running in a docker container, we can use the
40 # --add-host (extra_hosts in docker-compose.yml) to refer to the host IP.
41 return HOST_DOCKER_NAME
42 # Docker (Desktop) for Windows (WSL2) uses a special networking magic
43 # to refer to the host machine as `localhost` when exposing ports.
44 # In all other cases, assume we're executing directly inside conda on the host.
45 return "127.0.0.1" # "localhost"
48@pytest.fixture
49def mock_env(tunable_groups: TunableGroups) -> MockEnv:
50 """Test fixture for MockEnv."""
51 return MockEnv(
52 name="Test Env",
53 config={
54 "tunable_params": ["provision", "boot", "kernel"],
55 "mock_env_seed": SEED,
56 "mock_env_range": [60, 120],
57 "mock_env_metrics": ["score"],
58 },
59 tunables=tunable_groups,
60 )
63@pytest.fixture
64def mock_env_no_noise(tunable_groups: TunableGroups) -> MockEnv:
65 """Test fixture for MockEnv."""
66 return MockEnv(
67 name="Test Env No Noise",
68 config={
69 "tunable_params": ["provision", "boot", "kernel"],
70 "mock_env_seed": -1,
71 "mock_env_range": [60, 120],
72 "mock_env_metrics": ["score", "other_score"],
73 },
74 tunables=tunable_groups,
75 )
78# Fixtures to configure the pytest-docker plugin.
79@pytest.fixture(scope="session")
80def docker_setup() -> list[str] | str:
81 """Setup for docker services."""
82 if sys.platform == "darwin" or os.environ.get("HOST_OSTYPE", "").lower().startswith("darwin"):
83 # Workaround an oddity on macOS where the "docker-compose up"
84 # command always recreates the containers.
85 # That leads to races when multiple workers are trying to
86 # start and use the same services.
87 return ["up --build -d --no-recreate"]
88 else:
89 return ["up --build -d"]
92@pytest.fixture(scope="session")
93def docker_compose_file(pytestconfig: pytest.Config) -> list[str]:
94 """
95 Returns the path to the docker-compose file.
97 Parameters
98 ----------
99 pytestconfig : pytest.Config
101 Returns
102 -------
103 str
104 Path to the docker-compose file.
105 """
106 _ = pytestconfig # unused
107 return [
108 os.path.join(os.path.dirname(__file__), "services", "remote", "ssh", "docker-compose.yml"),
109 os.path.join(os.path.dirname(__file__), "storage", "sql", "docker-compose.yml"),
110 # Add additional configs as necessary here.
111 ]
114@pytest.fixture(scope="session")
115def docker_compose_project_name(short_testrun_uid: str) -> str:
116 """
117 Returns the name of the docker-compose project.
119 Returns
120 -------
121 str
122 Name of the docker-compose project.
123 """
124 # Use the xdist testrun UID to ensure that the docker-compose project name
125 # is unique across sessions, but shared amongst workers.
126 return f"mlos_bench-test-{short_testrun_uid}"
129@pytest.fixture(scope="session")
130def docker_services_lock(
131 shared_temp_dir: str,
132 short_testrun_uid: str,
133) -> InterProcessReaderWriterLock:
134 """
135 Gets a pytest session lock for xdist workers to mark when they're using the docker
136 services.
138 Yields
139 ------
140 A lock to ensure that setup/teardown operations don't happen while a
141 worker is using the docker services.
142 """
143 return InterProcessReaderWriterLock(
144 f"{shared_temp_dir}/pytest_docker_services-{short_testrun_uid}.lock"
145 )
148@pytest.fixture(scope="session")
149def docker_setup_teardown_lock(shared_temp_dir: str, short_testrun_uid: str) -> InterProcessLock:
150 """
151 Gets a pytest session lock between xdist workers for the docker setup/teardown
152 operations.
154 Yields
155 ------
156 A lock to ensure that only one worker is doing setup/teardown at a time.
157 """
158 return InterProcessLock(
159 f"{shared_temp_dir}/pytest_docker_services-setup-teardown-{short_testrun_uid}.lock"
160 )
163@pytest.fixture(scope="session")
164def locked_docker_services(
165 docker_compose_command: Any,
166 docker_compose_file: Any,
167 docker_compose_project_name: Any,
168 docker_setup: Any,
169 docker_cleanup: Any,
170 docker_setup_teardown_lock: InterProcessLock,
171 docker_services_lock: InterProcessReaderWriterLock,
172) -> Generator[DockerServices, Any, None]:
173 """A locked version of the docker_services fixture to implement xdist single
174 instance locking.
175 """
176 # pylint: disable=too-many-arguments,too-many-positional-arguments
177 # Mark the services as in use with the reader lock.
178 docker_services_lock.acquire_read_lock()
179 # Acquire the setup lock to prevent multiple setup operations at once.
180 docker_setup_teardown_lock.acquire()
181 # This "with get_docker_services(...)"" pattern is in the default fixture.
182 # We call it instead of docker_services() to avoid pytest complaints about
183 # calling fixtures directly.
184 with get_docker_services(
185 docker_compose_command,
186 docker_compose_file,
187 docker_compose_project_name,
188 docker_setup,
189 docker_cleanup,
190 ) as docker_services:
191 # Release the setup/tear down lock in order to let the setup operation
192 # continue for other workers (should be a no-op at this point).
193 docker_setup_teardown_lock.release()
194 # Yield the services so that tests within this worker can use them.
195 yield docker_services
196 # Now tests that use those services get to run on this worker...
197 # Once the tests are done, release the read lock that marks the services as in use.
198 docker_services_lock.release_read_lock()
199 # Now as we prepare to execute the cleanup code on context exit we need
200 # to acquire the setup/teardown lock again.
201 # First we attempt to get the write lock so that we wait for other
202 # readers to finish and guard against a lock inversion possibility.
203 docker_services_lock.acquire_write_lock()
204 # Next, acquire the setup/teardown lock
205 # First one here is the one to do actual work, everyone else is basically a no-op.
206 # Upon context exit, we should execute the docker_cleanup code.
207 # And try to get the setup/tear down lock again.
208 docker_setup_teardown_lock.acquire()
209 # Finally, after the docker_cleanup code has finished, remove both locks.
210 docker_setup_teardown_lock.release()
211 docker_services_lock.release_write_lock()