testing.py 19 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574
  1. from __future__ import annotations
  2. import collections.abc as cabc
  3. import contextlib
  4. import io
  5. import os
  6. import shlex
  7. import sys
  8. import tempfile
  9. import typing as t
  10. from types import TracebackType
  11. from . import _compat
  12. from . import formatting
  13. from . import termui
  14. from . import utils
  15. from ._compat import _find_binary_reader
  16. if t.TYPE_CHECKING:
  17. from _typeshed import ReadableBuffer
  18. from .core import Command
  19. class EchoingStdin:
  20. def __init__(self, input: t.BinaryIO, output: t.BinaryIO) -> None:
  21. self._input = input
  22. self._output = output
  23. self._paused = False
  24. def __getattr__(self, x: str) -> t.Any:
  25. return getattr(self._input, x)
  26. def _echo(self, rv: bytes) -> bytes:
  27. if not self._paused:
  28. self._output.write(rv)
  29. return rv
  30. def read(self, n: int = -1) -> bytes:
  31. return self._echo(self._input.read(n))
  32. def read1(self, n: int = -1) -> bytes:
  33. return self._echo(self._input.read1(n)) # type: ignore
  34. def readline(self, n: int = -1) -> bytes:
  35. return self._echo(self._input.readline(n))
  36. def readlines(self) -> list[bytes]:
  37. return [self._echo(x) for x in self._input.readlines()]
  38. def __iter__(self) -> cabc.Iterator[bytes]:
  39. return iter(self._echo(x) for x in self._input)
  40. def __repr__(self) -> str:
  41. return repr(self._input)
  42. @contextlib.contextmanager
  43. def _pause_echo(stream: EchoingStdin | None) -> cabc.Iterator[None]:
  44. if stream is None:
  45. yield
  46. else:
  47. stream._paused = True
  48. yield
  49. stream._paused = False
  50. class BytesIOCopy(io.BytesIO):
  51. """Patch ``io.BytesIO`` to let the written stream be copied to another.
  52. .. versionadded:: 8.2
  53. """
  54. def __init__(self, copy_to: io.BytesIO) -> None:
  55. super().__init__()
  56. self.copy_to = copy_to
  57. def flush(self) -> None:
  58. super().flush()
  59. self.copy_to.flush()
  60. def write(self, b: ReadableBuffer) -> int:
  61. self.copy_to.write(b)
  62. return super().write(b)
  63. class StreamMixer:
  64. """Mixes `<stdout>` and `<stderr>` streams.
  65. The result is available in the ``output`` attribute.
  66. .. versionadded:: 8.2
  67. """
  68. def __init__(self) -> None:
  69. self.output: io.BytesIO = io.BytesIO()
  70. self.stdout: io.BytesIO = BytesIOCopy(copy_to=self.output)
  71. self.stderr: io.BytesIO = BytesIOCopy(copy_to=self.output)
  72. class _NamedTextIOWrapper(io.TextIOWrapper):
  73. def __init__(
  74. self, buffer: t.BinaryIO, name: str, mode: str, **kwargs: t.Any
  75. ) -> None:
  76. super().__init__(buffer, **kwargs)
  77. self._name = name
  78. self._mode = mode
  79. def close(self) -> None:
  80. """
  81. The buffer this object contains belongs to some other object, so
  82. prevent the default __del__ implementation from closing that buffer.
  83. .. versionadded:: 8.3.2
  84. """
  85. ...
  86. @property
  87. def name(self) -> str:
  88. return self._name
  89. @property
  90. def mode(self) -> str:
  91. return self._mode
  92. def make_input_stream(
  93. input: str | bytes | t.IO[t.Any] | None, charset: str
  94. ) -> t.BinaryIO:
  95. # Is already an input stream.
  96. if hasattr(input, "read"):
  97. rv = _find_binary_reader(t.cast("t.IO[t.Any]", input))
  98. if rv is not None:
  99. return rv
  100. raise TypeError("Could not find binary reader for input stream.")
  101. if input is None:
  102. input = b""
  103. elif isinstance(input, str):
  104. input = input.encode(charset)
  105. return io.BytesIO(input)
  106. class Result:
  107. """Holds the captured result of an invoked CLI script.
  108. :param runner: The runner that created the result
  109. :param stdout_bytes: The standard output as bytes.
  110. :param stderr_bytes: The standard error as bytes.
  111. :param output_bytes: A mix of ``stdout_bytes`` and ``stderr_bytes``, as the
  112. user would see it in its terminal.
  113. :param return_value: The value returned from the invoked command.
  114. :param exit_code: The exit code as integer.
  115. :param exception: The exception that happened if one did.
  116. :param exc_info: Exception information (exception type, exception instance,
  117. traceback type).
  118. .. versionchanged:: 8.2
  119. ``stderr_bytes`` no longer optional, ``output_bytes`` introduced and
  120. ``mix_stderr`` has been removed.
  121. .. versionadded:: 8.0
  122. Added ``return_value``.
  123. """
  124. def __init__(
  125. self,
  126. runner: CliRunner,
  127. stdout_bytes: bytes,
  128. stderr_bytes: bytes,
  129. output_bytes: bytes,
  130. return_value: t.Any,
  131. exit_code: int,
  132. exception: BaseException | None,
  133. exc_info: tuple[type[BaseException], BaseException, TracebackType]
  134. | None = None,
  135. ):
  136. self.runner = runner
  137. self.stdout_bytes = stdout_bytes
  138. self.stderr_bytes = stderr_bytes
  139. self.output_bytes = output_bytes
  140. self.return_value = return_value
  141. self.exit_code = exit_code
  142. self.exception = exception
  143. self.exc_info = exc_info
  144. @property
  145. def output(self) -> str:
  146. """The terminal output as unicode string, as the user would see it.
  147. .. versionchanged:: 8.2
  148. No longer a proxy for ``self.stdout``. Now has its own independent stream
  149. that is mixing `<stdout>` and `<stderr>`, in the order they were written.
  150. """
  151. return self.output_bytes.decode(self.runner.charset, "replace").replace(
  152. "\r\n", "\n"
  153. )
  154. @property
  155. def stdout(self) -> str:
  156. """The standard output as unicode string."""
  157. return self.stdout_bytes.decode(self.runner.charset, "replace").replace(
  158. "\r\n", "\n"
  159. )
  160. @property
  161. def stderr(self) -> str:
  162. """The standard error as unicode string.
  163. .. versionchanged:: 8.2
  164. No longer raise an exception, always returns the `<stderr>` string.
  165. """
  166. return self.stderr_bytes.decode(self.runner.charset, "replace").replace(
  167. "\r\n", "\n"
  168. )
  169. def __repr__(self) -> str:
  170. exc_str = repr(self.exception) if self.exception else "okay"
  171. return f"<{type(self).__name__} {exc_str}>"
  172. class CliRunner:
  173. """The CLI runner provides functionality to invoke a Click command line
  174. script for unittesting purposes in a isolated environment. This only
  175. works in single-threaded systems without any concurrency as it changes the
  176. global interpreter state.
  177. :param charset: the character set for the input and output data.
  178. :param env: a dictionary with environment variables for overriding.
  179. :param echo_stdin: if this is set to `True`, then reading from `<stdin>` writes
  180. to `<stdout>`. This is useful for showing examples in
  181. some circumstances. Note that regular prompts
  182. will automatically echo the input.
  183. :param catch_exceptions: Whether to catch any exceptions other than
  184. ``SystemExit`` when running :meth:`~CliRunner.invoke`.
  185. .. versionchanged:: 8.2
  186. Added the ``catch_exceptions`` parameter.
  187. .. versionchanged:: 8.2
  188. ``mix_stderr`` parameter has been removed.
  189. """
  190. def __init__(
  191. self,
  192. charset: str = "utf-8",
  193. env: cabc.Mapping[str, str | None] | None = None,
  194. echo_stdin: bool = False,
  195. catch_exceptions: bool = True,
  196. ) -> None:
  197. self.charset = charset
  198. self.env: cabc.Mapping[str, str | None] = env or {}
  199. self.echo_stdin = echo_stdin
  200. self.catch_exceptions = catch_exceptions
  201. def get_default_prog_name(self, cli: Command) -> str:
  202. """Given a command object it will return the default program name
  203. for it. The default is the `name` attribute or ``"root"`` if not
  204. set.
  205. """
  206. return cli.name or "root"
  207. def make_env(
  208. self, overrides: cabc.Mapping[str, str | None] | None = None
  209. ) -> cabc.Mapping[str, str | None]:
  210. """Returns the environment overrides for invoking a script."""
  211. rv = dict(self.env)
  212. if overrides:
  213. rv.update(overrides)
  214. return rv
  215. @contextlib.contextmanager
  216. def isolation(
  217. self,
  218. input: str | bytes | t.IO[t.Any] | None = None,
  219. env: cabc.Mapping[str, str | None] | None = None,
  220. color: bool = False,
  221. ) -> cabc.Iterator[tuple[io.BytesIO, io.BytesIO, io.BytesIO]]:
  222. """A context manager that sets up the isolation for invoking of a
  223. command line tool. This sets up `<stdin>` with the given input data
  224. and `os.environ` with the overrides from the given dictionary.
  225. This also rebinds some internals in Click to be mocked (like the
  226. prompt functionality).
  227. This is automatically done in the :meth:`invoke` method.
  228. :param input: the input stream to put into `sys.stdin`.
  229. :param env: the environment overrides as dictionary.
  230. :param color: whether the output should contain color codes. The
  231. application can still override this explicitly.
  232. .. versionadded:: 8.2
  233. An additional output stream is returned, which is a mix of
  234. `<stdout>` and `<stderr>` streams.
  235. .. versionchanged:: 8.2
  236. Always returns the `<stderr>` stream.
  237. .. versionchanged:: 8.0
  238. `<stderr>` is opened with ``errors="backslashreplace"``
  239. instead of the default ``"strict"``.
  240. .. versionchanged:: 4.0
  241. Added the ``color`` parameter.
  242. """
  243. bytes_input = make_input_stream(input, self.charset)
  244. echo_input = None
  245. old_stdin = sys.stdin
  246. old_stdout = sys.stdout
  247. old_stderr = sys.stderr
  248. old_forced_width = formatting.FORCED_WIDTH
  249. formatting.FORCED_WIDTH = 80
  250. env = self.make_env(env)
  251. stream_mixer = StreamMixer()
  252. if self.echo_stdin:
  253. bytes_input = echo_input = t.cast(
  254. t.BinaryIO, EchoingStdin(bytes_input, stream_mixer.stdout)
  255. )
  256. sys.stdin = text_input = _NamedTextIOWrapper(
  257. bytes_input, encoding=self.charset, name="<stdin>", mode="r"
  258. )
  259. if self.echo_stdin:
  260. # Force unbuffered reads, otherwise TextIOWrapper reads a
  261. # large chunk which is echoed early.
  262. text_input._CHUNK_SIZE = 1 # type: ignore
  263. sys.stdout = _NamedTextIOWrapper(
  264. stream_mixer.stdout, encoding=self.charset, name="<stdout>", mode="w"
  265. )
  266. sys.stderr = _NamedTextIOWrapper(
  267. stream_mixer.stderr,
  268. encoding=self.charset,
  269. name="<stderr>",
  270. mode="w",
  271. errors="backslashreplace",
  272. )
  273. @_pause_echo(echo_input) # type: ignore
  274. def visible_input(prompt: str | None = None) -> str:
  275. sys.stdout.write(prompt or "")
  276. try:
  277. val = next(text_input).rstrip("\r\n")
  278. except StopIteration as e:
  279. raise EOFError() from e
  280. sys.stdout.write(f"{val}\n")
  281. sys.stdout.flush()
  282. return val
  283. @_pause_echo(echo_input) # type: ignore
  284. def hidden_input(prompt: str | None = None) -> str:
  285. sys.stdout.write(f"{prompt or ''}\n")
  286. sys.stdout.flush()
  287. try:
  288. return next(text_input).rstrip("\r\n")
  289. except StopIteration as e:
  290. raise EOFError() from e
  291. @_pause_echo(echo_input) # type: ignore
  292. def _getchar(echo: bool) -> str:
  293. char = sys.stdin.read(1)
  294. if echo:
  295. sys.stdout.write(char)
  296. sys.stdout.flush()
  297. return char
  298. default_color = color
  299. def should_strip_ansi(
  300. stream: t.IO[t.Any] | None = None, color: bool | None = None
  301. ) -> bool:
  302. if color is None:
  303. return not default_color
  304. return not color
  305. old_visible_prompt_func = termui.visible_prompt_func
  306. old_hidden_prompt_func = termui.hidden_prompt_func
  307. old__getchar_func = termui._getchar
  308. old_should_strip_ansi = utils.should_strip_ansi # type: ignore
  309. old__compat_should_strip_ansi = _compat.should_strip_ansi
  310. termui.visible_prompt_func = visible_input
  311. termui.hidden_prompt_func = hidden_input
  312. termui._getchar = _getchar
  313. utils.should_strip_ansi = should_strip_ansi # type: ignore
  314. _compat.should_strip_ansi = should_strip_ansi
  315. old_env = {}
  316. try:
  317. for key, value in env.items():
  318. old_env[key] = os.environ.get(key)
  319. if value is None:
  320. try:
  321. del os.environ[key]
  322. except Exception:
  323. pass
  324. else:
  325. os.environ[key] = value
  326. yield (stream_mixer.stdout, stream_mixer.stderr, stream_mixer.output)
  327. finally:
  328. for key, value in old_env.items():
  329. if value is None:
  330. try:
  331. del os.environ[key]
  332. except Exception:
  333. pass
  334. else:
  335. os.environ[key] = value
  336. sys.stdout = old_stdout
  337. sys.stderr = old_stderr
  338. sys.stdin = old_stdin
  339. termui.visible_prompt_func = old_visible_prompt_func
  340. termui.hidden_prompt_func = old_hidden_prompt_func
  341. termui._getchar = old__getchar_func
  342. utils.should_strip_ansi = old_should_strip_ansi # type: ignore
  343. _compat.should_strip_ansi = old__compat_should_strip_ansi
  344. formatting.FORCED_WIDTH = old_forced_width
  345. def invoke(
  346. self,
  347. cli: Command,
  348. args: str | cabc.Sequence[str] | None = None,
  349. input: str | bytes | t.IO[t.Any] | None = None,
  350. env: cabc.Mapping[str, str | None] | None = None,
  351. catch_exceptions: bool | None = None,
  352. color: bool = False,
  353. **extra: t.Any,
  354. ) -> Result:
  355. """Invokes a command in an isolated environment. The arguments are
  356. forwarded directly to the command line script, the `extra` keyword
  357. arguments are passed to the :meth:`~clickpkg.Command.main` function of
  358. the command.
  359. This returns a :class:`Result` object.
  360. :param cli: the command to invoke
  361. :param args: the arguments to invoke. It may be given as an iterable
  362. or a string. When given as string it will be interpreted
  363. as a Unix shell command. More details at
  364. :func:`shlex.split`.
  365. :param input: the input data for `sys.stdin`.
  366. :param env: the environment overrides.
  367. :param catch_exceptions: Whether to catch any other exceptions than
  368. ``SystemExit``. If :data:`None`, the value
  369. from :class:`CliRunner` is used.
  370. :param extra: the keyword arguments to pass to :meth:`main`.
  371. :param color: whether the output should contain color codes. The
  372. application can still override this explicitly.
  373. .. versionadded:: 8.2
  374. The result object has the ``output_bytes`` attribute with
  375. the mix of ``stdout_bytes`` and ``stderr_bytes``, as the user would
  376. see it in its terminal.
  377. .. versionchanged:: 8.2
  378. The result object always returns the ``stderr_bytes`` stream.
  379. .. versionchanged:: 8.0
  380. The result object has the ``return_value`` attribute with
  381. the value returned from the invoked command.
  382. .. versionchanged:: 4.0
  383. Added the ``color`` parameter.
  384. .. versionchanged:: 3.0
  385. Added the ``catch_exceptions`` parameter.
  386. .. versionchanged:: 3.0
  387. The result object has the ``exc_info`` attribute with the
  388. traceback if available.
  389. """
  390. exc_info = None
  391. if catch_exceptions is None:
  392. catch_exceptions = self.catch_exceptions
  393. with self.isolation(input=input, env=env, color=color) as outstreams:
  394. return_value = None
  395. exception: BaseException | None = None
  396. exit_code = 0
  397. if isinstance(args, str):
  398. args = shlex.split(args)
  399. try:
  400. prog_name = extra.pop("prog_name")
  401. except KeyError:
  402. prog_name = self.get_default_prog_name(cli)
  403. try:
  404. return_value = cli.main(args=args or (), prog_name=prog_name, **extra)
  405. except SystemExit as e:
  406. exc_info = sys.exc_info()
  407. e_code = t.cast("int | t.Any | None", e.code)
  408. if e_code is None:
  409. e_code = 0
  410. if e_code != 0:
  411. exception = e
  412. if not isinstance(e_code, int):
  413. sys.stdout.write(str(e_code))
  414. sys.stdout.write("\n")
  415. e_code = 1
  416. exit_code = e_code
  417. except Exception as e:
  418. if not catch_exceptions:
  419. raise
  420. exception = e
  421. exit_code = 1
  422. exc_info = sys.exc_info()
  423. finally:
  424. sys.stdout.flush()
  425. sys.stderr.flush()
  426. stdout = outstreams[0].getvalue()
  427. stderr = outstreams[1].getvalue()
  428. output = outstreams[2].getvalue()
  429. return Result(
  430. runner=self,
  431. stdout_bytes=stdout,
  432. stderr_bytes=stderr,
  433. output_bytes=output,
  434. return_value=return_value,
  435. exit_code=exit_code,
  436. exception=exception,
  437. exc_info=exc_info, # type: ignore
  438. )
  439. @contextlib.contextmanager
  440. def isolated_filesystem(
  441. self, temp_dir: str | os.PathLike[str] | None = None
  442. ) -> cabc.Iterator[str]:
  443. """A context manager that creates a temporary directory and
  444. changes the current working directory to it. This isolates tests
  445. that affect the contents of the CWD to prevent them from
  446. interfering with each other.
  447. :param temp_dir: Create the temporary directory under this
  448. directory. If given, the created directory is not removed
  449. when exiting.
  450. .. versionchanged:: 8.0
  451. Added the ``temp_dir`` parameter.
  452. """
  453. cwd = os.getcwd()
  454. dt = tempfile.mkdtemp(dir=temp_dir)
  455. os.chdir(dt)
  456. try:
  457. yield dt
  458. finally:
  459. os.chdir(cwd)
  460. if temp_dir is None:
  461. import shutil
  462. try:
  463. shutil.rmtree(dt)
  464. except OSError:
  465. pass