Coverage for mlos_bench/mlos_bench/tests/services/remote/ssh/fixtures.py: 98%
45 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"""
6Fixtures for the SSH service tests.
8Note: these are not in the conftest.py file because they are also used by remote_ssh_env_test.py
9"""
11import os
12import tempfile
13from collections.abc import Generator
14from subprocess import run
16import pytest
17from pytest_docker.plugin import Services as DockerServices
19from mlos_bench.services.remote.ssh.ssh_fileshare import SshFileShareService
20from mlos_bench.services.remote.ssh.ssh_host_service import SshHostService
21from mlos_bench.tests import wait_docker_service_socket
22from mlos_bench.tests.services.remote.ssh import (
23 ALT_TEST_SERVER_NAME,
24 REBOOT_TEST_SERVER_NAME,
25 SSH_TEST_SERVER_NAME,
26 SshTestServerInfo,
27)
29# pylint: disable=redefined-outer-name
32@pytest.fixture(scope="session")
33def docker_compose_file(pytestconfig: pytest.Config) -> list[str]:
34 """
35 Fixture for the path to the docker-compose file.
37 Parameters
38 ----------
39 pytestconfig : pytest.Config
41 Returns
42 -------
43 list[str]
44 List of paths to docker-compose files.
45 """
46 _ = pytestconfig # unused
47 return [
48 os.path.join(os.path.dirname(__file__), "docker-compose.yml"),
49 ]
52@pytest.fixture(scope="session")
53def docker_compose_project_name(short_testrun_uid: str) -> str:
54 """
55 Returns the name of the docker-compose project.
57 Returns
58 -------
59 str
60 Name of the docker-compose project.
61 """
62 # Use the xdist testrun UID to ensure that the docker-compose project name
63 # is unique across sessions, but shared amongst workers.
64 return f"""mlos_bench-test-{short_testrun_uid}-{__name__.replace(".", "-")}"""
67@pytest.fixture(scope="session")
68def ssh_test_server(
69 docker_hostname: str,
70 docker_compose_project_name: str,
71 locked_docker_services: DockerServices,
72) -> Generator[SshTestServerInfo]:
73 """
74 Fixture for getting the ssh test server services setup via docker-compose using
75 pytest-docker.
77 Yields the (hostname, port, username, id_rsa_path) of the test server.
79 Once the session is over, the docker containers are torn down, and the temporary
80 file holding the dynamically generated private key of the test server is deleted.
81 """
82 # Get a copy of the ssh id_rsa key from the test ssh server.
83 with tempfile.NamedTemporaryFile() as id_rsa_file:
84 ssh_test_server_info = SshTestServerInfo(
85 compose_project_name=docker_compose_project_name,
86 service_name=SSH_TEST_SERVER_NAME,
87 hostname=docker_hostname,
88 username="root",
89 id_rsa_path=id_rsa_file.name,
90 )
91 wait_docker_service_socket(
92 locked_docker_services,
93 ssh_test_server_info.hostname,
94 ssh_test_server_info.get_port(),
95 )
96 id_rsa_src = f"/{ssh_test_server_info.username}/.ssh/id_rsa"
97 docker_cp_cmd = (
98 f"docker compose -p {docker_compose_project_name} "
99 f"cp {SSH_TEST_SERVER_NAME}:{id_rsa_src} {id_rsa_file.name}"
100 )
101 cmd = run(
102 docker_cp_cmd.split(),
103 check=True,
104 cwd=os.path.dirname(__file__),
105 capture_output=True,
106 text=True,
107 )
108 if cmd.returncode != 0:
109 raise RuntimeError(
110 f"Failed to copy ssh key from {SSH_TEST_SERVER_NAME} container "
111 + f"[return={cmd.returncode}]: {str(cmd.stderr)}"
112 )
113 os.chmod(id_rsa_file.name, 0o600)
114 yield ssh_test_server_info
115 # NamedTempFile deleted on context exit
118@pytest.fixture(scope="session")
119def alt_test_server(
120 ssh_test_server: SshTestServerInfo,
121 locked_docker_services: DockerServices,
122) -> SshTestServerInfo:
123 """
124 Fixture for getting the second ssh test server info from the docker-compose.yml.
126 See additional notes in the ssh_test_server fixture above.
127 """
128 # Note: The alt-server uses the same image as the ssh-server container, so
129 # the id_rsa key and username should all match.
130 # Only the host port it is allocate is different.
131 alt_test_server_info = SshTestServerInfo(
132 compose_project_name=ssh_test_server.compose_project_name,
133 service_name=ALT_TEST_SERVER_NAME,
134 hostname=ssh_test_server.hostname,
135 username=ssh_test_server.username,
136 id_rsa_path=ssh_test_server.id_rsa_path,
137 )
138 wait_docker_service_socket(
139 locked_docker_services,
140 alt_test_server_info.hostname,
141 alt_test_server_info.get_port(),
142 )
143 return alt_test_server_info
146@pytest.fixture(scope="session")
147def reboot_test_server(
148 ssh_test_server: SshTestServerInfo,
149 locked_docker_services: DockerServices,
150) -> SshTestServerInfo:
151 """
152 Fixture for getting the third ssh test server info from the docker-compose.yml.
154 See additional notes in the ssh_test_server fixture above.
155 """
156 # Note: The reboot-server uses the same image as the ssh-server container, so
157 # the id_rsa key and username should all match.
158 # Only the host port it is allocate is different.
159 reboot_test_server_info = SshTestServerInfo(
160 compose_project_name=ssh_test_server.compose_project_name,
161 service_name=REBOOT_TEST_SERVER_NAME,
162 hostname=ssh_test_server.hostname,
163 username=ssh_test_server.username,
164 id_rsa_path=ssh_test_server.id_rsa_path,
165 )
166 wait_docker_service_socket(
167 locked_docker_services,
168 reboot_test_server_info.hostname,
169 reboot_test_server_info.get_port(),
170 )
171 return reboot_test_server_info
174@pytest.fixture
175def ssh_host_service(ssh_test_server: SshTestServerInfo) -> SshHostService:
176 """Generic SshHostService fixture."""
177 return SshHostService(
178 config={
179 "ssh_username": ssh_test_server.username,
180 "ssh_priv_key_path": ssh_test_server.id_rsa_path,
181 },
182 )
185@pytest.fixture
186def ssh_fileshare_service() -> SshFileShareService:
187 """Generic SshFileShareService fixture."""
188 return SshFileShareService(
189 config={
190 # Left blank to make sure we test per connection overrides.
191 },
192 )