gabriel / muse public
cli_test_helper.py python
203 lines 6.7 KB
Raw
sha256:660fcac1df3ab28f61862e961890bd2ca8b754fa0242079d93ca1e25037ec8a6 chore(tests): add docstring to tests/__init__.py so rc14 tr… Human 26 days ago
1 """Argparse-compatible CliRunner replacement for Muse test suite.
2
3 Replaces ``typer.testing.CliRunner`` so tests can call ``runner.invoke(cli,
4 args)`` without modification after the typer → argparse migration. The first
5 argument (``cli``) is always ``None`` (a stub) after migration; it is accepted
6 but ignored, and ``muse.cli.app.main`` is always the target.
7 """
8
9 from __future__ import annotations
10
11 import contextlib
12 import io
13 import os
14 import re
15 import sys
16 import threading
17 import traceback
18
19 from muse.cli.app import main
20 from muse.core._types import Manifest
21
22 type _EnvSaved = dict[str, str | None]
23
24 _ANSI_ESCAPE = re.compile(r"\x1b\[[0-9;]*m")
25
26
27 def _strip_ansi(text: str) -> str:
28 """Remove ANSI escape sequences — typer's CliRunner did this automatically."""
29 return _ANSI_ESCAPE.sub("", text)
30
31
32 class _StdinWithBuffer(io.StringIO):
33 """Text-mode stdin backed by StringIO, with a ``.buffer`` BytesIO sibling.
34
35 Accepts either a ``str`` (text) or ``bytes`` (binary). When binary is
36 provided the text surface is left empty — commands that read from
37 ``sys.stdin.buffer`` (e.g. ``unpack-objects``) get the raw bytes while
38 text-mode reads get an empty string.
39 """
40
41 def __init__(self, text: str | bytes) -> None:
42 if isinstance(text, bytes):
43 super().__init__("")
44 self.buffer = io.BytesIO(text)
45 else:
46 super().__init__(text)
47 self.buffer = io.BytesIO(text.encode())
48
49 def isatty(self) -> bool:
50 return False
51
52
53 class _StdoutCapture(io.StringIO):
54 """Text-mode stdout backed by StringIO, with a ``.buffer`` BytesIO sibling.
55
56 Some plumbing commands (e.g. ``cat-object``) write raw bytes to
57 ``sys.stdout.buffer``. Subclassing ``StringIO`` makes this assignable to
58 ``sys.stdout`` (and passable to ``contextlib.redirect_stdout``) without
59 any type annotation workaround. Binary output is decoded and appended to
60 the text output in ``getvalue()``.
61 """
62
63 def __init__(self) -> None:
64 super().__init__()
65 self.buffer = io.BytesIO()
66
67 def isatty(self) -> bool:
68 return False
69
70 def getvalue(self) -> str:
71 text_out = super().getvalue()
72 bytes_out = self.buffer.getvalue()
73 if bytes_out:
74 try:
75 text_out += bytes_out.decode("utf-8", errors="replace")
76 except Exception:
77 pass
78 return text_out
79
80
81 def _restore_env(saved: _EnvSaved) -> None:
82 """Restore environment variables to their pre-invoke state."""
83 for k, orig in saved.items():
84 if orig is None:
85 os.environ.pop(k, None)
86 else:
87 os.environ[k] = orig
88
89
90 class InvokeResult:
91 """Mirrors the fields that typer.testing.Result exposed."""
92
93 def __init__(
94 self,
95 exit_code: int,
96 output: str,
97 stderr_output: str = "",
98 stdout_bytes: bytes = b"",
99 ) -> None:
100 self.exit_code = exit_code
101 self.output = output
102 self.stdout = output
103 self.stderr = stderr_output
104 self.stdout_bytes = stdout_bytes
105 self.exception: BaseException | None = None
106
107 def __repr__(self) -> str:
108 return f"InvokeResult(exit_code={self.exit_code}, output={self.output!r})"
109
110
111 class CliRunner:
112 """Drop-in replacement for ``typer.testing.CliRunner``.
113
114 Captures stdout and stderr, calls ``main(args)``, and returns an
115 ``InvokeResult`` whose interface matches the typer equivalent closely
116 enough for the existing test suite to run without changes.
117
118 Honoured parameters:
119 - ``env``: key/value pairs set in ``os.environ`` for the duration of the
120 call and restored afterward.
121 - ``input``: string fed to ``sys.stdin`` (needed by ``unpack-objects``).
122 - ``catch_exceptions``: when False, exceptions propagate to the caller.
123
124 Thread safety: ``contextlib.redirect_stdout`` mutates ``sys.stdout``
125 globally, so concurrent ``invoke`` calls on different threads would
126 overwrite each other's capture buffer. A class-level lock serialises
127 the stdout/stderr redirect + main() call so each invocation gets its
128 own isolated capture.
129 """
130
131 _lock: threading.Lock = threading.Lock()
132
133 def invoke(
134 self,
135 _cli: None,
136 args: list[str],
137 catch_exceptions: bool = True,
138 input: str | bytes | None = None,
139 env: Manifest | None = None,
140 ) -> InvokeResult:
141 """Invoke ``main(args)`` and return captured output + exit code."""
142 # Apply caller-supplied env overrides; restore originals when done.
143 stdout_cap = _StdoutCapture()
144 stderr_buf = io.StringIO()
145 exit_code = 0
146
147 with self._lock:
148 # All global-state mutations (env, sys.stdin, sys.stdout, sys.stderr)
149 # are confined to the lock so concurrent invoke() calls on different
150 # threads do not race on shared process-global variables.
151 saved: _EnvSaved = {}
152 if env:
153 for k, v in env.items():
154 saved[k] = os.environ.get(k)
155 os.environ[k] = v
156
157 orig_stdin = sys.stdin
158 if input is not None:
159 sys.stdin = _StdinWithBuffer(input)
160
161 try:
162 with contextlib.redirect_stdout(stdout_cap):
163 with contextlib.redirect_stderr(stderr_buf):
164 main(list(args))
165 except SystemExit as exc:
166 raw = exc.code
167 if isinstance(raw, int):
168 exit_code = raw
169 elif hasattr(raw, "value"):
170 exit_code = int(raw.value)
171 elif raw is None:
172 exit_code = 0
173 else:
174 exit_code = int(raw)
175 except Exception as exc:
176 if not catch_exceptions:
177 sys.stdin = orig_stdin
178 _restore_env(saved)
179 raise
180 stderr_buf.write(traceback.format_exc())
181 exit_code = 1
182 result = InvokeResult(
183 exit_code,
184 _strip_ansi(stdout_cap.getvalue() + stderr_buf.getvalue()),
185 _strip_ansi(stderr_buf.getvalue()),
186 stdout_bytes=stdout_cap.buffer.getvalue(),
187 )
188 result.exception = exc
189 sys.stdin = orig_stdin
190 _restore_env(saved)
191 return result
192 finally:
193 sys.stdin = orig_stdin
194 _restore_env(saved)
195
196 raw_bytes = stdout_cap.buffer.getvalue()
197 combined = _strip_ansi(stdout_cap.getvalue() + stderr_buf.getvalue())
198 return InvokeResult(
199 exit_code,
200 combined,
201 _strip_ansi(stderr_buf.getvalue()),
202 stdout_bytes=raw_bytes,
203 )
File History 5 commits
sha256:660fcac1df3ab28f61862e961890bd2ca8b754fa0242079d93ca1e25037ec8a6 chore(tests): add docstring to tests/__init__.py so rc14 tr… Human 26 days ago
sha256:d8316ffae901be06347e16ab55be11868eb519dd16ade3e8aa16a99e662f7e62 baseline: rc14 re-baseline after rc3 store corruption recovery Human patch 26 days ago
sha256:50d413a635f57761c3fce1f70251b957b6423821525bf330747ef4007339222e Merge branch 'dev' into main Human 29 days ago
sha256:fb67fed5a4d3e40de84bdd163de94ef1386570bef1dd1a020a732c8a038962ce Merge branch 'dev' into main Human 48 days ago
sha256:1c4b3e3a9a1f300774c3ee662b572a698d5fd405bf765a71e6011a2e9c3eaaaa feat: Muse — version control for the agent era Human 100 days ago