Coverage for mlos_bench/mlos_bench/tests/services/remote/ssh/fixtures.py: 97%
38 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"""
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 ssh_test_server(
34 docker_hostname: str,
35 docker_compose_project_name: str,
36 locked_docker_services: DockerServices,
37) -> Generator[SshTestServerInfo]:
38 """
39 Fixture for getting the ssh test server services setup via docker-compose using
40 pytest-docker.
42 Yields the (hostname, port, username, id_rsa_path) of the test server.
44 Once the session is over, the docker containers are torn down, and the temporary
45 file holding the dynamically generated private key of the test server is deleted.
46 """
47 # Get a copy of the ssh id_rsa key from the test ssh server.
48 with tempfile.NamedTemporaryFile() as id_rsa_file:
49 ssh_test_server_info = SshTestServerInfo(
50 compose_project_name=docker_compose_project_name,
51 service_name=SSH_TEST_SERVER_NAME,
52 hostname=docker_hostname,
53 username="root",
54 id_rsa_path=id_rsa_file.name,
55 )
56 wait_docker_service_socket(
57 locked_docker_services,
58 ssh_test_server_info.hostname,
59 ssh_test_server_info.get_port(),
60 )
61 id_rsa_src = f"/{ssh_test_server_info.username}/.ssh/id_rsa"
62 docker_cp_cmd = (
63 f"docker compose -p {docker_compose_project_name} "
64 f"cp {SSH_TEST_SERVER_NAME}:{id_rsa_src} {id_rsa_file.name}"
65 )
66 cmd = run(
67 docker_cp_cmd.split(),
68 check=True,
69 cwd=os.path.dirname(__file__),
70 capture_output=True,
71 text=True,
72 )
73 if cmd.returncode != 0:
74 raise RuntimeError(
75 f"Failed to copy ssh key from {SSH_TEST_SERVER_NAME} container "
76 + f"[return={cmd.returncode}]: {str(cmd.stderr)}"
77 )
78 os.chmod(id_rsa_file.name, 0o600)
79 yield ssh_test_server_info
80 # NamedTempFile deleted on context exit
83@pytest.fixture(scope="session")
84def alt_test_server(
85 ssh_test_server: SshTestServerInfo,
86 locked_docker_services: DockerServices,
87) -> SshTestServerInfo:
88 """
89 Fixture for getting the second ssh test server info from the docker-compose.yml.
91 See additional notes in the ssh_test_server fixture above.
92 """
93 # Note: The alt-server uses the same image as the ssh-server container, so
94 # the id_rsa key and username should all match.
95 # Only the host port it is allocate is different.
96 alt_test_server_info = SshTestServerInfo(
97 compose_project_name=ssh_test_server.compose_project_name,
98 service_name=ALT_TEST_SERVER_NAME,
99 hostname=ssh_test_server.hostname,
100 username=ssh_test_server.username,
101 id_rsa_path=ssh_test_server.id_rsa_path,
102 )
103 wait_docker_service_socket(
104 locked_docker_services,
105 alt_test_server_info.hostname,
106 alt_test_server_info.get_port(),
107 )
108 return alt_test_server_info
111@pytest.fixture(scope="session")
112def reboot_test_server(
113 ssh_test_server: SshTestServerInfo,
114 locked_docker_services: DockerServices,
115) -> SshTestServerInfo:
116 """
117 Fixture for getting the third ssh test server info from the docker-compose.yml.
119 See additional notes in the ssh_test_server fixture above.
120 """
121 # Note: The reboot-server uses the same image as the ssh-server container, so
122 # the id_rsa key and username should all match.
123 # Only the host port it is allocate is different.
124 reboot_test_server_info = SshTestServerInfo(
125 compose_project_name=ssh_test_server.compose_project_name,
126 service_name=REBOOT_TEST_SERVER_NAME,
127 hostname=ssh_test_server.hostname,
128 username=ssh_test_server.username,
129 id_rsa_path=ssh_test_server.id_rsa_path,
130 )
131 wait_docker_service_socket(
132 locked_docker_services,
133 reboot_test_server_info.hostname,
134 reboot_test_server_info.get_port(),
135 )
136 return reboot_test_server_info
139@pytest.fixture
140def ssh_host_service(ssh_test_server: SshTestServerInfo) -> SshHostService:
141 """Generic SshHostService fixture."""
142 return SshHostService(
143 config={
144 "ssh_username": ssh_test_server.username,
145 "ssh_priv_key_path": ssh_test_server.id_rsa_path,
146 },
147 )
150@pytest.fixture
151def ssh_fileshare_service() -> SshFileShareService:
152 """Generic SshFileShareService fixture."""
153 return SshFileShareService(
154 config={
155 # Left blank to make sure we test per connection overrides.
156 },
157 )