|
| 1 | +""" |
| 2 | +A fake execution connector that simulates command execution without touching a |
| 3 | +real target. Both fact-gathering and operation execution funnel through |
| 4 | +``run_shell_command``, so a single interception point drives both. |
| 5 | +
|
| 6 | +Responses are configurable via inventory data (the ``fake_responses`` key), which |
| 7 | +maps a command matcher to a canned result. Anything not matched succeeds with |
| 8 | +empty output by default. |
| 9 | +""" |
| 10 | + |
| 11 | +from __future__ import annotations |
| 12 | + |
| 13 | +import os |
| 14 | +import random |
| 15 | +import re |
| 16 | +from typing import TYPE_CHECKING, Any |
| 17 | + |
| 18 | +import gevent |
| 19 | +from typing_extensions import Unpack, override |
| 20 | + |
| 21 | +from pyinfra import logger |
| 22 | +from pyinfra.api.output import echo |
| 23 | +from pyinfra.connectors.base import BaseConnector, ConnectorData, DataMeta |
| 24 | +from pyinfra.connectors.util import CommandOutput, OutputLine |
| 25 | + |
| 26 | +if TYPE_CHECKING: |
| 27 | + from collections.abc import Iterator |
| 28 | + from io import IOBase |
| 29 | + |
| 30 | + from pyinfra.api.arguments import ConnectorArguments |
| 31 | + from pyinfra.api.command import StringCommand |
| 32 | + |
| 33 | +# Default simulated command duration (seconds), so progress bars behave like a |
| 34 | +# real deploy rather than completing instantly. Override via env vars or per-host |
| 35 | +# ``fake_delay`` / ``fake_delay_jitter`` inventory data. |
| 36 | +DEFAULT_DELAY = float(os.environ.get("PYINFRA_FAKE_DELAY", "0.5")) |
| 37 | +DEFAULT_DELAY_JITTER = float(os.environ.get("PYINFRA_FAKE_DELAY_JITTER", "0.4")) |
| 38 | + |
| 39 | + |
| 40 | +class FakeConnectorData(ConnectorData, total=False): |
| 41 | + fake_responses: dict[str | re.Pattern, Any] |
| 42 | + fake_delay: float |
| 43 | + fake_delay_jitter: float |
| 44 | + |
| 45 | + |
| 46 | +class FakeConnector(BaseConnector): |
| 47 | + """ |
| 48 | + The ``@fake`` connector simulates execution locally without running anything, |
| 49 | + which is handy for demos, screenshots, documentation examples and tests. |
| 50 | +
|
| 51 | + Every command "succeeds" (exit code 0) with empty output by default. Specific |
| 52 | + responses can be scripted via the host ``fake_responses`` data, a mapping of a |
| 53 | + command matcher to a canned response. |
| 54 | +
|
| 55 | + The matcher (mapping key) is either: |
| 56 | +
|
| 57 | + + a ``str`` — matched as a **substring** of the command, or |
| 58 | + + a compiled ``re.Pattern`` (``re.compile(...)``) — matched with |
| 59 | + ``pattern.search(command)``, so regular expressions are supported. |
| 60 | +
|
| 61 | + The response (mapping value) is either: |
| 62 | +
|
| 63 | + + a ``str`` / ``list[str]`` of stdout lines, or |
| 64 | + + a ``dict`` with optional ``stdout`` (``str``/``list``), ``stderr`` |
| 65 | + (``str``/``list``) and ``success`` (``bool``) keys. |
| 66 | +
|
| 67 | + Matchers are tried in insertion order; the first match wins. |
| 68 | +
|
| 69 | + Each simulated command/transfer also takes a short, slightly randomised |
| 70 | + amount of time (so progress bars behave like a real deploy rather than |
| 71 | + finishing instantly). Tune this with the ``fake_delay`` / |
| 72 | + ``fake_delay_jitter`` host data, or the ``PYINFRA_FAKE_DELAY`` / |
| 73 | + ``PYINFRA_FAKE_DELAY_JITTER`` environment variables. Set the delay to |
| 74 | + ``0`` for instant execution (e.g. in tests). |
| 75 | + """ |
| 76 | + |
| 77 | + __examples_doc__ = """ |
| 78 | + Run any command or operation against one or more fake hosts, with no real |
| 79 | + target: |
| 80 | +
|
| 81 | + .. code:: shell |
| 82 | +
|
| 83 | + # A single fake host |
| 84 | + pyinfra @fake exec -- echo "hello world" |
| 85 | +
|
| 86 | + # Multiple named fake hosts (comma separated) |
| 87 | + pyinfra @fake/web-1,@fake/web-2 server.shell "echo hi" |
| 88 | +
|
| 89 | + Script what specific commands return with the ``fake_responses`` host data in |
| 90 | + an inventory file (``inventory.py``). Each key is a matcher, each value the |
| 91 | + canned response: |
| 92 | +
|
| 93 | + .. code:: python |
| 94 | +
|
| 95 | + import re |
| 96 | +
|
| 97 | + hosts = [ |
| 98 | + ( |
| 99 | + "@fake/web-1", |
| 100 | + { |
| 101 | + "fake_responses": { |
| 102 | + # substring match |
| 103 | + "command -v git": {"success": False}, |
| 104 | + "git --version": "git version 2.40.0", |
| 105 | + # regexp match (re.Pattern keys use pattern.search()) |
| 106 | + re.compile(r"^apt-get .*install"): { |
| 107 | + "success": False, |
| 108 | + "stderr": "E: locked", |
| 109 | + }, |
| 110 | + }, |
| 111 | + # instant execution for this host |
| 112 | + "fake_delay": 0, |
| 113 | + }, |
| 114 | + ), |
| 115 | + ] |
| 116 | +
|
| 117 | + A response value is either a ``str`` / ``list[str]`` of stdout lines, or a |
| 118 | + ``dict`` with optional ``stdout``, ``stderr`` and ``success`` keys. Matchers |
| 119 | + are tried in insertion order; the first match wins, and unmatched commands |
| 120 | + succeed with no output. |
| 121 | + """ |
| 122 | + |
| 123 | + handles_execution = True |
| 124 | + |
| 125 | + data_cls = FakeConnectorData |
| 126 | + data_meta = { |
| 127 | + "fake_responses": DataMeta( |
| 128 | + "Mapping of command matcher (substring str or re.Pattern) to a canned response.", |
| 129 | + default={}, |
| 130 | + ), |
| 131 | + "fake_delay": DataMeta( |
| 132 | + "Base duration (seconds) to simulate for each command/transfer.", |
| 133 | + default=DEFAULT_DELAY, |
| 134 | + ), |
| 135 | + "fake_delay_jitter": DataMeta( |
| 136 | + "Extra random duration (seconds, 0..jitter) added to each delay.", |
| 137 | + default=DEFAULT_DELAY_JITTER, |
| 138 | + ), |
| 139 | + } |
| 140 | + |
| 141 | + #: Commands executed against this connector instance, recorded for tests/debug. |
| 142 | + executed_commands: list[str] |
| 143 | + |
| 144 | + def __init__(self, state, host) -> None: |
| 145 | + super().__init__(state, host) |
| 146 | + self.executed_commands = [] |
| 147 | + |
| 148 | + @override |
| 149 | + @staticmethod |
| 150 | + def make_names_data(name: str | None = None) -> Iterator[tuple[str, dict, list[str]]]: |
| 151 | + if not name: |
| 152 | + yield "@fake", {}, ["@fake"] |
| 153 | + return |
| 154 | + |
| 155 | + for sub_name in name.split(","): |
| 156 | + sub_name = sub_name.strip() |
| 157 | + if not sub_name: |
| 158 | + continue |
| 159 | + yield f"@fake/{sub_name}", {}, ["@fake"] |
| 160 | + |
| 161 | + def _lookup_response(self, command_str: str) -> tuple[bool, CommandOutput]: |
| 162 | + fake_responses = self.data.get("fake_responses") or {} |
| 163 | + |
| 164 | + for matcher, response in fake_responses.items(): |
| 165 | + if isinstance(matcher, re.Pattern): |
| 166 | + if not matcher.search(command_str): |
| 167 | + continue |
| 168 | + elif matcher not in command_str: |
| 169 | + continue |
| 170 | + |
| 171 | + success = True |
| 172 | + stdout: list[str] = [] |
| 173 | + stderr: list[str] = [] |
| 174 | + |
| 175 | + if isinstance(response, dict): |
| 176 | + success = bool(response.get("success", True)) |
| 177 | + stdout = _as_lines(response.get("stdout")) |
| 178 | + stderr = _as_lines(response.get("stderr")) |
| 179 | + else: |
| 180 | + stdout = _as_lines(response) |
| 181 | + |
| 182 | + lines = [OutputLine("stdout", line) for line in stdout] |
| 183 | + lines += [OutputLine("stderr", line) for line in stderr] |
| 184 | + return success, CommandOutput(lines) |
| 185 | + |
| 186 | + # Default: succeed with no output. |
| 187 | + return True, CommandOutput([]) |
| 188 | + |
| 189 | + def _sleep(self) -> None: |
| 190 | + """Simulate a realistic, non-instant task duration. |
| 191 | +
|
| 192 | + Uses ``gevent.sleep`` so other host greenlets and the progress bar |
| 193 | + continue to run cooperatively while this "command" is in flight. |
| 194 | + """ |
| 195 | + delay = self.data.get("fake_delay") |
| 196 | + if delay is None: |
| 197 | + delay = DEFAULT_DELAY |
| 198 | + jitter = self.data.get("fake_delay_jitter") |
| 199 | + if jitter is None: |
| 200 | + jitter = DEFAULT_DELAY_JITTER |
| 201 | + |
| 202 | + duration = float(delay) + random.uniform(0, max(0.0, float(jitter))) |
| 203 | + if duration > 0: |
| 204 | + gevent.sleep(duration) |
| 205 | + |
| 206 | + @override |
| 207 | + def run_shell_command( |
| 208 | + self, |
| 209 | + command: str | StringCommand, |
| 210 | + print_output: bool = False, |
| 211 | + print_input: bool = False, |
| 212 | + **arguments: Unpack[ConnectorArguments], |
| 213 | + ) -> tuple[bool, CommandOutput]: |
| 214 | + if isinstance(command, str): |
| 215 | + command_str = command |
| 216 | + else: |
| 217 | + command_str = command.get_masked_value() |
| 218 | + self.executed_commands.append(command_str) |
| 219 | + |
| 220 | + logger.debug("[fake] simulating command on %s: %s", self.host.name, command_str) |
| 221 | + |
| 222 | + if print_input: |
| 223 | + echo(f"{self.host.print_prefix}>>> {command_str}", err=True) |
| 224 | + |
| 225 | + self._sleep() |
| 226 | + |
| 227 | + status, output = self._lookup_response(command_str) |
| 228 | + |
| 229 | + if print_output: |
| 230 | + for line in output.output_lines: |
| 231 | + echo(f"{self.host.print_prefix}{line}", err=True) |
| 232 | + |
| 233 | + return status, output |
| 234 | + |
| 235 | + @override |
| 236 | + def put_file( |
| 237 | + self, |
| 238 | + filename_or_io: str | IOBase, |
| 239 | + remote_filename: str, |
| 240 | + remote_temp_filename: str | None = None, |
| 241 | + print_output: bool = False, |
| 242 | + print_input: bool = False, |
| 243 | + **arguments: Unpack[ConnectorArguments], |
| 244 | + ) -> bool: |
| 245 | + logger.debug("[fake] simulating put_file on %s: %s", self.host.name, remote_filename) |
| 246 | + self._sleep() |
| 247 | + return True |
| 248 | + |
| 249 | + @override |
| 250 | + def get_file( |
| 251 | + self, |
| 252 | + remote_filename: str, |
| 253 | + filename_or_io: str | IOBase, |
| 254 | + remote_temp_filename: str | None = None, |
| 255 | + print_output: bool = False, |
| 256 | + print_input: bool = False, |
| 257 | + **arguments: Unpack[ConnectorArguments], |
| 258 | + ) -> bool: |
| 259 | + logger.debug("[fake] simulating get_file on %s: %s", self.host.name, remote_filename) |
| 260 | + self._sleep() |
| 261 | + return True |
| 262 | + |
| 263 | + |
| 264 | +def _as_lines(value: Any) -> list[str]: |
| 265 | + if value is None: |
| 266 | + return [] |
| 267 | + if isinstance(value, str): |
| 268 | + return value.splitlines() |
| 269 | + if isinstance(value, (list, tuple)): |
| 270 | + return [str(line) for line in value] |
| 271 | + return [str(value)] |
0 commit comments