Coverage for mlos_bench/mlos_bench/tests/conftest.py: 100%
44 statements
« prev ^ index » next coverage.py v7.6.1, created at 2024-10-07 01:52 +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
8from typing import Any, Generator, List
10import pytest
11from fasteners import InterProcessLock, InterProcessReaderWriterLock
12from pytest_docker.plugin import Services as DockerServices
13from pytest_docker.plugin import get_docker_services
15from mlos_bench.environments.mock_env import MockEnv
16from mlos_bench.tests import SEED, tunable_groups_fixtures
17from mlos_bench.tunables.tunable_groups import TunableGroups
19# pylint: disable=redefined-outer-name
20# -- Ignore pylint complaints about pytest references to
21# `tunable_groups` fixture as both a function and a parameter.
23# Expose some of those as local names so they can be picked up as fixtures by pytest.
24tunable_groups_config = tunable_groups_fixtures.tunable_groups_config
25tunable_groups = tunable_groups_fixtures.tunable_groups
26mixed_numerics_tunable_groups = tunable_groups_fixtures.mixed_numerics_tunable_groups
27covariant_group = tunable_groups_fixtures.covariant_group
30@pytest.fixture
31def mock_env(tunable_groups: TunableGroups) -> MockEnv:
32 """Test fixture for MockEnv."""
33 return MockEnv(
34 name="Test Env",
35 config={
36 "tunable_params": ["provision", "boot", "kernel"],
37 "mock_env_seed": SEED,
38 "mock_env_range": [60, 120],
39 "mock_env_metrics": ["score"],
40 },
41 tunables=tunable_groups,
42 )
45@pytest.fixture
46def mock_env_no_noise(tunable_groups: TunableGroups) -> MockEnv:
47 """Test fixture for MockEnv."""
48 return MockEnv(
49 name="Test Env No Noise",
50 config={
51 "tunable_params": ["provision", "boot", "kernel"],
52 "mock_env_seed": -1,
53 "mock_env_range": [60, 120],
54 "mock_env_metrics": ["score", "other_score"],
55 },
56 tunables=tunable_groups,
57 )
60# Fixtures to configure the pytest-docker plugin.
63@pytest.fixture(scope="session")
64def docker_compose_file(pytestconfig: pytest.Config) -> List[str]:
65 """
66 Returns the path to the docker-compose file.
68 Parameters
69 ----------
70 pytestconfig : pytest.Config
72 Returns
73 -------
74 str
75 Path to the docker-compose file.
76 """
77 _ = pytestconfig # unused
78 return [
79 os.path.join(os.path.dirname(__file__), "services", "remote", "ssh", "docker-compose.yml"),
80 # Add additional configs as necessary here.
81 ]
84@pytest.fixture(scope="session")
85def docker_compose_project_name(short_testrun_uid: str) -> str:
86 """
87 Returns the name of the docker-compose project.
89 Returns
90 -------
91 str
92 Name of the docker-compose project.
93 """
94 # Use the xdist testrun UID to ensure that the docker-compose project name
95 # is unique across sessions, but shared amongst workers.
96 return f"mlos_bench-test-{short_testrun_uid}"
99@pytest.fixture(scope="session")
100def docker_services_lock(
101 shared_temp_dir: str,
102 short_testrun_uid: str,
103) -> InterProcessReaderWriterLock:
104 """
105 Gets a pytest session lock for xdist workers to mark when they're using the docker
106 services.
108 Yields
109 ------
110 A lock to ensure that setup/teardown operations don't happen while a
111 worker is using the docker services.
112 """
113 return InterProcessReaderWriterLock(
114 f"{shared_temp_dir}/pytest_docker_services-{short_testrun_uid}.lock"
115 )
118@pytest.fixture(scope="session")
119def docker_setup_teardown_lock(shared_temp_dir: str, short_testrun_uid: str) -> InterProcessLock:
120 """
121 Gets a pytest session lock between xdist workers for the docker setup/teardown
122 operations.
124 Yields
125 ------
126 A lock to ensure that only one worker is doing setup/teardown at a time.
127 """
128 return InterProcessLock(
129 f"{shared_temp_dir}/pytest_docker_services-setup-teardown-{short_testrun_uid}.lock"
130 )
133@pytest.fixture(scope="session")
134def locked_docker_services(
135 docker_compose_command: Any,
136 docker_compose_file: Any,
137 docker_compose_project_name: Any,
138 docker_setup: Any,
139 docker_cleanup: Any,
140 docker_setup_teardown_lock: InterProcessLock,
141 docker_services_lock: InterProcessReaderWriterLock,
142) -> Generator[DockerServices, Any, None]:
143 """A locked version of the docker_services fixture to implement xdist single
144 instance locking.
145 """
146 # pylint: disable=too-many-arguments
147 # Mark the services as in use with the reader lock.
148 docker_services_lock.acquire_read_lock()
149 # Acquire the setup lock to prevent multiple setup operations at once.
150 docker_setup_teardown_lock.acquire()
151 # This "with get_docker_services(...)"" pattern is in the default fixture.
152 # We call it instead of docker_services() to avoid pytest complaints about
153 # calling fixtures directly.
154 with get_docker_services(
155 docker_compose_command,
156 docker_compose_file,
157 docker_compose_project_name,
158 docker_setup,
159 docker_cleanup,
160 ) as docker_services:
161 # Release the setup/tear down lock in order to let the setup operation
162 # continue for other workers (should be a no-op at this point).
163 docker_setup_teardown_lock.release()
164 # Yield the services so that tests within this worker can use them.
165 yield docker_services
166 # Now tests that use those services get to run on this worker...
167 # Once the tests are done, release the read lock that marks the services as in use.
168 docker_services_lock.release_read_lock()
169 # Now as we prepare to execute the cleanup code on context exit we need
170 # to acquire the setup/teardown lock again.
171 # First we attempt to get the write lock so that we wait for other
172 # readers to finish and guard against a lock inversion possibility.
173 docker_services_lock.acquire_write_lock()
174 # Next, acquire the setup/teardown lock
175 # First one here is the one to do actual work, everyone else is basically a no-op.
176 # Upon context exit, we should execute the docker_cleanup code.
177 # And try to get the setup/tear down lock again.
178 docker_setup_teardown_lock.acquire()
179 # Finally, after the docker_cleanup code has finished, remove both locks.
180 docker_setup_teardown_lock.release()
181 docker_services_lock.release_write_lock()