Coverage for /usr/local/lib/python3.12/site-packages/prefect/testing/cli.py: 0%

1from __future__ import annotations 

2 

3import contextlib 

4import re 

5import textwrap 

6from typing import Iterable 

7 

8import readchar 

9from rich.console import Console 

10from typer.testing import CliRunner, Result # type: ignore 

11 

12from prefect.cli import app 

13from prefect.utilities.asyncutils import in_async_main_thread 

14 

15 

16def check_contains(cli_result: Result, content: str, should_contain: bool) -> None: 

17 """ 

18 Utility function to see if content is or is not in a CLI result. 

19 

20 Args: 

21 should_contain: if True, checks that content is in cli_result, 

22 if False, checks that content is not in cli_result 

23 """ 

24 stdout_output = cli_result.stdout.strip() 

25 

26 # Try to get stderr, but handle the case where it's not captured separately 

27 stderr_output = "" 

28 try: 

29 stderr_output = getattr(cli_result, "stderr", "").strip() 

30 except ValueError: 

31 # In some Click/Typer versions, stderr is not separately captured 

32 pass 

33 

34 # Combine both stdout and stderr for checking 

35 output = stdout_output + stderr_output 

36 

37 content = textwrap.dedent(content).strip() 

38 

39 if should_contain: 

40 section_heading = "------ desired content ------" 

41 else: 

42 section_heading = "------ undesired content ------" 

43 

44 print(section_heading) 

45 print(content) 

46 print() 

47 

48 if len(content) > 20: 

49 display_content = content[:20] + "..." 

50 else: 

51 display_content = content 

52 

53 if should_contain: 

54 assert content in output, ( 

55 f"Desired contents {display_content!r} not found in CLI output" 

56 ) 

57 else: 

58 assert content not in output, ( 

59 f"Undesired contents {display_content!r} found in CLI output" 

60 ) 

61 

62 

63def invoke_and_assert( 

64 command: str | list[str], 

65 user_input: str | None = None, 

66 prompts_and_responses: list[tuple[str, str] | tuple[str, str, str]] | None = None, 

67 expected_output: str | None = None, 

68 expected_output_contains: str | Iterable[str] | None = None, 

69 expected_output_does_not_contain: str | Iterable[str] | None = None, 

70 expected_line_count: int | None = None, 

71 expected_code: int | None = 0, 

72 echo: bool = True, 

73 temp_dir: str | None = None, 

74) -> Result: 

75 """ 

76 Test utility for the Prefect CLI application, asserts exact match with CLI output. 

77 

78 Args: 

79 command: Command passed to the Typer CliRunner 

80 user_input: User input passed to the Typer CliRunner when running interactive 

81 commands. 

82 expected_output: Used when you expect the CLI output to be an exact match with 

83 the provided text. 

84 expected_output_contains: Used when you expect the CLI output to contain the 

85 string or strings. 

86 expected_output_does_not_contain: Used when you expect the CLI output to not 

87 contain the string or strings. 

88 expected_code: 0 if we expect the app to exit cleanly, else 1 if we expect 

89 the app to exit with an error. 

90 temp_dir: if provided, the CLI command will be run with this as its present 

91 working directory. 

92 """ 

93 prompts_and_responses = prompts_and_responses or [] 

94 if in_async_main_thread(): 

95 raise RuntimeError( 

96 textwrap.dedent( 

97 """ 

98 You cannot run `invoke_and_assert` directly from an async 

99 function. If you need to run `invoke_and_assert` in an async 

100 function, run it with `run_sync_in_worker_thread`. 

101 

102 Example: 

103 run_sync_in_worker_thread( 

104 invoke_and_assert, 

105 command=['my', 'command'], 

106 expected_code=0, 

107 ) 

108 """ 

109 ) 

110 ) 

111 runner = CliRunner() 

112 if temp_dir: 

113 ctx = runner.isolated_filesystem(temp_dir=temp_dir) 

114 else: 

115 ctx = contextlib.nullcontext() 

116 

117 if user_input and prompts_and_responses: 

118 raise ValueError("Cannot provide both user_input and prompts_and_responses") 

119 

120 if prompts_and_responses: 

121 user_input = ( 

122 ("\n".join(response for (_, response, *_) in prompts_and_responses) + "\n") 

123 .replace("↓", readchar.key.DOWN) 

124 .replace("↑", readchar.key.UP) 

125 ) 

126 

127 with ctx: 

128 result = runner.invoke(app, command, catch_exceptions=False, input=user_input) 

129 

130 if echo: 

131 print("\n------ CLI output ------") 

132 print(result.stdout) 

133 

134 if expected_code is not None: 

135 assertion_error_message = ( 

136 f"Expected code {expected_code} but got {result.exit_code}\n" 

137 "Output from CLI command:\n" 

138 "-----------------------\n" 

139 f"{result.stdout}" 

140 ) 

141 assert result.exit_code == expected_code, assertion_error_message 

142 

143 if expected_output is not None: 

144 output = result.stdout.strip() 

145 expected_output = textwrap.dedent(expected_output).strip() 

146 

147 compare_string = ( 

148 "------ expected ------\n" 

149 f"{expected_output}\n" 

150 "------ actual ------\n" 

151 f"{output}\n" 

152 "------ end ------\n" 

153 ) 

154 assert output == expected_output, compare_string 

155 

156 if prompts_and_responses: 

157 output = result.stdout.strip() 

158 cursor = 0 

159 

160 for item in prompts_and_responses: 

161 prompt = item[0] 

162 selected_option = item[2] if len(item) == 3 else None 

163 

164 prompt_re = rf"{re.escape(prompt)}.*?" 

165 if not selected_option: 

166 # If we're not prompting for a table, then expect that the 

167 # prompt ends with a colon. 

168 prompt_re += ":" 

169 

170 match = re.search(prompt_re, output[cursor:]) 

171 if not match: 

172 raise AssertionError(f"Prompt '{prompt}' not found in CLI output") 

173 cursor = cursor + match.end() 

174 

175 if selected_option: 

176 option_re = re.escape(f"│ > │ {selected_option}") 

177 match = re.search(option_re, output[cursor:]) 

178 if not match: 

179 raise AssertionError( 

180 f"Option '{selected_option}' not found after prompt '{prompt}'" 

181 ) 

182 cursor = cursor + match.end() 

183 

184 if expected_output_contains is not None: 

185 if isinstance(expected_output_contains, str): 

186 check_contains(result, expected_output_contains, should_contain=True) 

187 else: 

188 for contents in expected_output_contains: 

189 check_contains(result, contents, should_contain=True) 

190 

191 if expected_output_does_not_contain is not None: 

192 if isinstance(expected_output_does_not_contain, str): 

193 check_contains( 

194 result, expected_output_does_not_contain, should_contain=False 

195 ) 

196 else: 

197 for contents in expected_output_does_not_contain: 

198 check_contains(result, contents, should_contain=False) 

199 

200 if expected_line_count is not None: 

201 line_count = len(result.stdout.splitlines()) 

202 assert expected_line_count == line_count, ( 

203 f"Expected {expected_line_count} lines of CLI output, only" 

204 f" {line_count} lines present" 

205 ) 

206 

207 return result 

208 

209 

210@contextlib.contextmanager 

211def temporary_console_width(console: Console, width: int): 

212 original = console.width 

213 

214 try: 

215 console._width = width # type: ignore 

216 yield 

217 finally: 

218 console._width = original # type: ignore