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

6Tests for mlos_bench.services.remote.ssh.SshService base class. 

7""" 

8 

9import asyncio 

10from importlib.metadata import version, PackageNotFoundError 

11import time 

12 

13from subprocess import run 

14from threading import Thread 

15 

16import pytest 

17from pytest_lazy_fixtures.lazy_fixture import lf as lazy_fixture 

18 

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

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

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

22 

23from mlos_bench.tests import requires_docker, requires_ssh, check_socket, resolve_host_name 

24from mlos_bench.tests.services.remote.ssh import SshTestServerInfo, ALT_TEST_SERVER_NAME, SSH_TEST_SERVER_NAME 

25 

26 

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

28 try: 

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

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

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

32 raise UserWarning("pytest-lazy-fixture conflicts with pytest>=8.0.0. Please remove it.") 

33 except PackageNotFoundError: 

34 # OK: pytest-lazy-fixture not installed 

35 pass 

36 

37 

38@requires_docker 

39@requires_ssh 

40@pytest.mark.parametrize(["ssh_test_server_info", "server_name"], [ 

41 (lazy_fixture("ssh_test_server"), SSH_TEST_SERVER_NAME), 

42 (lazy_fixture("alt_test_server"), ALT_TEST_SERVER_NAME), 

43]) 

44def test_ssh_service_test_infra(ssh_test_server_info: SshTestServerInfo, 

45 server_name: str) -> None: 

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

47 assert ssh_test_server_info.service_name == server_name 

48 

49 ip_addr = resolve_host_name(ssh_test_server_info.hostname) 

50 assert ip_addr is not None 

51 

52 local_port = ssh_test_server_info.get_port() 

53 assert check_socket(ip_addr, local_port) 

54 ssh_cmd = "ssh -o BatchMode=yes -o StrictHostKeyChecking=accept-new " \ 

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

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

57 cmd = run(ssh_cmd.split(), 

58 capture_output=True, 

59 text=True, 

60 check=True) 

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

62 

63 

64@pytest.mark.filterwarnings("ignore:.*(coroutine 'sleep' was never awaited).*:RuntimeWarning:.*event_loop_context_test.*:0") 

65def test_ssh_service_context_handler() -> None: 

66 """ 

67 Test the SSH service context manager handling. 

68 See Also: test_event_loop_context 

69 """ 

70 # pylint: disable=protected-access 

71 

72 # Should start with no event loop thread. 

73 assert SshService._EVENT_LOOP_CONTEXT._event_loop_thread is None 

74 

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

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

77 assert ssh_host_service 

78 assert not ssh_host_service._in_context 

79 assert ssh_host_service._EVENT_LOOP_CONTEXT._event_loop_thread is None 

80 

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

82 with ssh_host_service: 

83 assert ssh_host_service._in_context 

84 assert isinstance(SshService._EVENT_LOOP_CONTEXT._event_loop_thread, Thread) # type: ignore[unreachable] 

85 # Give the thread a chance to start. 

86 # Mostly important on the underpowered Windows CI machines. 

87 time.sleep(0.25) 

88 assert SshService._EVENT_LOOP_THREAD_SSH_CLIENT_CACHE is not None 

89 

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

91 assert ssh_fileshare_service 

92 assert not ssh_fileshare_service._in_context 

93 

94 with ssh_fileshare_service: 

95 assert ssh_fileshare_service._in_context 

96 assert ssh_host_service._in_context 

97 assert SshService._EVENT_LOOP_CONTEXT._event_loop_thread \ 

98 is ssh_host_service._EVENT_LOOP_CONTEXT._event_loop_thread \ 

99 is ssh_fileshare_service._EVENT_LOOP_CONTEXT._event_loop_thread 

100 assert SshService._EVENT_LOOP_THREAD_SSH_CLIENT_CACHE \ 

101 is ssh_host_service._EVENT_LOOP_THREAD_SSH_CLIENT_CACHE \ 

102 is ssh_fileshare_service._EVENT_LOOP_THREAD_SSH_CLIENT_CACHE 

103 

104 assert not ssh_fileshare_service._in_context 

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

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

107 future = ssh_fileshare_service._run_coroutine(asyncio.sleep(0.1, result='foo')) 

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

109 

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

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

112 assert SshService._EVENT_LOOP_THREAD_SSH_CLIENT_CACHE is not None 

113 

114 

115if __name__ == '__main__': 

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

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