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

1# 

2# Copyright (c) Microsoft Corporation. 

3# Licensed under the MIT License. 

4# 

5"""Tests for mlos_bench.services.remote.ssh.SshService base class.""" 

6 

7import asyncio 

8import time 

9from importlib.metadata import PackageNotFoundError, version 

10from subprocess import run 

11from threading import Thread 

12 

13import pytest 

14from pytest_lazy_fixtures.lazy_fixture import lf as lazy_fixture 

15 

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

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

18from mlos_bench.services.remote.ssh.ssh_service import SshService 

19from mlos_bench.tests import ( 

20 check_socket, 

21 requires_docker, 

22 requires_ssh, 

23 resolve_host_name, 

24) 

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

26 ALT_TEST_SERVER_NAME, 

27 SSH_TEST_SERVER_NAME, 

28 SshTestServerInfo, 

29) 

30 

31if version("pytest") >= "8.0.0": 

32 try: 

33 # We replaced pytest-lazy-fixture with pytest-lazy-fixtures: 

34 # https://github.com/TvoroG/pytest-lazy-fixture/issues/65 

35 if version("pytest-lazy-fixture"): 

36 raise UserWarning( 

37 "pytest-lazy-fixture conflicts with pytest>=8.0.0. Please remove it." 

38 ) 

39 except PackageNotFoundError: 

40 # OK: pytest-lazy-fixture not installed 

41 pass 

42 

43 

44@requires_docker 

45@requires_ssh 

46@pytest.mark.parametrize( 

47 ["ssh_test_server_info", "server_name"], 

48 [ 

49 (lazy_fixture("ssh_test_server"), SSH_TEST_SERVER_NAME), 

50 (lazy_fixture("alt_test_server"), ALT_TEST_SERVER_NAME), 

51 ], 

52) 

53def test_ssh_service_test_infra(ssh_test_server_info: SshTestServerInfo, server_name: str) -> None: 

54 """Check for the pytest-docker ssh test infra.""" 

55 assert ssh_test_server_info.service_name == server_name 

56 

57 ip_addr = resolve_host_name(ssh_test_server_info.hostname) 

58 assert ip_addr is not None 

59 

60 local_port = ssh_test_server_info.get_port() 

61 assert check_socket(ip_addr, local_port) 

62 ssh_cmd = ( 

63 "ssh -o BatchMode=yes -o StrictHostKeyChecking=accept-new " 

64 + f"-l {ssh_test_server_info.username} -i {ssh_test_server_info.id_rsa_path} " 

65 + f"-p {local_port} {ssh_test_server_info.hostname} hostname" 

66 ) 

67 cmd = run(ssh_cmd.split(), capture_output=True, text=True, check=True) 

68 assert cmd.stdout.strip() == server_name 

69 

70 

71@pytest.mark.filterwarnings( 

72 "ignore:.*(coroutine 'sleep' was never awaited).*:RuntimeWarning:.*event_loop_context_test.*:0" 

73) 

74def test_ssh_service_context_handler() -> None: 

75 """ 

76 Test the SSH service context manager handling. 

77 

78 See Also: test_event_loop_context 

79 """ 

80 # pylint: disable=protected-access 

81 

82 # Should start with no event loop thread. 

83 assert SshService._EVENT_LOOP_CONTEXT._event_loop_thread is None 

84 

85 # The background thread should only be created upon context entry. 

86 ssh_host_service = SshHostService(config={}, global_config={}, parent=None) 

87 assert ssh_host_service 

88 assert not ssh_host_service._in_context 

89 assert ssh_host_service._EVENT_LOOP_CONTEXT._event_loop_thread is None 

90 

91 # After we enter the SshService instance context, we should have a background thread. 

92 with ssh_host_service: 

93 assert ssh_host_service._in_context 

94 assert ( # type: ignore[unreachable] 

95 isinstance(SshService._EVENT_LOOP_CONTEXT._event_loop_thread, Thread) 

96 ) 

97 # Give the thread a chance to start. 

98 # Mostly important on the underpowered Windows CI machines. 

99 time.sleep(0.25) 

100 assert SshService._EVENT_LOOP_THREAD_SSH_CLIENT_CACHE is not None 

101 

102 ssh_fileshare_service = SshFileShareService(config={}, global_config={}, parent=None) 

103 assert ssh_fileshare_service 

104 assert not ssh_fileshare_service._in_context 

105 

106 with ssh_fileshare_service: 

107 assert ssh_fileshare_service._in_context 

108 assert ssh_host_service._in_context 

109 assert ( 

110 SshService._EVENT_LOOP_CONTEXT._event_loop_thread 

111 is ssh_host_service._EVENT_LOOP_CONTEXT._event_loop_thread 

112 is ssh_fileshare_service._EVENT_LOOP_CONTEXT._event_loop_thread 

113 ) 

114 assert ( 

115 SshService._EVENT_LOOP_THREAD_SSH_CLIENT_CACHE 

116 is ssh_host_service._EVENT_LOOP_THREAD_SSH_CLIENT_CACHE 

117 is ssh_fileshare_service._EVENT_LOOP_THREAD_SSH_CLIENT_CACHE 

118 ) 

119 

120 assert not ssh_fileshare_service._in_context 

121 # And that instance should be unusable after we are outside the context. 

122 with pytest.raises( 

123 AssertionError 

124 ): # , pytest.warns(RuntimeWarning, match=r".*coroutine 'sleep' was never awaited"): 

125 future = ssh_fileshare_service._run_coroutine(asyncio.sleep(0.1, result="foo")) 

126 raise ValueError(f"Future should not have been available to wait on {future.result()}") 

127 

128 # The background thread should remain running since we have another context still open. 

129 assert isinstance(SshService._EVENT_LOOP_CONTEXT._event_loop_thread, Thread) 

130 assert SshService._EVENT_LOOP_THREAD_SSH_CLIENT_CACHE is not None 

131 

132 

133if __name__ == "__main__": 

134 # For debugging in Windows which has issues with pytest detection in vscode. 

135 pytest.main(["-n1", "--dist=no", "-k", "test_ssh_service_background_thread"])