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""" 

6Common data classes for the SSH service tests. 

7""" 

8 

9from dataclasses import dataclass 

10from subprocess import run 

11from typing import Optional 

12 

13from pytest_docker.plugin import Services as DockerServices 

14 

15from mlos_bench.tests import check_socket 

16 

17 

18# The SSH test server port and name. 

19# See Also: docker-compose.yml 

20SSH_TEST_SERVER_PORT = 2254 

21SSH_TEST_SERVER_NAME = 'ssh-server' 

22ALT_TEST_SERVER_NAME = 'alt-server' 

23REBOOT_TEST_SERVER_NAME = 'reboot-server' 

24 

25 

26@dataclass 

27class SshTestServerInfo: 

28 """ 

29 A data class for SshTestServerInfo. 

30 """ 

31 

32 compose_project_name: str 

33 service_name: str 

34 hostname: str 

35 username: str 

36 id_rsa_path: str 

37 _port: Optional[int] = None 

38 

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

40 """ 

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

42 

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

44 """ 

45 if self._port is None or uncached: 

46 port_cmd = run(f"docker compose -p {self.compose_project_name} port {self.service_name} {SSH_TEST_SERVER_PORT}", 

47 shell=True, check=True, capture_output=True) 

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

49 return self._port 

50 

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

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

53 return { 

54 "ssh_hostname": self.hostname, 

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

56 "ssh_username": self.username, 

57 "ssh_priv_key_path": self.id_rsa_path, 

58 } 

59 

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

61 """ 

62 Convert to a connect_params dict for SshClient. 

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

64 """ 

65 return { 

66 "host": self.hostname, 

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

68 "username": self.username, 

69 } 

70 

71 

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

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

74 docker_services.wait_until_responsive( 

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

76 timeout=30.0, 

77 pause=0.5, 

78 )