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

6Unit tests for the service to run the scripts locally. 

7""" 

8import sys 

9import tempfile 

10 

11import pytest 

12import pandas 

13 

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

15from mlos_bench.services.config_persistence import ConfigPersistenceService 

16from mlos_bench.util import path_join 

17 

18# pylint: disable=redefined-outer-name 

19# -- Ignore pylint complaints about pytest references to 

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

21 

22 

23def test_split_cmdline() -> None: 

24 """ 

25 Test splitting a commandline into subcommands. 

26 """ 

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

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

29 ['.', 'env.sh'], 

30 ['&&'], 

31 ['('], 

32 ['echo', 'hello'], 

33 ['&&'], 

34 ['echo', 'world'], 

35 ['|'], 

36 ['tee'], 

37 ['>'], 

38 ['/tmp/test'], 

39 ['||'], 

40 ['echo', 'foo'], 

41 ['&&'], 

42 ['echo', '$var'], 

43 [';'], 

44 ['true'], 

45 [')'], 

46 ] 

47 

48 

49@pytest.fixture 

50def local_exec_service() -> LocalExecService: 

51 """ 

52 Test fixture for LocalExecService. 

53 """ 

54 config = { 

55 "abort_on_error": True, 

56 } 

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

58 

59 

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

61 """ 

62 Test local script resolution logic with complex subcommand names. 

63 """ 

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

65 script_abspath = local_exec_service.config_loader_service.resolve_path(script) 

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

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

68 subcmds_tokens = split_cmdline(orig_cmdline) 

69 # pylint: disable=protected-access 

70 subcmds_tokens = [local_exec_service._resolve_cmdline_script_path(subcmd_tokens) for subcmd_tokens in subcmds_tokens] 

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

72 expanded_cmdline = " ".join(cmdline_tokens) 

73 assert expanded_cmdline == expected_cmdline 

74 

75 

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

77 """ 

78 Run a script locally and check the results. 

79 """ 

80 # `echo` should work on all platforms 

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

82 assert return_code == 0 

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

84 assert stderr.strip() == "" 

85 

86 

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

88 """ 

89 Run a multiline script locally and check the results. 

90 """ 

91 # `echo` should work on all platforms 

92 (return_code, stdout, stderr) = local_exec_service.local_exec([ 

93 "echo hello", 

94 "echo world" 

95 ]) 

96 assert return_code == 0 

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

98 assert stderr.strip() == "" 

99 

100 

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

102 """ 

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

104 """ 

105 # `echo` should work on all platforms 

106 (return_code, stdout, stderr) = local_exec_service.local_exec([ 

107 r"echo $var", # Unix shell 

108 r"echo %var%" # Windows cmd 

109 ], env={"var": "VALUE", "int_var": 10}) 

110 assert return_code == 0 

111 if sys.platform == 'win32': 

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

113 else: 

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

115 assert stderr.strip() == "" 

116 

117 

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

119 """ 

120 Run a script locally and read the resulting CSV file. 

121 """ 

122 with local_exec_service.temp_dir_context() as temp_dir: 

123 

124 (return_code, stdout, stderr) = local_exec_service.local_exec([ 

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

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

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

128 ], cwd=temp_dir) 

129 

130 assert return_code == 0 

131 assert stdout.strip() == "" 

132 assert stderr.strip() == "" 

133 

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

135 if sys.platform == 'win32': 

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

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

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

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

140 data.rename(str.rstrip, axis='columns', inplace=True) 

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

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

143 

144 

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

146 """ 

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

148 """ 

149 with local_exec_service.temp_dir_context() as temp_dir: 

150 

151 input_file = "input.txt" 

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

153 fh_input.write("hello\n") 

154 

155 (return_code, stdout, stderr) = local_exec_service.local_exec([ 

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

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

158 ], cwd=temp_dir) 

159 

160 assert return_code == 0 

161 assert stdout.strip() == "" 

162 assert stderr.strip() == "" 

163 

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

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

166 

167 

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

169 """ 

170 Try to run a non-existent command. 

171 """ 

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

173 assert return_code != 0 

174 assert stdout.strip() == "" 

175 

176 

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

178 """ 

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

180 """ 

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

182 "echo hello", 

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

184 "echo world", 

185 ]) 

186 assert return_code != 0 

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

188 

189 

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

191 """ 

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

193 """ 

194 local_exec_service.abort_on_error = False 

195 (return_code, stdout, _stderr) = local_exec_service.local_exec([ 

196 "echo hello", 

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

198 "echo world", 

199 ]) 

200 assert return_code == 0 

201 assert stdout.splitlines() == [ 

202 "hello", 

203 "world", 

204 ] 

205 

206 

207def test_temp_dir_path_expansion() -> None: 

208 """ 

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

210 """ 

211 # Create a temp dir for the test. 

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

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

214 # the fact. 

215 with tempfile.TemporaryDirectory() as temp_dir: 

216 global_config = { 

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

218 } 

219 config = { 

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

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

222 } 

223 local_exec_service = LocalExecService(config, global_config, parent=ConfigPersistenceService()) 

224 # pylint: disable=protected-access 

225 assert isinstance(local_exec_service._temp_dir, str) 

226 assert path_join(local_exec_service._temp_dir, abs_path=True) == path_join(temp_dir, "temp", abs_path=True)