Coverage for mlos_bench/mlos_bench/tests/services/remote/ssh/test_ssh_service.py: 96%
56 statements
« prev ^ index » next coverage.py v7.6.9, created at 2024-12-20 00:44 +0000
« prev ^ index » next coverage.py v7.6.9, created at 2024-12-20 00:44 +0000
1#
2# Copyright (c) Microsoft Corporation.
3# Licensed under the MIT License.
4#
5"""Tests for mlos_bench.services.remote.ssh.SshService base class."""
7import asyncio
8import time
9from importlib.metadata import PackageNotFoundError, version
10from subprocess import run
11from threading import Thread
13import pytest
14from pytest_lazy_fixtures.lazy_fixture import lf as lazy_fixture
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)
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
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
57 ip_addr = resolve_host_name(ssh_test_server_info.hostname)
58 assert ip_addr is not None
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
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.
78 See Also: test_event_loop_context
79 """
80 # pylint: disable=protected-access
82 # Should start with no event loop thread.
83 assert SshService._EVENT_LOOP_CONTEXT._event_loop_thread is None
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
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
102 ssh_fileshare_service = SshFileShareService(config={}, global_config={}, parent=None)
103 assert ssh_fileshare_service
104 assert not ssh_fileshare_service._in_context
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 )
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()}")
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
133if __name__ == "__main__":
134 # For debugging in Windows which has issues with pytest detection in vscode.
135 pytest.main(["-n0", "--dist=no", "-k", "test_ssh_service_context_handler"])