Coverage for mlos_bench/mlos_bench/tests/services/remote/ssh/fixtures.py: 96%

1# 

2# Copyright (c) Microsoft Corporation. 

3# Licensed under the MIT License. 

4# 

5""" 

6Fixtures for the SSH service tests. 

7 

8Note: these are not in the conftest.py file because they are also used by remote_ssh_env_test.py 

9""" 

10 

11from typing import Generator 

12from subprocess import run 

13 

14import os 

15import sys 

16import tempfile 

17 

18import pytest 

19from pytest_docker.plugin import Services as DockerServices 

20 

21from mlos_bench.services.remote.ssh.ssh_host_service import SshHostService 

22from mlos_bench.services.remote.ssh.ssh_fileshare import SshFileShareService 

23 

24from mlos_bench.tests import resolve_host_name 

25from mlos_bench.tests.services.remote.ssh import (SshTestServerInfo, 

26 ALT_TEST_SERVER_NAME, 

27 REBOOT_TEST_SERVER_NAME, 

28 SSH_TEST_SERVER_NAME, 

29 wait_docker_service_socket) 

30 

31# pylint: disable=redefined-outer-name 

32 

33HOST_DOCKER_NAME = 'host.docker.internal' 

34 

35 

36@pytest.fixture(scope="session") 

37def ssh_test_server_hostname() -> str: 

38 """Returns the local hostname to use to connect to the test ssh server.""" 

39 if sys.platform != 'win32' and resolve_host_name(HOST_DOCKER_NAME): 

40 # On Linux, if we're running in a docker container, we can use the 

41 # --add-host (extra_hosts in docker-compose.yml) to refer to the host IP. 

42 return HOST_DOCKER_NAME 

43 # Docker (Desktop) for Windows (WSL2) uses a special networking magic 

44 # to refer to the host machine as `localhost` when exposing ports. 

45 # In all other cases, assume we're executing directly inside conda on the host. 

46 return 'localhost' 

47 

48 

49@pytest.fixture(scope="session") 

50def ssh_test_server(ssh_test_server_hostname: str, 

51 docker_compose_project_name: str, 

52 locked_docker_services: DockerServices) -> Generator[SshTestServerInfo, None, None]: 

53 """ 

54 Fixture for getting the ssh test server services setup via docker-compose 

55 using pytest-docker. 

56 

57 Yields the (hostname, port, username, id_rsa_path) of the test server. 

58 

59 Once the session is over, the docker containers are torn down, and the 

60 temporary file holding the dynamically generated private key of the test 

61 server is deleted. 

62 """ 

63 # Get a copy of the ssh id_rsa key from the test ssh server. 

64 with tempfile.NamedTemporaryFile() as id_rsa_file: 

65 ssh_test_server_info = SshTestServerInfo( 

66 compose_project_name=docker_compose_project_name, 

67 service_name=SSH_TEST_SERVER_NAME, 

68 hostname=ssh_test_server_hostname, 

69 username='root', 

70 id_rsa_path=id_rsa_file.name) 

71 wait_docker_service_socket(locked_docker_services, ssh_test_server_info.hostname, ssh_test_server_info.get_port()) 

72 id_rsa_src = f"/{ssh_test_server_info.username}/.ssh/id_rsa" 

73 docker_cp_cmd = f"docker compose -p {docker_compose_project_name} cp {SSH_TEST_SERVER_NAME}:{id_rsa_src} {id_rsa_file.name}" 

74 cmd = run(docker_cp_cmd.split(), check=True, cwd=os.path.dirname(__file__), capture_output=True, text=True) 

75 if cmd.returncode != 0: 

76 raise RuntimeError(f"Failed to copy ssh key from {SSH_TEST_SERVER_NAME} container " 

77 + f"[return={cmd.returncode}]: {str(cmd.stderr)}") 

78 os.chmod(id_rsa_file.name, 0o600) 

79 yield ssh_test_server_info 

80 # NamedTempFile deleted on context exit 

81 

82 

83@pytest.fixture(scope="session") 

84def alt_test_server(ssh_test_server: SshTestServerInfo, 

85 locked_docker_services: DockerServices) -> SshTestServerInfo: 

86 """ 

87 Fixture for getting the second ssh test server info from the docker-compose.yml. 

88 See additional notes in the ssh_test_server fixture above. 

89 """ 

90 # Note: The alt-server uses the same image as the ssh-server container, so 

91 # the id_rsa key and username should all match. 

92 # Only the host port it is allocate is different. 

93 alt_test_server_info = SshTestServerInfo( 

94 compose_project_name=ssh_test_server.compose_project_name, 

95 service_name=ALT_TEST_SERVER_NAME, 

96 hostname=ssh_test_server.hostname, 

97 username=ssh_test_server.username, 

98 id_rsa_path=ssh_test_server.id_rsa_path) 

99 wait_docker_service_socket(locked_docker_services, alt_test_server_info.hostname, alt_test_server_info.get_port()) 

100 return alt_test_server_info 

101 

102 

103@pytest.fixture(scope="session") 

104def reboot_test_server(ssh_test_server: SshTestServerInfo, 

105 locked_docker_services: DockerServices) -> SshTestServerInfo: 

106 """ 

107 Fixture for getting the third ssh test server info from the docker-compose.yml. 

108 See additional notes in the ssh_test_server fixture above. 

109 """ 

110 # Note: The reboot-server uses the same image as the ssh-server container, so 

111 # the id_rsa key and username should all match. 

112 # Only the host port it is allocate is different. 

113 reboot_test_server_info = SshTestServerInfo( 

114 compose_project_name=ssh_test_server.compose_project_name, 

115 service_name=REBOOT_TEST_SERVER_NAME, 

116 hostname=ssh_test_server.hostname, 

117 username=ssh_test_server.username, 

118 id_rsa_path=ssh_test_server.id_rsa_path) 

119 wait_docker_service_socket(locked_docker_services, reboot_test_server_info.hostname, reboot_test_server_info.get_port()) 

120 return reboot_test_server_info 

121 

122 

123@pytest.fixture 

124def ssh_host_service(ssh_test_server: SshTestServerInfo) -> SshHostService: 

125 """Generic SshHostService fixture.""" 

126 return SshHostService( 

127 config={ 

128 "ssh_username": ssh_test_server.username, 

129 "ssh_priv_key_path": ssh_test_server.id_rsa_path, 

130 }, 

131 ) 

132 

133 

134@pytest.fixture 

135def ssh_fileshare_service() -> SshFileShareService: 

136 """Generic SshFileShareService fixture.""" 

137 return SshFileShareService( 

138 config={ 

139 # Left blank to make sure we test per connection overrides. 

140 }, 

141 )