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 

9from typing import Optional 

10 

11from pytest_docker.plugin import Services as DockerServices 

12 

13from mlos_bench.tests import check_socket 

14 

15# The SSH test server port and name. 

16# See Also: docker-compose.yml 

17SSH_TEST_SERVER_PORT = 2254 

18SSH_TEST_SERVER_NAME = "ssh-server" 

19ALT_TEST_SERVER_NAME = "alt-server" 

20REBOOT_TEST_SERVER_NAME = "reboot-server" 

21 

22 

23@dataclass 

24class SshTestServerInfo: 

25 """A data class for SshTestServerInfo.""" 

26 

27 compose_project_name: str 

28 service_name: str 

29 hostname: str 

30 username: str 

31 id_rsa_path: str 

32 _port: Optional[int] = None 

33 

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

35 """ 

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

37 

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

39 the DockerServices. 

40 """ 

41 if self._port is None or uncached: 

42 port_cmd = run( 

43 ( 

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

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

46 ), 

47 shell=True, 

48 check=True, 

49 capture_output=True, 

50 ) 

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

52 return self._port 

53 

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

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

56 return { 

57 "ssh_hostname": self.hostname, 

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

59 "ssh_username": self.username, 

60 "ssh_priv_key_path": self.id_rsa_path, 

61 } 

62 

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

64 """ 

65 Convert to a connect_params dict for SshClient. 

66 

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

68 """ 

69 return { 

70 "host": self.hostname, 

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

72 "username": self.username, 

73 } 

74 

75 

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

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

78 docker_services.wait_until_responsive( 

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

80 timeout=30.0, 

81 pause=0.5, 

82 )