Coverage for mlos_bench/mlos_bench/tests/services/remote/ssh/__init__.py: 100%

1# 

2# Copyright (c) Microsoft Corporation. 

3# Licensed under the MIT License. 

4# 

5"""Common data classes for the SSH service tests.""" 

6 

7from dataclasses import dataclass 

8from subprocess import run 

9 

10from pytest_docker.plugin import Services as DockerServices 

11 

12from mlos_bench.tests import check_socket 

13 

14# The SSH test server port and name. 

15# See Also: docker-compose.yml 

16SSH_TEST_SERVER_PORT = 2254 

17SSH_TEST_SERVER_NAME = "ssh-server" 

18ALT_TEST_SERVER_NAME = "alt-server" 

19REBOOT_TEST_SERVER_NAME = "reboot-server" 

20 

21 

22@dataclass 

23class SshTestServerInfo: 

24 """A data class for SshTestServerInfo.""" 

25 

26 compose_project_name: str 

27 service_name: str 

28 hostname: str 

29 username: str 

30 id_rsa_path: str 

31 _port: int | None = None 

32 

33 def get_port(self, uncached: bool = False) -> int: 

34 """ 

35 Gets the port that the SSH test server is listening on. 

36 

37 Note: this value can change when the service restarts so we can't rely on 

38 the DockerServices. 

39 """ 

40 if self._port is None or uncached: 

41 port_cmd = run( 

42 ( 

43 f"docker compose -p {self.compose_project_name} " 

44 f"port {self.service_name} {SSH_TEST_SERVER_PORT}" 

45 ), 

46 shell=True, 

47 check=True, 

48 capture_output=True, 

49 ) 

50 self._port = int(port_cmd.stdout.decode().strip().split(":")[1]) 

51 return self._port 

52 

53 def to_ssh_service_config(self, uncached: bool = False) -> dict: 

54 """Convert to a config dict for SshService.""" 

55 return { 

56 "ssh_hostname": self.hostname, 

57 "ssh_port": self.get_port(uncached), 

58 "ssh_username": self.username, 

59 "ssh_priv_key_path": self.id_rsa_path, 

60 } 

61 

62 def to_connect_params(self, uncached: bool = False) -> dict: 

63 """ 

64 Convert to a connect_params dict for SshClient. 

65 

66 See Also: mlos_bench.services.remote.ssh.ssh_service.SshService._get_connect_params() 

67 """ 

68 return { 

69 "host": self.hostname, 

70 "port": self.get_port(uncached), 

71 "username": self.username, 

72 } 

73 

74 

75def wait_docker_service_socket(docker_services: DockerServices, hostname: str, port: int) -> None: 

76 """Wait until a docker service is ready.""" 

77 docker_services.wait_until_responsive( 

78 check=lambda: check_socket(hostname, port), 

79 timeout=30.0, 

80 pause=0.5, 

81 )