Coverage for mlos_bench/mlos_bench/tests/services/local/local_exec_test.py: 98%
86 statements
« prev ^ index » next coverage.py v7.8.0, created at 2025-04-01 00:52 +0000
1#
2# Copyright (c) Microsoft Corporation.
3# Licensed under the MIT License.
4#
5"""Unit tests for the service to run the scripts locally."""
6import sys
7import tempfile
9import pandas
10import pytest
12from mlos_bench.os_environ import environ
13from mlos_bench.services.config_persistence import ConfigPersistenceService
14from mlos_bench.services.local.local_exec import LocalExecService, split_cmdline
15from mlos_bench.util import path_join
17# pylint: disable=redefined-outer-name
18# -- Ignore pylint complaints about pytest references to
19# `local_exec_service` fixture as both a function and a parameter.
22def test_split_cmdline() -> None:
23 """Test splitting a commandline into subcommands."""
24 cmdline = (
25 ". env.sh && (echo hello && echo world | tee > /tmp/test || echo foo && echo $var; true)"
26 )
27 assert list(split_cmdline(cmdline)) == [
28 [".", "env.sh"],
29 ["&&"],
30 ["("],
31 ["echo", "hello"],
32 ["&&"],
33 ["echo", "world"],
34 ["|"],
35 ["tee"],
36 [">"],
37 ["/tmp/test"],
38 ["||"],
39 ["echo", "foo"],
40 ["&&"],
41 ["echo", "$var"],
42 [";"],
43 ["true"],
44 [")"],
45 ]
48@pytest.fixture
49def local_exec_service() -> LocalExecService:
50 """Test fixture for LocalExecService."""
51 config = {
52 "abort_on_error": True,
53 }
54 return LocalExecService(config, parent=ConfigPersistenceService())
57def test_resolve_script(local_exec_service: LocalExecService) -> None:
58 """Test local script resolution logic with complex subcommand names."""
59 script = "os/linux/runtime/scripts/local/generate_kernel_config_script.py"
60 script_abspath = local_exec_service.config_loader_service.resolve_path(script)
61 orig_cmdline = f". env.sh && {script} --input foo"
62 expected_cmdline = f". env.sh && {script_abspath} --input foo"
63 subcmds_tokens = split_cmdline(orig_cmdline)
64 # pylint: disable=protected-access
65 subcmds_tokens = [
66 local_exec_service._resolve_cmdline_script_path(subcmd_tokens)
67 for subcmd_tokens in subcmds_tokens
68 ]
69 cmdline_tokens = [token for subcmd_tokens in subcmds_tokens for token in subcmd_tokens]
70 expanded_cmdline = " ".join(cmdline_tokens)
71 assert expanded_cmdline == expected_cmdline
74def test_run_script(local_exec_service: LocalExecService) -> None:
75 """Run a script locally and check the results."""
76 # `echo` should work on all platforms
77 (return_code, stdout, stderr) = local_exec_service.local_exec(["echo hello"])
78 assert return_code == 0
79 assert stdout.strip() == "hello"
80 assert stderr.strip() == ""
83def test_run_script_multiline(local_exec_service: LocalExecService) -> None:
84 """Run a multiline script locally and check the results."""
85 # `echo` should work on all platforms
86 (return_code, stdout, stderr) = local_exec_service.local_exec(["echo hello", "echo world"])
87 assert return_code == 0
88 assert stdout.strip().split() == ["hello", "world"]
89 assert stderr.strip() == ""
92def test_run_script_multiline_env(local_exec_service: LocalExecService) -> None:
93 """Run a multiline script locally and pass the environment variables to it."""
94 # `echo` should work on all platforms
95 environ["LOCAL_VAR"] = "LOCAL_VALUE" # Make sure parent env is passed to child
96 (return_code, stdout, stderr) = local_exec_service.local_exec(
97 [
98 r"echo $var $int_var $LOCAL_VAR", # Unix shell
99 r"echo %var% %int_var% %LOCAL_VAR%", # Windows cmd
100 ],
101 env={"var": "VALUE", "int_var": 10},
102 )
103 assert return_code == 0
104 if sys.platform == "win32":
105 expected = ["$var", "$int_var", "$LOCAL_VAR", "VALUE", "10", "LOCAL_VALUE"]
106 else:
107 expected = ["VALUE", "10", "LOCAL_VALUE", "%var%", "%int_var%", "%LOCAL_VAR%"]
108 assert stdout.strip().split() == expected
109 assert stderr.strip() == ""
112def test_run_script_read_csv(local_exec_service: LocalExecService) -> None:
113 """Run a script locally and read the resulting CSV file."""
114 with local_exec_service.temp_dir_context() as temp_dir:
116 (return_code, stdout, stderr) = local_exec_service.local_exec(
117 [
118 "echo 'col1,col2'> output.csv", # No space before '>' to make it work on Windows
119 "echo '111,222' >> output.csv",
120 "echo '333,444' >> output.csv",
121 ],
122 cwd=temp_dir,
123 )
125 assert return_code == 0
126 assert stdout.strip() == ""
127 assert stderr.strip() == ""
129 data = pandas.read_csv(path_join(temp_dir, "output.csv"))
130 if sys.platform == "win32":
131 # Workaround for Python's subprocess module on Windows adding a
132 # space inbetween the col1,col2 arg and the redirect symbol which
133 # cmd poorly interprets as being part of the original string arg.
134 # Without this, we get "col2 " as the second column name.
135 data.rename(str.rstrip, axis="columns", inplace=True)
136 assert all(data.col1 == [111, 333])
137 assert all(data.col2 == [222, 444])
140def test_run_script_write_read_txt(local_exec_service: LocalExecService) -> None:
141 """Write data a temp location and run a script that updates it there."""
142 with local_exec_service.temp_dir_context() as temp_dir:
144 input_file = "input.txt"
145 with open(path_join(temp_dir, input_file), "w", encoding="utf-8") as fh_input:
146 fh_input.write("hello\n")
148 (return_code, stdout, stderr) = local_exec_service.local_exec(
149 [
150 f"echo 'world' >> {input_file}",
151 f"echo 'test' >> {input_file}",
152 ],
153 cwd=temp_dir,
154 )
156 assert return_code == 0
157 assert stdout.strip() == ""
158 assert stderr.strip() == ""
160 with open(path_join(temp_dir, input_file), encoding="utf-8") as fh_input:
161 assert fh_input.read().split() == ["hello", "world", "test"]
164def test_run_script_fail(local_exec_service: LocalExecService) -> None:
165 """Try to run a non-existent command."""
166 (return_code, stdout, _stderr) = local_exec_service.local_exec(["foo_bar_baz hello"])
167 assert return_code != 0
168 assert stdout.strip() == ""
171def test_run_script_middle_fail_abort(local_exec_service: LocalExecService) -> None:
172 """Try to run a series of commands, one of which fails, and abort early."""
173 (return_code, stdout, _stderr) = local_exec_service.local_exec(
174 [
175 "echo hello",
176 "cmd /c 'exit 1'" if sys.platform == "win32" else "false",
177 "echo world",
178 ]
179 )
180 assert return_code != 0
181 assert stdout.strip() == "hello"
184def test_run_script_middle_fail_pass(local_exec_service: LocalExecService) -> None:
185 """Try to run a series of commands, one of which fails, but let it pass."""
186 local_exec_service.abort_on_error = False
187 (return_code, stdout, _stderr) = local_exec_service.local_exec(
188 [
189 "echo hello",
190 "cmd /c 'exit 1'" if sys.platform == "win32" else "false",
191 "echo world",
192 ]
193 )
194 assert return_code == 0
195 assert stdout.splitlines() == [
196 "hello",
197 "world",
198 ]
201def test_temp_dir_path_expansion() -> None:
202 """Test that we can control the temp_dir path using globals expansion."""
203 # Create a temp dir for the test.
204 # Normally this would be a real path set on the CLI or in a global config,
205 # but for test purposes we still want it to be dynamic and cleaned up after
206 # the fact.
207 with tempfile.TemporaryDirectory() as temp_dir:
208 global_config = {
209 "workdir": temp_dir, # e.g., "." or "/tmp/mlos_bench"
210 }
211 config = {
212 # The temp_dir for the LocalExecService should get expanded via workdir global config.
213 "temp_dir": "$workdir/temp",
214 }
215 local_exec_service = LocalExecService(
216 config, global_config, parent=ConfigPersistenceService()
217 )
218 # pylint: disable=protected-access
219 assert isinstance(local_exec_service._temp_dir, str)
220 assert path_join(local_exec_service._temp_dir, abs_path=True) == path_join(
221 temp_dir,
222 "temp",
223 abs_path=True,
224 )