Coverage for mlos_bench/mlos_bench/tests/services/local/local_exec_test.py: 98%

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 

8 

9import pandas 

10import pytest 

11 

12from mlos_bench.services.config_persistence import ConfigPersistenceService 

13from mlos_bench.services.local.local_exec import LocalExecService, split_cmdline 

14from mlos_bench.util import path_join 

15 

16# pylint: disable=redefined-outer-name 

17# -- Ignore pylint complaints about pytest references to 

18# `local_exec_service` fixture as both a function and a parameter. 

19 

20 

21def test_split_cmdline() -> None: 

22 """Test splitting a commandline into subcommands.""" 

23 cmdline = ( 

24 ". env.sh && (echo hello && echo world | tee > /tmp/test || echo foo && echo $var; true)" 

25 ) 

26 assert list(split_cmdline(cmdline)) == [ 

27 [".", "env.sh"], 

28 ["&&"], 

29 ["("], 

30 ["echo", "hello"], 

31 ["&&"], 

32 ["echo", "world"], 

33 ["|"], 

34 ["tee"], 

35 [">"], 

36 ["/tmp/test"], 

37 ["||"], 

38 ["echo", "foo"], 

39 ["&&"], 

40 ["echo", "$var"], 

41 [";"], 

42 ["true"], 

43 [")"], 

44 ] 

45 

46 

47@pytest.fixture 

48def local_exec_service() -> LocalExecService: 

49 """Test fixture for LocalExecService.""" 

50 config = { 

51 "abort_on_error": True, 

52 } 

53 return LocalExecService(config, parent=ConfigPersistenceService()) 

54 

55 

56def test_resolve_script(local_exec_service: LocalExecService) -> None: 

57 """Test local script resolution logic with complex subcommand names.""" 

58 script = "os/linux/runtime/scripts/local/generate_kernel_config_script.py" 

59 script_abspath = local_exec_service.config_loader_service.resolve_path(script) 

60 orig_cmdline = f". env.sh && {script} --input foo" 

61 expected_cmdline = f". env.sh && {script_abspath} --input foo" 

62 subcmds_tokens = split_cmdline(orig_cmdline) 

63 # pylint: disable=protected-access 

64 subcmds_tokens = [ 

65 local_exec_service._resolve_cmdline_script_path(subcmd_tokens) 

66 for subcmd_tokens in subcmds_tokens 

67 ] 

68 cmdline_tokens = [token for subcmd_tokens in subcmds_tokens for token in subcmd_tokens] 

69 expanded_cmdline = " ".join(cmdline_tokens) 

70 assert expanded_cmdline == expected_cmdline 

71 

72 

73def test_run_script(local_exec_service: LocalExecService) -> None: 

74 """Run a script locally and check the results.""" 

75 # `echo` should work on all platforms 

76 (return_code, stdout, stderr) = local_exec_service.local_exec(["echo hello"]) 

77 assert return_code == 0 

78 assert stdout.strip() == "hello" 

79 assert stderr.strip() == "" 

80 

81 

82def test_run_script_multiline(local_exec_service: LocalExecService) -> None: 

83 """Run a multiline script locally and check the results.""" 

84 # `echo` should work on all platforms 

85 (return_code, stdout, stderr) = local_exec_service.local_exec(["echo hello", "echo world"]) 

86 assert return_code == 0 

87 assert stdout.strip().split() == ["hello", "world"] 

88 assert stderr.strip() == "" 

89 

90 

91def test_run_script_multiline_env(local_exec_service: LocalExecService) -> None: 

92 """Run a multiline script locally and pass the environment variables to it.""" 

93 # `echo` should work on all platforms 

94 (return_code, stdout, stderr) = local_exec_service.local_exec( 

95 [r"echo $var", r"echo %var%"], # Unix shell # Windows cmd 

96 env={"var": "VALUE", "int_var": 10}, 

97 ) 

98 assert return_code == 0 

99 if sys.platform == "win32": 

100 assert stdout.strip().split() == ["$var", "VALUE"] 

101 else: 

102 assert stdout.strip().split() == ["VALUE", "%var%"] 

103 assert stderr.strip() == "" 

104 

105 

106def test_run_script_read_csv(local_exec_service: LocalExecService) -> None: 

107 """Run a script locally and read the resulting CSV file.""" 

108 with local_exec_service.temp_dir_context() as temp_dir: 

109 

110 (return_code, stdout, stderr) = local_exec_service.local_exec( 

111 [ 

112 "echo 'col1,col2'> output.csv", # No space before '>' to make it work on Windows 

113 "echo '111,222' >> output.csv", 

114 "echo '333,444' >> output.csv", 

115 ], 

116 cwd=temp_dir, 

117 ) 

118 

119 assert return_code == 0 

120 assert stdout.strip() == "" 

121 assert stderr.strip() == "" 

122 

123 data = pandas.read_csv(path_join(temp_dir, "output.csv")) 

124 if sys.platform == "win32": 

125 # Workaround for Python's subprocess module on Windows adding a 

126 # space inbetween the col1,col2 arg and the redirect symbol which 

127 # cmd poorly interprets as being part of the original string arg. 

128 # Without this, we get "col2 " as the second column name. 

129 data.rename(str.rstrip, axis="columns", inplace=True) 

130 assert all(data.col1 == [111, 333]) 

131 assert all(data.col2 == [222, 444]) 

132 

133 

134def test_run_script_write_read_txt(local_exec_service: LocalExecService) -> None: 

135 """Write data a temp location and run a script that updates it there.""" 

136 with local_exec_service.temp_dir_context() as temp_dir: 

137 

138 input_file = "input.txt" 

139 with open(path_join(temp_dir, input_file), "wt", encoding="utf-8") as fh_input: 

140 fh_input.write("hello\n") 

141 

142 (return_code, stdout, stderr) = local_exec_service.local_exec( 

143 [ 

144 f"echo 'world' >> {input_file}", 

145 f"echo 'test' >> {input_file}", 

146 ], 

147 cwd=temp_dir, 

148 ) 

149 

150 assert return_code == 0 

151 assert stdout.strip() == "" 

152 assert stderr.strip() == "" 

153 

154 with open(path_join(temp_dir, input_file), "rt", encoding="utf-8") as fh_input: 

155 assert fh_input.read().split() == ["hello", "world", "test"] 

156 

157 

158def test_run_script_fail(local_exec_service: LocalExecService) -> None: 

159 """Try to run a non-existent command.""" 

160 (return_code, stdout, _stderr) = local_exec_service.local_exec(["foo_bar_baz hello"]) 

161 assert return_code != 0 

162 assert stdout.strip() == "" 

163 

164 

165def test_run_script_middle_fail_abort(local_exec_service: LocalExecService) -> None: 

166 """Try to run a series of commands, one of which fails, and abort early.""" 

167 (return_code, stdout, _stderr) = local_exec_service.local_exec( 

168 [ 

169 "echo hello", 

170 "cmd /c 'exit 1'" if sys.platform == "win32" else "false", 

171 "echo world", 

172 ] 

173 ) 

174 assert return_code != 0 

175 assert stdout.strip() == "hello" 

176 

177 

178def test_run_script_middle_fail_pass(local_exec_service: LocalExecService) -> None: 

179 """Try to run a series of commands, one of which fails, but let it pass.""" 

180 local_exec_service.abort_on_error = False 

181 (return_code, stdout, _stderr) = local_exec_service.local_exec( 

182 [ 

183 "echo hello", 

184 "cmd /c 'exit 1'" if sys.platform == "win32" else "false", 

185 "echo world", 

186 ] 

187 ) 

188 assert return_code == 0 

189 assert stdout.splitlines() == [ 

190 "hello", 

191 "world", 

192 ] 

193 

194 

195def test_temp_dir_path_expansion() -> None: 

196 """Test that we can control the temp_dir path using globals expansion.""" 

197 # Create a temp dir for the test. 

198 # Normally this would be a real path set on the CLI or in a global config, 

199 # but for test purposes we still want it to be dynamic and cleaned up after 

200 # the fact. 

201 with tempfile.TemporaryDirectory() as temp_dir: 

202 global_config = { 

203 "workdir": temp_dir, # e.g., "." or "/tmp/mlos_bench" 

204 } 

205 config = { 

206 # The temp_dir for the LocalExecService should get expanded via workdir global config. 

207 "temp_dir": "$workdir/temp", 

208 } 

209 local_exec_service = LocalExecService( 

210 config, global_config, parent=ConfigPersistenceService() 

211 ) 

212 # pylint: disable=protected-access 

213 assert isinstance(local_exec_service._temp_dir, str) 

214 assert path_join(local_exec_service._temp_dir, abs_path=True) == path_join( 

215 temp_dir, 

216 "temp", 

217 abs_path=True, 

218 )