diff --git a/.gitignore b/.gitignore index ee9136a..40a9bf2 100644 --- a/.gitignore +++ b/.gitignore @@ -6,6 +6,7 @@ .dk-installer demo-config.json dk-*-credentials.txt +dk-*-install.json # Docker docker-compose.yml @@ -24,3 +25,10 @@ build/ dist/ *.egg-info/ dk-installer.spec + +# Claude +.claude/.local/ + +# UV +bin/ +uv.lock diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..30f513e --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,101 @@ +# CLAUDE.md + +This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. + +## Repository purpose + +This repo ships **one artifact**: `dk-installer.py` — a single-file, stdlib-only Python script that end users download and run to install/upgrade/start/demo the open-source DataKitchen products (TestGen and DataOps Observability) locally. TestGen supports both Docker Compose and pip-via-uv install modes; Observability is Docker-only. On Windows it is also packaged as `dk-installer.exe` via PyInstaller (see `.github/workflows/release_exe.yml`). + +The `demo/` directory is a separate deliverable: it is built into the `datakitchen/data-observability-demo` Docker image that `dk-installer.py` pulls at runtime to generate demo data. It is **not** imported by the installer. + +## Common commands + +```bash +pip install .[dev,test] # install ruff + pytest (project has no runtime deps) + +ruff check --show-fixes # lint (CI-enforced) +ruff format --check --diff # format check (CI-enforced) +ruff format # apply formatting + +pytest # run full test suite +pytest tests/test_tg_install.py # single file +pytest tests/test_action.py::test_name # single test +pytest -m unit # only unit-marked tests +pytest -m integration # only integration-marked tests +pytest --cov --cov-report=term-missing # with coverage (matches CI) + +python3 dk-installer.py --help # see installer CLI +python3 dk-installer.py tg install # run an action locally during dev +``` + +Building the Windows `.exe` happens automatically on every push to `main` (`release_exe.yml` → PyInstaller → GitHub Release tagged `latest`). For local builds on Windows, see `docs/build_windows_installer.md`. + +## Architecture: how `dk-installer.py` is organized + +The installer is a ~2300-line single file intentionally using only the Python stdlib — users run it without installing any packages. Do not introduce third-party runtime dependencies. + +Core abstractions (all in `dk-installer.py`): + +- **`Installer`** — top-level argparse wrapper. `get_installer_instance()` at the bottom of the file registers the two products (`obs`, `tg`) and their actions. Each product sets compose-file defaults (`compose_file_name`, `compose_project_name`) that flow into actions via argparse `set_defaults`. +- **`Action`** — base class for one CLI subcommand (e.g., `tg install`). Owns session-scoped concerns: creates a timestamped log folder under `.dk-installer/` (or `%LOCALAPPDATA%/DataKitchenApps/` on Windows), configures logging, zips logs on exit, wraps execution in `AnalyticsWrapper`, enforces `requirements` (list of `Requirement` objects that shell out to check `docker`, `docker compose`, etc.), and provides `run_cmd` / `run_cmd_retries` — always use these rather than raw `subprocess` so output is captured per-command into the session zip. +- **`MultiStepAction`** — `Action` subclass that declares a `steps: list[type[Step]]`. Each `Step` has `pre_execute` (run for all steps before any executes — validation phase) then `execute` (the actual work). On any step failure, remaining steps are skipped and `on_action_fail` runs in reverse order; on success, `on_action_success` runs in reverse order. **Most install/upgrade actions are `MultiStepAction`s** — when adding a new install phase, write a new `Step` class and add it to the list. +- **`Step`** — unit of work inside a `MultiStepAction`. Steps share state via `action.ctx` (a dict on the parent action). Raising `SkipStep` from `execute` marks it SKIPPED; raising any other exception marks it FAILED and aborts the action if `required = True`. +- **`ComposeActionMixin` / `ComposeDeleteAction` / `ComposePullImagesStep` / `ComposeStartStep` / `CreateComposeFileStepBase`** — shared building blocks for both products. `Obs*` and `TestGen*` classes specialize these. +- **`AnalyticsWrapper`** — sends anonymous Mixpanel events for each action (disabled with `--no-analytics` or `DK_INSTALLER_ANALYTICS=no`). Instance ID is persisted to `.dk-installer/instance.txt`. Don't log PII here. +- **`Console`** (global `CONSOLE`) — all user-facing output goes through this; don't use bare `print` for user messages (the menu code and `collect_user_input` are the exceptions). +- **`Menu`** / `show_menu` — only used when the frozen Windows `.exe` is launched with no arguments (double-click). Not part of the CLI flow on Unix. + +The action registry in `get_installer_instance()` is the authoritative list of user-facing commands — to add a new command, add an `Action` subclass there. + +### TestGen install modes + +TestGen has two install modes: `docker` (Compose) and `pip` (uv-managed venv with embedded Postgres). Mode is recorded at install time in a JSON marker file (`dk-tg-install.json`) so `tg upgrade` / `tg delete` / `tg start` / `tg run-demo` / `tg delete-demo` know which path to take. + +The five `Testgen*Action` classes that span both modes follow a unified pattern: +- `_per_invocation_attrs` includes `_resolved_mode` (and `steps` / `intro_text` for the `MultiStepAction`-based ones) so menu re-runs start clean. +- `check_requirements` resolves mode once via `_resolve_install_mode`, then calls `super().check_requirements`. +- `_resolve_install_mode` reads the marker (or runs auto-detect for `install`), sets `self._resolved_mode`, optionally records `analytics["install_mode"]`. Install/upgrade/start/run-demo abort when no install exists; delete and delete-demo are idempotent (return rather than raise). +- `get_requirements` reads `self._resolved_mode` — Docker reqs only when in Docker mode. +- `execute` branches on `self._resolved_mode`. For `MultiStepAction` subclasses, `self.steps` is also swapped at resolution time using class-level `pip_steps` / `docker_steps`. + +The pip path bootstraps a pinned `uv` from the astral-sh GitHub release if one isn't already on PATH (see "Bumping uv" below), then runs `uv tool install` to put `dataops-testgen` in a managed venv. After install, the app is auto-started via `start_testgen_app` (foreground until Ctrl+C); `tg start` brings it up again later. TestGen reads its config from `~/.testgen/config.env` — port, SSL, and `TESTGEN_LOG_FILE_PATH` are all written there at standalone-setup time. + +### Data locations at runtime + +- Unix: installer writes the compose file, credentials file, and `demo-config.json` next to `dk-installer.py`; logs go to `./.dk-installer/-.zip`. +- Windows: data and logs go to `%LOCALAPPDATA%/DataKitchenApps/`. + +### Demo container + +The `demo/` tree is built into a separate image (`datakitchen/data-observability-demo:latest`) via `demo/deploy/build-image`. `DemoContainerAction` in `dk-installer.py` pulls this image and mounts `demo-config.json` into it. Changes to `demo/*.py` don't affect the installer until that image is rebuilt and pushed. + +### Bumping uv + +The pip install path bootstraps a known version of `uv` from the astral-sh GitHub release. Two top-level constants govern this: + +- `UV_VERSION` — the pinned version (e.g., `"0.11.7"`). +- `UV_ASSETS` — a `(platform.system(), platform.machine()) → (asset_name, sha256)` map. Six entries: Linux x86_64/aarch64, Darwin x86_64/arm64, Windows AMD64/ARM64. + +To bump: + +1. Update `UV_VERSION`. +2. Pull the matching `dist-manifest.json` from `https://github.com/astral-sh/uv/releases/download//dist-manifest.json` and refresh the SHA256 for each of the 6 assets in `UV_ASSETS`. Each release also publishes a `.sha256` file you can `curl` directly if you'd rather pin one at a time. +3. Sanity-check: `pytest tests/test_uv_bootstrap.py`. The bootstrap step exercises hash verification and the asset-not-supported path. + +Do not skip the hash refresh — TLS verification is intentionally relaxed for the GitHub download (corp-proxy support), and the SHA256 pin is the security guarantee. + +## Testing conventions + +- `tests/installer.py` is a **symlink to `../dk-installer.py`** — tests import installer internals as `from tests.installer import ...`. Don't replace this with a copy. +- Heavy use of `unittest.mock.patch` to stub `subprocess` / `start_cmd` / `run_cmd`. The key fixtures live in `tests/conftest.py` — `action_cls` patches class-level attributes on `Action` so tests can instantiate actions without a real session folder, and `args_mock` provides a fully-populated `argparse.Namespace`. +- Tests are marked `@pytest.mark.unit` or `@pytest.mark.integration`. CI runs everything; use the markers locally to scope a run. + +## Style + +- Line length 120, double quotes, ruff-enforced (`pyproject.toml` restricts ruff's `include` to `dk-installer.py` only — the `demo/` and `tests/` trees are deliberately not linted by this project's ruff config). +- Pre-commit hooks run ruff on commit (`.pre-commit-config.yaml`). Install once with `pre-commit install`. +- Target Python is 3.9 (CI uses 3.9); avoid 3.10+ syntax like `match` statements or `X | Y` type unions in new code — the file uses `typing.Union` / `typing.Optional` deliberately for this reason. + +## CI + +`.github/workflows/pull_request.yml` runs ruff + pytest (with coverage comment) on every PR against `main`. `release_exe.yml` publishes the Windows `.exe` on every push to `main` by force-moving the `latest` tag and recreating the release — keep this in mind before merging, since each merge replaces the public download. diff --git a/README.md b/README.md index 5abff36..ae50f41 100644 --- a/README.md +++ b/README.md @@ -47,12 +47,21 @@ And it allows you to make fast, safe development changes. ### Install the required software -#### Requirements for TestGen & Observability +TestGen can be installed two ways: -| Software | Tested Versions | Command to check version | -|-------------------------|-----------------------------------------|-------------------------------| -| [Python](https://www.python.org/downloads/)
- Most Linux and macOS systems have Python pre-installed.
- On Windows machines, you will need to download and install it. | 3.9, 3.10, 3.11, 3.12, 3.13 | `python3 --version` | -| [Docker](https://docs.docker.com/get-docker/)
[Docker Compose](https://docs.docker.com/compose/install/) | 26.1, 27.5, 28.5
2.38, 2.39, 2.40 | `docker -v`
`docker compose version` | +- **Docker** (recommended when Docker is available) — deploys TestGen as a Docker Compose application. The most stable experience for persistent use; suited for team eval on a shared VM. +- **pip** — no Docker required. The installer downloads `uv`, installs Python 3.13 if needed, and installs TestGen in an isolated environment with an embedded Postgres database. Recommended for evaluation on machines where Docker isn't available. + +`tg install` with no flag prompts you to pick. If Docker isn't fully available it lists which prerequisites failed and recommends pip instead. Pass `--docker` or `--pip` to skip the prompt. + +Observability is always installed via Docker Compose. + +| Install mode | Required software | +|---|---| +| TestGen (pip) | [Python](https://www.python.org/downloads/) 3.9+ (only needed to run the installer itself — TestGen will use Python 3.13 via `uv`). | +| TestGen (Docker) + Observability | [Python](https://www.python.org/downloads/) 3.9+, [Docker](https://docs.docker.com/get-docker/) 27+, [Docker Compose](https://docs.docker.com/compose/install/) 5.0+. | + +Check versions with `python3 --version`, `docker -v`, `docker compose version`. ### Download the installer @@ -77,16 +86,18 @@ Before going through the quickstart, complete the prequisites above and then the ### Install the TestGen application -The installation downloads the latest Docker images for TestGen and deploys a new Docker Compose application. The process may take 5~10 minutes depending on your machine and network connection. - ```shell python3 dk-installer.py tg install ``` -The `--port` option may be used to set a custom localhost port for the application (default: 8501). -To enable SSL for HTTPS support, use the `--ssl-cert-file` and `--ssl-key-file` options to specify local file paths to your SSL certificate and key files. +With no flag, the installer probes Docker, shows which prerequisites are met, and prompts you to pick Docker or pip. Pass `--docker` or `--pip` to skip the prompt. -Once the installation completes, verify that you can login to the UI with the URL and credentials provided in the output. +* **pip mode** — downloads `uv` (if not already on your PATH), uses it to install Python 3.13 (if needed) and TestGen in an isolated environment. Typically takes 4-8 minutes. +* **Docker mode** — deploys TestGen as a Docker Compose application. Typically takes 5-10 minutes. + +On completion, the installer writes credentials to `dk-tg-credentials.txt`, generates demo data, and opens the TestGen UI in your default browser. Use `--no-demo` to skip demo generation. `--port` sets the UI port (default 8501); `--api-port` sets the API/MCP port (default 8530); `--ssl-cert-file` / `--ssl-key-file` enable HTTPS. + +Either install mode can later be upgraded with `python3 dk-installer.py tg upgrade` and restarted with `python3 dk-installer.py tg start` — the installer detects which flavor is present and routes accordingly. ### Install the Observability application @@ -136,15 +147,19 @@ Leave this process running, and continue with the [quickstart guide](https://doc ## Useful Commands -### DataOps TestGen +### DataOps TestGen (pip install) + +Start the app: `python3 dk-installer.py tg start` (reachable at `http://localhost:8501`, blocks until Ctrl+C) -The [docker compose CLI](https://docs.docker.com/compose/reference/) can be used to operate the installed TestGen application. All commands must be run in the same folder that contains the `docker-compose.yaml` file generated by the installation. +Stop the app: `Ctrl+C` in the terminal running `tg start` + +Upgrade the app to latest version: `python3 dk-installer.py tg upgrade` -Access the _testgen_ CLI: `docker compose exec engine bash` (use `exit` to return to the regular terminal) +### DataOps TestGen (Docker install) -Stop the app: `docker compose down` +Start the app: `python3 dk-installer.py tg start` (or `docker compose up` from the install folder) -Restart the app: `docker compose up` +Stop the app: `docker compose down` from the install folder containing `docker-compose.yaml` Upgrade the app to latest version: `python3 dk-installer.py tg upgrade` diff --git a/dk-installer.py b/dk-installer.py index ce14beb..9200142 100755 --- a/dk-installer.py +++ b/dk-installer.py @@ -19,14 +19,19 @@ import random import re import secrets +import shutil +import socket import ssl +import stat import string import subprocess import sys +import tarfile import textwrap import time import urllib.request import urllib.parse +import webbrowser import zipfile import typing @@ -46,8 +51,9 @@ BASE_API_URL_TPL = "http://host.docker.internal:{}/api" CREDENTIALS_FILE = "dk-{}-credentials.txt" -TESTGEN_LATEST_TAG = "v5" -TESTGEN_DEFAULT_IMAGE = f"datakitchen/dataops-testgen:{TESTGEN_LATEST_TAG}" +TESTGEN_MAJOR_VERSION = "5" +TESTGEN_PYTHON_VERSION = "3.13" +TESTGEN_DEFAULT_IMAGE = f"datakitchen/dataops-testgen:v{TESTGEN_MAJOR_VERSION}" TESTGEN_PULL_TIMEOUT = 5 TESTGEN_PULL_RETRIES = 3 TESTGEN_DEFAULT_PORT = 8501 @@ -55,6 +61,50 @@ TESTGEN_LATEST_VERSIONS_URL = ( "https://dk-support-external.s3.us-east-1.amazonaws.com/testgen-observability/testgen-latest-versions.json" ) +TESTGEN_PIP_PACKAGE = "dataops-testgen" +TESTGEN_COMPOSE_FILE = "docker-compose.yml" +TESTGEN_LOG_FILE_PATH = pathlib.Path.home() / ".testgen" / "logs" / "app.log" +TESTGEN_CONFIG_ENV_PATH = pathlib.Path.home() / ".testgen" / "config.env" +TESTGEN_APP_READY_TIMEOUT = 120 +INSTALL_MARKER_FILE = "dk-{}-install.json" +INSTALL_MODE_DOCKER = "docker" +INSTALL_MODE_PIP = "pip" + +UV_VERSION = "0.11.7" +UV_RELEASE_URL_TPL = "https://github.com/astral-sh/uv/releases/download/{version}/{asset}" +UV_DOWNLOAD_TIMEOUT = 120 +UV_DOWNLOAD_RETRIES = 3 +UV_BIN_SUBDIR = "bin" +# To bump UV_VERSION, refresh the SHA256s here from the dist-manifest: +# https://github.com/astral-sh/uv/releases/download//dist-manifest.json +# See "Bumping uv" in CLAUDE.md. +UV_ASSETS: dict[tuple[str, str], tuple[str, str]] = { + # (platform.system(), platform.machine()) -> (asset_name, sha256) + ("Linux", "x86_64"): ( + "uv-x86_64-unknown-linux-gnu.tar.gz", + "6681d691eb7f9c00ac6a3af54252f7ab29ae72f0c8f95bdc7f9d1401c23ea868", + ), + ("Linux", "aarch64"): ( + "uv-aarch64-unknown-linux-gnu.tar.gz", + "f2ee1cde9aabb4c6e43bd3f341dadaf42189a54e001e521346dc31547310e284", + ), + ("Darwin", "x86_64"): ( + "uv-x86_64-apple-darwin.tar.gz", + "0a4bc8fcde4974ea3560be21772aeecab600a6f43fa6e58169f9fa7b3b71d302", + ), + ("Darwin", "arm64"): ( + "uv-aarch64-apple-darwin.tar.gz", + "66e37d91f839e12481d7b932a1eccbfe732560f42c1cfb89faddfa2454534ba8", + ), + ("Windows", "AMD64"): ( + "uv-x86_64-pc-windows-msvc.zip", + "fe0c7815acf4fc45f8a5eff58ed3cf7ae2e15c3cf1dceadbd10c816ec1690cc1", + ), + ("Windows", "ARM64"): ( + "uv-aarch64-pc-windows-msvc.zip", + "1387e1c94e15196351196b79fce4c1e6f4b30f19cdaaf9ff85fbd6b046018aa2", + ), +} OBS_LATEST_TAG = "v2" OBS_DEF_BE_IMAGE = f"datakitchen/dataops-observability-be:{OBS_LATEST_TAG}" @@ -84,15 +134,25 @@ LOG = logging.getLogger() COMPOSE_VAR_RE = re.compile(r"\$\{(\w+):-([^\}]*)\}") +TESTGEN_PIP_VERSION_RE = re.compile(rf"^{re.escape(TESTGEN_PIP_PACKAGE)}\s+v(\S+)") # # Utility functions # -def _get_tg_base_url(args): +def get_tg_url(args, port): protocol = "https" if args.ssl_cert_file and args.ssl_key_file else "http" - return f"{protocol}://localhost:{args.port}" + return f"{protocol}://localhost:{port}" + + +def open_app_in_browser(url: str) -> None: + """Best-effort open the URL in the user's default browser. Silent no-op + on headless / browser-less environments.""" + try: + webbrowser.open(url) + except Exception: + LOG.exception("Failed to open browser for %s", url) def collect_images_digest(action, images, env=None): @@ -132,9 +192,40 @@ def generate_password(): return password -def delete_file(file_path): - LOG.debug("Deleting [%s]", file_path.name) - file_path.unlink(missing_ok=True) +def remove_path(path: pathlib.Path, label: typing.Optional[str] = None) -> bool: + """Remove a file or directory tree if it exists. + When ``label`` is provided, success/failure is also reported via CONSOLE. + + Returns True if something was actually removed. + """ + if not (path.exists() or path.is_symlink()): + return False + LOG.debug("Removing path [%s]", path) + try: + if path.is_dir(): + # On Windows, files inside a Postgres data dir are often marked read-only, + # which causes shutil.rmtree to abort partway through. Clear the read-only + # bit and retry from the error callback. shutil.rmtree's `onerror` is + # deprecated in 3.12 and removed in 3.14, replaced by `onexc`; the callback + # signatures differ on the third arg but we ignore it either way. + def _retry(func, p, _exc): + os.chmod(p, stat.S_IWRITE) + func(p) + + if sys.version_info >= (3, 12): + shutil.rmtree(path, onexc=_retry) + else: + shutil.rmtree(path, onerror=_retry) + else: + path.unlink() + except OSError: + LOG.exception("Failed to remove %s", path) + if label: + CONSOLE.msg(f"Could not remove {label} ({path}); remove manually if needed.") + return False + if label: + CONSOLE.msg(f"Removed {label} ({path})") + return True @functools.cache @@ -145,13 +236,102 @@ def get_installer_version(): return "N/A" +def resolve_windows_redirected_path(path: pathlib.Path) -> pathlib.Path: + """If running under Microsoft Store Python, rewrite ``path`` from its + UWP-virtualized form (what Python sees via ``os.environ['LOCALAPPDATA']``) + to the real on-disk path users can navigate to in Explorer/PowerShell. + Returns ``path`` unchanged on non-Windows or non-Store Python. + """ + exe = pathlib.Path(sys.executable) + if "WindowsApps" not in exe.parts: + return path + parts = exe.parent.name.split("_") + if len(parts) < 2: + return path + pfn = f"{parts[0]}_{parts[-1]}" + try: + local_appdata = pathlib.Path(os.environ["LOCALAPPDATA"]) + rel = path.relative_to(local_appdata) + except (KeyError, ValueError): + return path + return local_appdata / "Packages" / pfn / "LocalCache" / "Local" / rel + + def simplify_path(path: pathlib.Path) -> pathlib.Path: + if platform.system() == "Windows": + path = resolve_windows_redirected_path(path) try: return path.relative_to(pathlib.Path().absolute()) except ValueError: return path +def command_hint(prod: str, subcmd: str, menu_label: str) -> str: + """Render a user-facing CLI hint. Under the frozen Windows .exe, the user + typically has no Python and runs via the menu, so we point them at the + menu label instead of a command they can't type. + """ + if getattr(sys, "frozen", False): + return f"select '{menu_label}' from the menu" + return f"run `python3 {INSTALLER_NAME} {prod} {subcmd}`" + + +class InstallMarker: + """Read/write the TestGen install marker file. Falls back to detecting + a legacy Docker install (compose file + credentials) from before the + marker was introduced. + """ + + def __init__(self, data_folder: pathlib.Path, prod: str, compose_file_name: typing.Optional[str] = None): + self._data_folder = data_folder + self._prod = prod + self._compose_file_name = compose_file_name + self.path = data_folder / INSTALL_MARKER_FILE.format(prod) + + def read(self) -> typing.Optional[str]: + if self.path.exists(): + try: + data = json.loads(self.path.read_text()) + except Exception: + LOG.exception("Failed to read install marker at %s", self.path) + else: + install_mode = data.get("install_mode") + if install_mode in (INSTALL_MODE_DOCKER, INSTALL_MODE_PIP): + return install_mode + LOG.warning("Install marker has unexpected install_mode: %r", install_mode) + if ( + self._compose_file_name + and (self._data_folder / self._compose_file_name).exists() + and (self._data_folder / CREDENTIALS_FILE.format(self._prod)).exists() + ): + LOG.info("No marker present; detected legacy Docker install in %s", self._data_folder) + return INSTALL_MODE_DOCKER + return None + + def write(self, mode: str, **extra) -> None: + if mode not in (INSTALL_MODE_DOCKER, INSTALL_MODE_PIP): + raise ValueError(f"Unknown install_mode: {mode}") + now = datetime.datetime.now(datetime.timezone.utc).isoformat() + created_on = now + if self.path.exists(): + try: + existing = json.loads(self.path.read_text()) + if isinstance(existing.get("created_on"), str): + created_on = existing["created_on"] + except Exception: + LOG.exception("Failed to read existing install marker at %s", self.path) + self.path.write_text( + json.dumps( + {"install_mode": mode, "created_on": created_on, "last_updated_on": now, **extra}, + indent=2, + ) + ) + + def unlink(self) -> None: + if self.path.exists(): + self.path.unlink() + + @contextlib.contextmanager def stream_iterator(proc: subprocess.Popen, stream_name: str, file_path: pathlib.Path, timeout: float = 1.0): comm_index, exc_attr = { @@ -217,8 +397,8 @@ def __init__(self): def title(self, text): LOG.info("Console title: [%s]", text) - if not self._last_is_space: - print("") + # Always blank-line before a title so they are separated from any input() prompts + print("") print(f" == {text}") print("") self._last_is_space = True @@ -298,8 +478,9 @@ class Requirement: key: str cmd: tuple[typing.Union[str, pathlib.Path], ...] fail_msg: tuple[str, ...] + label: typing.Optional[str] = None - def check_availability(self, action, args): + def check_availability(self, action, args, quiet=False): try: action.run_cmd_retries( *(seg.format(**args.__dict__) for seg in self.cmd), @@ -307,9 +488,10 @@ def check_availability(self, action, args): retries=1, ) except CommandFailed: - CONSOLE.space() - for line in self.fail_msg: - CONSOLE.msg(line.format(**args.__dict__)) + if not quiet: + CONSOLE.space() + for line in self.fail_msg: + CONSOLE.msg(line.format(**args.__dict__)) return False else: return True @@ -551,10 +733,19 @@ def _get_failed_cmd_log_file_path( return None, None def _msg_unexpected_error(self, exception: Exception) -> None: - exception, log_path = self._get_failed_cmd_log_file_path(exception) - if exception and log_path: - CONSOLE.msg(f"Command '{exception.cmd}' failed with code {exception.ret_code}. See the output below.") + cmd_exception, log_path = self._get_failed_cmd_log_file_path(exception) + if cmd_exception and log_path: + CONSOLE.msg( + f"Command '{cmd_exception.cmd}' failed with code {cmd_exception.ret_code}. See the output below." + ) CONSOLE.print_log(log_path) + else: + root = exception + while root.__cause__ is not None: + root = root.__cause__ + if str(root).strip(): + CONSOLE.space() + CONSOLE.msg(f"Error: {root}") msg_file_path = simplify_path(self.session_zip) CONSOLE.space() @@ -565,13 +756,24 @@ def _msg_unexpected_error(self, exception: Exception) -> None: def get_requirements(self, args) -> list[Requirement]: return self.requirements - def _check_requirements(self, args): + def check_requirements(self, args): missing_reqs = [req.key for req in self.get_requirements(args) if not req.check_availability(self, args)] if missing_reqs: self.analytics.additional_properties["missing_requirements"] = missing_reqs raise AbortAction + # Names of instance attributes that hold per-invocation state. Reset + # before each run so the same Action instance can be re-invoked cleanly + # in menu mode (Windows .exe) without state from the previous run leaking + # into the next. Subclasses extend this tuple with their own attrs. + _per_invocation_attrs: tuple[str, ...] = ("_cmd_idx",) + + def _reset_per_invocation_state(self): + for attr in self._per_invocation_attrs: + self.__dict__.pop(attr, None) + def execute_with_log(self, args): + self._reset_per_invocation_state() with ( self.init_session_folder(prefix=f"{args.prod}-{self.args_cmd}"), self.configure_logging(debug=args.debug), @@ -596,7 +798,7 @@ def execute_with_log(self, args): LOG.info("Installer version: %s", get_installer_version()) try: - self._check_requirements(args) + self.check_requirements(args) self.execute(args) except AbortAction: @@ -609,8 +811,9 @@ def execute_with_log(self, args): self._msg_unexpected_error(e) raise InstallerError from e except KeyboardInterrupt as e: - CONSOLE.space() - CONSOLE.msg("Processing interrupted. This may result in an inconsistent platform state.") + # Reset the cursor to column 0 — the terminal echoed `^C` mid-line. + print("") + CONSOLE.msg("Processing interrupted. This may result in an inconsistent application state.") raise AbortAction from e def get_parser(self, sub_parsers): @@ -683,12 +886,16 @@ def run_cmd( return json_lines @contextlib.contextmanager - def start_cmd(self, *cmd, raise_on_non_zero=True, env=None, **popen_args): + def start_cmd(self, *cmd, raise_on_non_zero=True, env=None, redact=(), **popen_args): started = time.time() self._cmd_idx += 1 - cmd_str = " ".join(str(part) for part in cmd) - LOG.debug("Command [%04d]: [%s]", self._cmd_idx, cmd_str) + # Censor secrets before they reach logs + log_str = " ".join(str(part) for part in cmd) + for secret in redact: + if secret: + log_str = log_str.replace(str(secret), "***") + LOG.debug("Command [%04d]: [%s]", self._cmd_idx, log_str) if isinstance(env, dict): LOG.debug("Command [%04d] extra ENV: [%s]", self._cmd_idx, ", ".join(env.keys())) @@ -705,9 +912,9 @@ def start_cmd(self, *cmd, raise_on_non_zero=True, env=None, **popen_args): ) except FileNotFoundError as e: LOG.error("Command [%04d] failed to find the executable", self._cmd_idx) - raise CommandFailed(self._cmd_idx, cmd_str, None) from e + raise CommandFailed(self._cmd_idx, log_str, None) from e - slug_cmd = re.sub(r"[^a-zA-Z]+", "-", cmd_str)[:100].strip("-") + slug_cmd = re.sub(r"[^a-zA-Z]+", "-", log_str)[:100].strip("-") stdout_path, stderr_path = [ self.session_folder.joinpath(f"{self._cmd_idx:04d}-{stream_name}-{slug_cmd}.txt") @@ -728,7 +935,7 @@ def start_cmd(self, *cmd, raise_on_non_zero=True, env=None, **popen_args): # We capture and raise CommandFailed to allow the client code to raise an empty CommandFailed exception # but still get a contextualized exception at the end except CommandFailed as e: - raise CommandFailed(self._cmd_idx, cmd_str, proc.returncode) from e.__cause__ + raise CommandFailed(self._cmd_idx, log_str, proc.returncode) from e.__cause__ finally: elapsed = time.time() - started LOG.info( @@ -771,6 +978,10 @@ def __init__(self): super().__init__() self.ctx = {} + def _reset_per_invocation_state(self): + super()._reset_per_invocation_state() + self.ctx = {} + def _print_intro_text(self, args): CONSOLE.space() for line in self.intro_text: @@ -960,11 +1171,13 @@ def run(self, parent=None): "DOCKER", ("docker", "-v"), ("The prerequisite Docker is not available.", "Install Docker and try again."), + label="Docker installed", ) REQ_DOCKER_DAEMON = Requirement( "DOCKER_ENGINE", ("docker", "system", "events", "--since=0m", "--until=0m"), ("The Docker engine is not running.", "Start the Docker engine and try again."), + label="Docker engine running", ) REQ_TESTGEN_IMAGE = Requirement( "TESTGEN_IMAGE", @@ -973,9 +1186,151 @@ def run(self, parent=None): "The Docker engine could not access TestGen's image.", "Make sure your networking policy allows Docker to pull the {image} image.", ), + label="TestGen image reachable", ) +def get_uv_asset(prod: str) -> tuple[str, str]: + """Return (asset_name, sha256) for the current platform, or raise AbortAction.""" + key = (platform.system(), platform.machine()) + try: + return UV_ASSETS[key] + except KeyError: + supported = ", ".join(f"{s}/{m}" for s, m in UV_ASSETS) + CONSOLE.msg(f"No prebuilt uv binary available for platform {key[0]}/{key[1]}.") + CONSOLE.msg(f"Supported: {supported}.") + CONSOLE.msg( + "Install uv manually (https://docs.astral.sh/uv/getting-started/installation/) and re-run, " + f"or {command_hint(prod, 'install --docker', 'Install TestGen')} to use Docker." + ) + raise AbortAction + + +def resolve_uv_path(data_folder: pathlib.Path) -> typing.Optional[str]: + """Return the path to a usable ``uv`` binary, preferring ``PATH`` then the + installer-local download from a prior bootstrap. Returns ``None`` if neither + is available — callers decide whether that's fatal. + """ + if uv_on_path := shutil.which("uv"): + return uv_on_path + bin_name = "uv.exe" if platform.system() == "Windows" else "uv" + local_uv = data_folder / UV_BIN_SUBDIR / bin_name + if local_uv.exists(): + return str(local_uv) + return None + + +class UvBootstrapStep(Step): + label = "Preparing the Python environment" + + def pre_execute(self, action, args): + # Resolve uv eagerly so later steps' pre_execute hooks can use it. + # If uv has to be downloaded, ctx stays unset until execute runs + # and this step's download path populates it. + if uv_path := resolve_uv_path(action.data_folder): + action.ctx["uv_path"] = uv_path + + def execute(self, action, args): + if uv_path := action.ctx.get("uv_path"): + LOG.info("Using existing uv at %s", uv_path) + action.analytics.additional_properties["uv_source"] = "existing" + self._capture_uv_version(action, uv_path) + return + + asset_name, expected_sha256 = get_uv_asset(args.prod) + url = UV_RELEASE_URL_TPL.format(version=UV_VERSION, asset=asset_name) + action.analytics.additional_properties["uv_source"] = "download" + + bin_name = "uv.exe" if platform.system() == "Windows" else "uv" + target_path = action.data_folder / UV_BIN_SUBDIR / bin_name + + last_exc = None + for attempt in range(1, UV_DOWNLOAD_RETRIES + 1): + try: + self._download_and_install(action, url, asset_name, expected_sha256) + except InstallerError: + # Deterministic failure (SHA256 mismatch, malformed archive, + # unknown asset format) — retrying won't help, and on a + # SHA256 mismatch, repeated attempts waste minutes only to + # report the same MITM-or-corrupted-release error. + if target_path.exists(): + target_path.unlink() + raise + except Exception as e: + LOG.warning("uv bootstrap attempt %d/%d failed: %s", attempt, UV_DOWNLOAD_RETRIES, e) + last_exc = e + if target_path.exists(): + target_path.unlink() + # Exponential backoff between attempts to ride out transient + # network blips and registry rate-limits. + if attempt < UV_DOWNLOAD_RETRIES: + time.sleep(2**attempt) + else: + action.ctx["uv_path"] = str(target_path) + self._capture_uv_version(action, str(target_path)) + return + raise InstallerError(f"Failed to bootstrap uv after {UV_DOWNLOAD_RETRIES} attempts: {last_exc}") from last_exc + + def _download_and_install(self, action, url: str, asset_name: str, expected_sha256: str) -> None: + bin_dir = action.data_folder / UV_BIN_SUBDIR + bin_dir.mkdir(parents=True, exist_ok=True) + archive_path = bin_dir / asset_name + + try: + LOG.info("Downloading uv %s from %s", UV_VERSION, url) + ssl_context = ssl.create_default_context() + ssl_context.check_hostname = False + ssl_context.verify_mode = ssl.CERT_NONE + with urllib.request.urlopen(url, timeout=UV_DOWNLOAD_TIMEOUT, context=ssl_context) as resp: + archive_path.write_bytes(resp.read()) + + actual_sha256 = hashlib.sha256(archive_path.read_bytes()).hexdigest() + if actual_sha256 != expected_sha256: + raise InstallerError( + f"SHA256 mismatch for {asset_name}: expected {expected_sha256}, got {actual_sha256}" + ) + + self._extract_uv_binary(archive_path, asset_name, bin_dir) + + if platform.system() != "Windows": + (bin_dir / "uv").chmod(0o755) + finally: + if archive_path.exists(): + archive_path.unlink() + + @staticmethod + def _extract_uv_binary(archive_path: pathlib.Path, asset_name: str, bin_dir: pathlib.Path) -> None: + if asset_name.endswith(".tar.gz"): + with tarfile.open(archive_path, "r:gz") as tf: + for member in tf.getmembers(): + if member.isfile() and pathlib.PurePosixPath(member.name).name == "uv": + src = tf.extractfile(member) + if src is None: + continue + (bin_dir / "uv").write_bytes(src.read()) + return + raise InstallerError(f"Could not find 'uv' binary in archive {asset_name}") + if asset_name.endswith(".zip"): + with zipfile.ZipFile(archive_path) as zf: + for name in zf.namelist(): + if pathlib.PurePosixPath(name).name == "uv.exe": + (bin_dir / "uv.exe").write_bytes(zf.read(name)) + return + raise InstallerError(f"Could not find 'uv.exe' in archive {asset_name}") + raise InstallerError(f"Unexpected asset format: {asset_name}") + + @staticmethod + def _capture_uv_version(action, uv_path: str) -> None: + try: + output = action.run_cmd(uv_path, "--version", capture_text=True, raise_on_non_zero=False) + except Exception: + LOG.exception("Failed to capture uv version") + return + # `uv --version` prints e.g. "uv 0.11.7 (abcd1234 2024-09-15)". + if output and (match := re.match(r"uv\s+(\S+)", output.strip())): + action.analytics.additional_properties["uv_version"] = match.group(1) + + class AnalyticsMultiStepAction(MultiStepAction): ANALYTICS_DISCLAIMER = [ "DataKitchen has enabled anonymous aggregate user behavior analytics.", @@ -1008,22 +1363,7 @@ def get_volumes(self, args) -> list[dict[str, str]]: volumes = self.run_cmd("docker", "volume", "list", "--format=json", capture_json_lines=True) return [v for v in volumes if label in v.get("Labels", "")] - -class ComposeDeleteAction(Action, ComposeActionMixin): - args_cmd = "delete" - requirements = [REQ_DOCKER, REQ_DOCKER_DAEMON] - - def execute(self, args): - if self.get_compose_file_path(args).exists(): - self._delete_containers(args) - self._delete_network() - else: - # Trying to delete the network before any exception - self._delete_network() - # Trying to delete dangling volumes - self._delete_volumes(args) - - def _delete_containers(self, args): + def delete_compose_containers(self, args): CONSOLE.title(f"Delete {args.prod_name} instance") try: self.run_cmd( @@ -1042,11 +1382,11 @@ def _delete_containers(self, args): raise AbortAction else: if not args.keep_config: - delete_file(self.get_compose_file_path(args)) - delete_file(self.data_folder / CREDENTIALS_FILE.format(args.prod)) + remove_path(self.get_compose_file_path(args)) + remove_path(self.data_folder / CREDENTIALS_FILE.format(args.prod)) CONSOLE.msg("Docker containers and volumes deleted") - def _delete_network(self): + def delete_compose_network(self): try: self.run_cmd("docker", "network", "rm", DOCKER_NETWORK, raise_on_non_zero=True) except CommandFailed: @@ -1054,7 +1394,7 @@ def _delete_network(self): else: CONSOLE.msg("Docker network deleted") - def _delete_volumes(self, args): + def delete_compose_volumes(self, args): if volumes := self.get_volumes(args): try: self.run_cmd( @@ -1069,6 +1409,21 @@ def _delete_volumes(self, args): else: CONSOLE.msg("Docker volumes deleted") + +class ComposeDeleteAction(Action, ComposeActionMixin): + args_cmd = "delete" + requirements = [REQ_DOCKER, REQ_DOCKER_DAEMON] + + def execute(self, args): + if self.get_compose_file_path(args).exists(): + self.delete_compose_containers(args) + self.delete_compose_network() + else: + # Trying to delete the network before any exception + self.delete_compose_network() + # Trying to delete dangling volumes + self.delete_compose_volumes(args) + def get_parser(self, sub_parsers): parser = super().get_parser(sub_parsers) parser.add_argument( @@ -1095,7 +1450,7 @@ def pre_execute(self, action, args): f"Found {args.prod_name} docker compose containers and/or volumes. If a previous attempt to run this", ) CONSOLE.msg( - f"installer failed, please run `python3 {INSTALLER_NAME} {args.prod} delete` before trying again." + f"installer failed, {command_hint(args.prod, 'delete', f'Uninstall {args.prod_name}')} before trying again." ) CONSOLE.space() if volumes: @@ -1261,7 +1616,7 @@ def on_action_success(self, action, args): def on_action_fail(self, action, args): # We keep the file around for inspection when in debug mode if not args.debug and not action.ctx.get("using_existing"): - delete_file(action.get_compose_file_path(args)) + remove_path(action.get_compose_file_path(args)) def get_observability_version(action, args): @@ -1375,7 +1730,7 @@ def on_action_success(self, action, args): console_tee(f"Username: {self._user_data['username']}") console_tee(f"Password: {self._user_data['password']}", skip_logging=True) - CONSOLE.msg(f"(Credentials also written to {cred_file_path.name} file)") + CONSOLE.msg(f"(Credentials also written to {simplify_path(cred_file_path)})") class ObsGenerateDemoConfigStep(Step): @@ -1635,7 +1990,9 @@ def execute(self, args): except Exception: CONSOLE.title("Demo FAILED") CONSOLE.space() - CONSOLE.msg(f"To retry the demo, first run `python3 {INSTALLER_NAME} {args.prod} delete-demo`") + CONSOLE.msg( + f"To retry the demo, first {command_hint(args.prod, 'delete-demo', f'Delete {args.prod_name} demo data')}." + ) else: CONSOLE.title("Demo SUCCEEDED") @@ -1657,7 +2014,9 @@ def execute(self, args): try: self.run_dk_demo_container("obs-heartbeat-demo") except KeyboardInterrupt: - CONSOLE.msg("Observability Heartbeat demo stopped") + # Reset the cursor to column 0 — the terminal echoed `^C` mid-line. + print("") + CONSOLE.msg("Observability Heartbeat demo stopped.") class UpdateComposeFileStep(Step): @@ -1668,6 +2027,7 @@ def __init__(self): self.update_analytics = False self.update_token = False self.update_base_url = False + self.update_api_port = False super().__init__() def pre_execute(self, action, args): @@ -1680,17 +2040,7 @@ def pre_execute(self, action, args): self.update_version = True else: try: - output = action.run_cmd( - "docker", - "compose", - "-f", - action.get_compose_file_path(args), - "exec", - "engine", - "testgen", - "--help", - capture_text=True, - ) + output = run_testgen_cli(action, args, "--help", capture_text=True) version_match = re.search(r"TestGen\s(?:[a-zA-Z]+\s)*([0-9.]*)", output) current_version = version_match.group(1) @@ -1765,7 +2115,7 @@ def execute(self, action, args): contents = action.get_compose_file_path(args).read_text() if self.update_version: - contents = re.sub(r"(image:\s*datakitchen.+:).+\n", rf"\1{TESTGEN_LATEST_TAG}\n", contents) + contents = re.sub(r"(image:\s*datakitchen.+:).+\n", rf"\1v{TESTGEN_MAJOR_VERSION}\n", contents) if self.update_analytics: if args.send_analytics_data: @@ -1834,13 +2184,13 @@ def on_action_success(self, action, args): super().on_action_success(action, args) cred_file_path = action.data_folder.joinpath(CREDENTIALS_FILE.format(args.prod)) with CONSOLE.tee(cred_file_path) as console_tee: - console_tee(f"User Interface: {_get_tg_base_url(args)}") - console_tee("CLI Access: docker compose exec engine bash") + console_tee(f"User Interface: {get_tg_url(args, args.port)}") + console_tee(f"API & MCP: {get_tg_url(args, args.api_port)}") console_tee("") console_tee(f"Username: {self.username}") - console_tee(f"Password: {self.password}") + console_tee(f"Password: {self.password}", skip_logging=True) - CONSOLE.msg(f"(Credentials also written to {cred_file_path.name} file)") + CONSOLE.msg(f"(Credentials also written to {simplify_path(cred_file_path)})") def get_credentials_from_compose_file(self, file_contents): username = None @@ -1893,7 +2243,7 @@ def get_compose_file_contents(self, action, args): TG_EXPORT_TO_OBSERVABILITY_VERIFY_SSL: no TG_INSTANCE_ID: {action.analytics.get_instance_id()} TG_ANALYTICS: {"yes" if args.send_analytics_data else "no"} - TG_UI_BASE_URL: {_get_tg_base_url(args)} + TG_UI_BASE_URL: {get_tg_url(args, args.port)} {ssl_variables} services: @@ -1968,24 +2318,14 @@ def execute(self, action, args): class TestGenSetupDatabaseStep(Step): - label = "Initializing the platform database" + label = "Initializing the application database" def execute(self, action, args): - action.run_cmd( - "docker", - "compose", - "-f", - action.get_compose_file_path(args), - "exec", - "engine", - "testgen", - "setup-system-db", - "--yes", - ) + run_testgen_cli(action, args, "setup-system-db", "--yes") class TestGenUpgradeDatabaseStep(Step): - label = "Upgrading the platform database" + label = "Upgrading the application database" def pre_execute(self, action, args): self.required = action.args_cmd == "upgrade" @@ -1993,37 +2333,324 @@ def pre_execute(self, action, args): def execute(self, action, args): if action.args_cmd == "install" and action.ctx.get("using_existing"): raise SkipStep - else: - action.run_cmd( - "docker", - "compose", - "-f", - action.get_compose_file_path(args), - "exec", - "engine", - "testgen", - "upgrade-system-version", - ) + run_testgen_cli(action, args, "upgrade-system-version") def on_action_success(self, action, args): - output = action.run_cmd( - "docker", - "compose", - "-f", - action.get_compose_file_path(args), - "exec", - "engine", - "testgen", - "--help", - capture_text=True, - ) + output = run_testgen_cli(action, args, "--help", capture_text=True) match = re.search(r"TestGen\s(?:[a-zA-Z]+\s)*([0-9.]*)", output) CONSOLE.msg(f"Application version: {match.group(1)}") +class UvToolInstallStep(Step): + label = "Installing TestGen" + + def execute(self, action, args): + uv_path = action.ctx["uv_path"] + major = int(TESTGEN_MAJOR_VERSION) + constraint = f"{TESTGEN_PIP_PACKAGE}[standalone]>={major},<{major + 1}" + action.run_cmd( + uv_path, + "tool", + "install", + "--force", + "--python", + TESTGEN_PYTHON_VERSION, + constraint, + ) + # Add uv's tool bin dir to the user's shell rc so future shells pick + # up the testgen entry point. Non-fatal if the shell isn't recognized. + action.run_cmd(uv_path, "tool", "update-shell", raise_on_non_zero=False) + if version := read_installed_testgen_version(action): + action.analytics.additional_properties["testgen_version"] = version + + +def read_installed_testgen_version(action) -> typing.Optional[str]: + """Parse ``uv tool list`` for the installed dataops-testgen version, or + ``None`` if uv is unavailable or the tool isn't listed. + """ + uv_path = action.ctx.get("uv_path") + if uv_path is None: + return None + try: + output = action.run_cmd(uv_path, "tool", "list", capture_text=True, raise_on_non_zero=False) + except Exception: + LOG.exception("Failed to read uv tool list") + return None + for line in (output or "").splitlines(): + match = TESTGEN_PIP_VERSION_RE.match(line.strip()) + if match: + return match.group(1) + return None + + +def run_testgen_cli(action, args, *cmd_args, **run_cmd_kwargs): + """Run a ``testgen`` CLI subcommand in the appropriate mode based on + ``action._resolved_mode``: pip mode invokes the testgen script directly; + Docker mode runs it inside the engine container via ``docker compose exec``. + Extra keyword arguments are forwarded to ``action.run_cmd`` (e.g. + ``capture_text=True``) and the return value of ``run_cmd`` is returned. + """ + if action._resolved_mode == INSTALL_MODE_PIP: + testgen_path = resolve_testgen_path(action, args) + return action.run_cmd(testgen_path, *cmd_args, **run_cmd_kwargs) + return action.run_cmd( + "docker", + "compose", + "-f", + action.get_compose_file_path(args), + "exec", + "engine", + "testgen", + *cmd_args, + **run_cmd_kwargs, + ) + + +def resolve_testgen_path(action, args) -> str: + """Return the absolute path to the ``testgen`` script that ``uv tool install`` + placed in uv's bin dir. Calling this script directly (instead of via + ``uv tool run --from ...``) executes inside the persistent tool venv — + necessary for steps like ``standalone-setup`` whose side effects + (e.g., Streamlit patching) must apply to the venv ``run-app`` later uses. + """ + ctx = getattr(action, "ctx", None) or {} + uv_path = ctx.get("uv_path") or resolve_uv_path(action.data_folder) + if uv_path is None: + CONSOLE.msg(f"uv not found. To install TestGen, {command_hint(args.prod, 'install', 'Install TestGen')}.") + raise AbortAction + bin_dir = action.run_cmd(uv_path, "tool", "dir", "--bin", capture_text=True) + if not bin_dir: + raise InstallerError("Could not determine uv tool bin directory.") + bin_name = "testgen.exe" if platform.system() == "Windows" else "testgen" + testgen_path = pathlib.Path(bin_dir.strip()) / bin_name + if not testgen_path.exists(): + CONSOLE.msg(f"testgen script not found at {testgen_path}.") + CONSOLE.msg( + f"Try {command_hint(args.prod, 'delete', 'Uninstall TestGen')}, " + f"then {command_hint(args.prod, 'install', 'Install TestGen')}." + ) + raise AbortAction + return str(testgen_path) + + +def wait_for_tcp_port(port: int, timeout: int) -> bool: + """Poll ``localhost:port`` until a TCP connection succeeds or timeout elapses.""" + deadline = time.time() + timeout + while time.time() < deadline: + try: + with socket.create_connection(("localhost", port), timeout=1): + return True + except OSError: + time.sleep(0.5) + return False + + +def read_testgen_config_env() -> dict[str, str]: + """Parse ``~/.testgen/config.env`` (key=value lines). The source of truth + for the port + SSL settings TestGen uses, since ``standalone-setup`` + persists them there at install time. + """ + config: dict[str, str] = {} + if not TESTGEN_CONFIG_ENV_PATH.exists(): + return config + for line in TESTGEN_CONFIG_ENV_PATH.read_text().splitlines(): + line = line.strip() + if not line or line.startswith("#"): + continue + key, sep, value = line.partition("=") + if sep: + config[key.strip()] = value.strip().strip('"').strip("'") + return config + + +def start_testgen_app(action, args) -> None: + """Start ``testgen run-app`` and block until the user interrupts. + + stdout/stderr are discarded — TestGen writes its own logs to + ``TESTGEN_LOG_FILE_PATH`` (configured at standalone-setup time) and the + App Logs dialog in the UI surfaces them. Capturing the subprocess streams + here would just duplicate that and bloat the support zip. + + Using ``Popen`` directly with ``DEVNULL`` rather than going through + ``start_cmd`` — for an indefinite-running process, ``start_cmd``'s pipe- + based capture would deadlock once the OS pipe buffer fills. + """ + testgen_path = resolve_testgen_path(action, args) + # Resolve port + SSL state from the standalone-setup-persisted config so + # the URL we display matches what TestGen actually binds to (and so this + # works for ``tg start`` where args has no port flags registered). + config = read_testgen_config_env() + port = int(config.get("TG_UI_PORT") or TESTGEN_DEFAULT_PORT) + has_ssl = bool(config.get("SSL_CERT_FILE") and config.get("SSL_KEY_FILE")) + url = f"{'https' if has_ssl else 'http'}://localhost:{port}" + + LOG.debug("Starting TestGen: %s run-app", testgen_path) + + CONSOLE.space() + CONSOLE.msg("Starting TestGen...") + + try: + proc = subprocess.Popen( + [testgen_path, "run-app"], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + stdin=subprocess.DEVNULL, + ) + except FileNotFoundError as e: + raise InstallerError(f"Could not start TestGen: {e}") from e + + try: + if not wait_for_tcp_port(port, timeout=TESTGEN_APP_READY_TIMEOUT): + proc.terminate() + with contextlib.suppress(subprocess.TimeoutExpired): + proc.wait(timeout=5) + raise InstallerError( + f"TestGen did not start within {TESTGEN_APP_READY_TIMEOUT} seconds. " + f"See {simplify_path(TESTGEN_LOG_FILE_PATH)} for details." + ) + + CONSOLE.msg(f"TestGen is running at {url}.") + # During ``tg install`` the credentials file already lists the log path; + # only repeat it for ``tg start`` invocations where there's no creds file print. + if action.args_cmd == "start": + CONSOLE.msg(f"Logs: {simplify_path(TESTGEN_LOG_FILE_PATH)}") + CONSOLE.space() + CONSOLE.msg("Press Ctrl+C to stop the app.") + CONSOLE.msg(f"To start it again later, {command_hint(args.prod, 'start', 'Start TestGen')}.") + CONSOLE.space() + + try: + proc.wait() + except KeyboardInterrupt: + # Reset the cursor to column 0 — the terminal echoed `^C` mid-line. + print("") + CONSOLE.msg("Stopping TestGen...") + proc.terminate() + try: + proc.wait(timeout=10) + except subprocess.TimeoutExpired: + proc.kill() + proc.wait() + CONSOLE.msg("TestGen stopped.") + CONSOLE.msg(f"To start it again, {command_hint(args.prod, 'start', 'Start TestGen')}.") + finally: + if proc.poll() is None: + proc.terminate() + with contextlib.suppress(subprocess.TimeoutExpired): + proc.wait(timeout=5) + + +class UvToolUpgradeStep(Step): + label = "Upgrading TestGen" + + def __init__(self): + self.current_version = None + + def pre_execute(self, action, args): + self.current_version = read_installed_testgen_version(action) + if self.current_version: + CONSOLE.msg(f"Current version: v{self.current_version}") + + def execute(self, action, args): + uv_path = action.ctx["uv_path"] + # ``--no-cache`` (top-level option) bypasses uv's cached PyPI index for + # this one invocation so a release that was just published is picked up + # immediately instead of being served from a stale cache. + action.run_cmd( + uv_path, + "--no-cache", + "tool", + "upgrade", + TESTGEN_PIP_PACKAGE, + ) + + def on_action_success(self, action, args): + new_version = read_installed_testgen_version(action) + if new_version is None: + return + action.analytics.additional_properties["testgen_version"] = new_version + if new_version == self.current_version: + CONSOLE.msg(f"Application is already up-to-date (v{new_version}).") + else: + CONSOLE.msg(f"Updated to v{new_version}.") + + +class TestgenStandaloneSetupStep(Step): + label = "Initializing TestGen" + + def __init__(self): + self.username = None + self.password = None + + def pre_execute(self, action, args): + self.username = DEFAULT_USER_DATA["username"] + self.password = generate_password() + + def execute(self, action, args): + # standalone-setup persists these env vars to ~/.testgen/config.env so + # subsequent ``testgen run-app`` invocations pick them up automatically. + # TESTGEN_LOG_FILE_PATH lets the App Logs dialog in the UI surface logs. + env = { + "TG_UI_PORT": str(args.port), + "TG_API_PORT": str(args.api_port), + "TESTGEN_LOG_FILE_PATH": str(TESTGEN_LOG_FILE_PATH), + } + if args.ssl_cert_file: + env["SSL_CERT_FILE"] = args.ssl_cert_file + if args.ssl_key_file: + env["SSL_KEY_FILE"] = args.ssl_key_file + run_testgen_cli( + action, + args, + "standalone-setup", + "--username", + self.username, + "--password", + self.password, + env=env, + redact=(self.password,), + ) + + def on_action_success(self, action, args): + cred_file_path = action.data_folder.joinpath(CREDENTIALS_FILE.format(args.prod)) + log_path = simplify_path(TESTGEN_LOG_FILE_PATH) + with CONSOLE.tee(cred_file_path) as console_tee: + console_tee(f"User Interface: {get_tg_url(args, args.port)}") + console_tee(f"API & MCP: {get_tg_url(args, args.api_port)}") + console_tee(f"Logs: {log_path}") + console_tee("") + console_tee(f"Username: {self.username}") + console_tee(f"Password: {self.password}", skip_logging=True) + CONSOLE.msg(f"(Credentials also written to {simplify_path(cred_file_path)})") + CONSOLE.space() + + +class TestgenQuickStartStep(Step): + """Generate demo data so the user has something to look at right after + install. Mode-agnostic — dispatches via ``run_testgen_cli``. Non-blocking: + failure here logs and continues; the user can run ``tg run-demo`` later. + """ + + label = "Generating demo data" + required = False + + def execute(self, action, args): + if not getattr(args, "generate_demo", True): + raise SkipStep + run_testgen_cli(action, args, "quick-start") + + class TestgenInstallAction(ComposeActionMixin, AnalyticsMultiStepAction): - steps = [ + """Install TestGen via either pip (uv-managed) or Docker Compose. + + Mode is chosen by ``--pip`` / ``--docker``, or auto-detected when neither + is given (defaults to pip if Docker prerequisites are not met, otherwise + prompts). + """ + + pip_steps = [UvBootstrapStep, UvToolInstallStep, TestgenStandaloneSetupStep, TestgenQuickStartStep] + docker_steps = [ ComposeVerifyExistingInstallStep, DockerNetworkStep, TestGenCreateDockerComposeFileStep, @@ -2031,17 +2658,42 @@ class TestgenInstallAction(ComposeActionMixin, AnalyticsMultiStepAction): ComposeStartStep, TestGenSetupDatabaseStep, TestGenUpgradeDatabaseStep, + TestgenQuickStartStep, + ] + pip_intro = [ + "Installing TestGen with pip.", + "The process may take 4~8 minutes depending on your system resources and network speed.", ] + docker_intro = [ + "Installing TestGen with Docker Compose.", + "This process may take 5~10 minutes depending on your system resources and network speed.", + ] + docker_requirements = [REQ_DOCKER, REQ_DOCKER_DAEMON, REQ_TESTGEN_IMAGE] + args_cmd = "install" label = "Installation" title = "Install TestGen" - intro_text = ["This process may take 5~10 minutes depending on your system resources and network speed."] - - args_cmd = "install" - requirements = [REQ_DOCKER, REQ_DOCKER_DAEMON, REQ_TESTGEN_IMAGE] + _per_invocation_attrs = (*MultiStepAction._per_invocation_attrs, "_resolved_mode", "steps", "intro_text") def get_parser(self, sub_parsers): parser = super().get_parser(sub_parsers) + mode_group = parser.add_mutually_exclusive_group() + mode_group.add_argument( + "--pip", + dest="install_mode", + action="store_const", + const=INSTALL_MODE_PIP, + help="Install TestGen with pip.", + ) + mode_group.add_argument( + "--docker", + dest="install_mode", + action="store_const", + const=INSTALL_MODE_DOCKER, + help="Install TestGen with Docker Compose.", + ) + parser.set_defaults(install_mode=None) + # Args supported by both modes parser.add_argument( "--port", dest="port", @@ -2056,12 +2708,33 @@ def get_parser(self, sub_parsers): default=TESTGEN_DEFAULT_API_PORT, help="Which port will be used to access TestGen's API and MCP server. Defaults to %(default)s", ) + parser.add_argument( + "--ssl-cert-file", + dest="ssl_cert_file", + action="store", + default=None, + help="Path to SSL certificate file.", + ) + parser.add_argument( + "--ssl-key-file", + dest="ssl_key_file", + action="store", + default=None, + help="Path to SSL key file.", + ) + parser.add_argument( + "--no-demo", + dest="generate_demo", + action="store_false", + help="Skip generating demo data after install. Default is to generate.", + ) + # Docker-only args parser.add_argument( "--image", dest="image", action="store", default=TESTGEN_DEFAULT_IMAGE, - help="TestGen image to use for the install. Defaults to %(default)s", + help="(Docker mode only) TestGen image to use for the install. Defaults to %(default)s", ) parser.add_argument( "--pull-timeout", @@ -2069,29 +2742,114 @@ def get_parser(self, sub_parsers): action="store", default=TESTGEN_PULL_TIMEOUT, help=( - "Maximum amount of time in minutes that Docker will be allowed to pull the images. " + "(Docker mode only) Maximum amount of time in minutes that Docker will be allowed to pull the images. " "Defaults to '%(default)s'" ), ) - parser.add_argument( - "--ssl-cert-file", - dest="ssl_cert_file", - action="store", - default=None, - help="Path to SSL certificate file.", - ) - parser.add_argument( - "--ssl-key-file", - dest="ssl_key_file", - action="store", - default=None, - help="Path to SSL key file.", - ) return parser + def check_requirements(self, args): + if not hasattr(self, "_resolved_mode"): + self._resolve_install_mode(args) + super().check_requirements(args) + + def get_requirements(self, args): + if self._resolved_mode == INSTALL_MODE_DOCKER: + return self.docker_requirements + return [] + + def _resolve_install_mode(self, args): + existing = InstallMarker(self.data_folder, args.prod, args.compose_file_name).read() + if existing: + CONSOLE.msg(f"Found an existing TestGen {existing} installation in {self.data_folder}.") + CONSOLE.space() + CONSOLE.msg(f"To update it, {command_hint(args.prod, 'upgrade', 'Upgrade TestGen')}.") + CONSOLE.msg(f"To remove it and start over, {command_hint(args.prod, 'delete', 'Uninstall TestGen')}.") + CONSOLE.space() + raise AbortAction + + if args.install_mode in (INSTALL_MODE_PIP, INSTALL_MODE_DOCKER): + mode = args.install_mode + else: + mode = self._auto_select_mode(args) + + self._resolved_mode = mode + self.steps = self.pip_steps if mode == INSTALL_MODE_PIP else self.docker_steps + self.analytics.additional_properties["install_mode"] = mode + LOG.info("tg install resolved to %s mode", mode) + + def _auto_select_mode(self, args): + # Probe each Docker prerequisite individually so we can show per-prereq + # status to the user. If a prereq fails (e.g. Docker engine not running), + # the user can decide whether to fix it or fall back to pip. + prereq_results = [(req, req.check_availability(self, args, quiet=True)) for req in self.docker_requirements] + docker_ready = all(ok for _, ok in prereq_results) + + CONSOLE.space() + CONSOLE.msg("TestGen offers two installation modes:") + CONSOLE.space() + CONSOLE.msg("[d] Docker Compose (Recommended)") + CONSOLE.msg(" The most stable TestGen experience for persistent use.") + CONSOLE.msg(" Provides a fully managed environment with an isolated PostgreSQL container.") + prereq_status = " ".join(f"{'(✓)' if ok else '(X)'} {req.label or req.key}" for req, ok in prereq_results) + CONSOLE.msg(f" Prerequisites: {prereq_status}") + CONSOLE.space() + CONSOLE.msg("[p] Pip + embedded PostgreSQL") + CONSOLE.msg(" A lightweight Python installation suited for evaluation.") + CONSOLE.msg( + " Sets up an isolated Python environment and manages the PostgreSQL database on the file system." + ) + CONSOLE.space() + + if docker_ready: + prompt = f"{CONSOLE.MARGIN}Install with Docker [d] or pip [p]? (default: d): " + valid_default = "d" + else: + CONSOLE.msg("To install with Docker, fix the prerequisites and run the install again.") + CONSOLE.space() + prompt = f"{CONSOLE.MARGIN}Install with pip [p]? (default: p): " + valid_default = "p" + + while True: + try: + choice = input(prompt).strip().lower() + except (KeyboardInterrupt, EOFError): + print("") + raise AbortAction + if choice == "": + choice = valid_default + if docker_ready and choice in ("d", "docker"): + return INSTALL_MODE_DOCKER + if choice in ("p", "pip"): + return INSTALL_MODE_PIP + print(f"'{choice}' is not a valid option.") + + def execute(self, args): + self.intro_text = self.pip_intro if self._resolved_mode == INSTALL_MODE_PIP else self.docker_intro + super().execute(args) + InstallMarker(self.data_folder, args.prod).write(self._resolved_mode) + # Pip mode: keep the app running so the user has a one-command install + # experience. Docker mode already runs as detached containers via + # ``docker compose up --wait``, so no need to start anything here — + # but open the browser for parity with pip's Streamlit auto-launch. + if self._resolved_mode == INSTALL_MODE_PIP: + start_testgen_app(self, args) + else: + open_app_in_browser(get_tg_url(args, args.port)) + + +class TestgenStandaloneUpgradeStep(Step): + label = "Upgrading the application database" + + def execute(self, action, args): + run_testgen_cli(action, args, "upgrade-system-version") + class TestgenUpgradeAction(ComposeActionMixin, AnalyticsMultiStepAction): - steps = [ + """Upgrade an existing TestGen install. Mode is read from the install marker.""" + + pip_steps = [UvBootstrapStep, UvToolUpgradeStep, TestgenStandaloneUpgradeStep] + docker_steps = [ UpdateComposeFileStep, ComposeStopStep, ComposePullImagesStep, @@ -2100,13 +2858,41 @@ class TestgenUpgradeAction(ComposeActionMixin, AnalyticsMultiStepAction): TestGenUpgradeDatabaseStep, ] + args_cmd = "upgrade" label = "Upgrade" title = "Upgrade TestGen" intro_text = ["This process may take 5~10 minutes depending on your system resources and network speed."] + _per_invocation_attrs = (*MultiStepAction._per_invocation_attrs, "_resolved_mode", "steps", "intro_text") - args_cmd = "upgrade" + def get_parser(self, sub_parsers): + parser = super().get_parser(sub_parsers) + # Docker-only args - ignored in pip mode + parser.add_argument( + "--skip-verify", + dest="skip_verify", + action="store_true", + help="(Docker mode only) Whether to skip the version check before upgrading.", + ) + parser.add_argument( + "--pull-timeout", + type=int, + action="store", + default=TESTGEN_PULL_TIMEOUT, + help=( + "(Docker mode only) Maximum amount of time in minutes that Docker will be allowed to pull the images. " + "Defaults to '%(default)s'" + ), + ) + return parser + + def check_requirements(self, args): + if not hasattr(self, "_resolved_mode"): + self._resolve_install_mode(args) + super().check_requirements(args) def get_requirements(self, args): + if self._resolved_mode == INSTALL_MODE_PIP: + return [] return [ REQ_DOCKER, REQ_DOCKER_DAEMON, @@ -2120,34 +2906,198 @@ def get_requirements(self, args): "config", ), ( - f"TestGen's Docker configuration file is not available at {self.data_folder.joinpath(self.get_compose_file_path(args))}.", + f"TestGen's Docker configuration file is not available at " + f"{self.data_folder.joinpath(self.get_compose_file_path(args))}.", "Re-install TestGen and try again.", ), ), ] + def _resolve_install_mode(self, args): + mode = InstallMarker(self.data_folder, args.prod, args.compose_file_name).read() + if mode is None: + CONSOLE.msg(f"No TestGen installation found in {self.data_folder}.") + CONSOLE.msg(f"To install TestGen, {command_hint(args.prod, 'install', 'Install TestGen')}.") + CONSOLE.space() + raise AbortAction + self._resolved_mode = mode + self.steps = self.pip_steps if mode == INSTALL_MODE_PIP else self.docker_steps + self.analytics.additional_properties["install_mode"] = mode + LOG.info("tg upgrade resolved to %s mode", mode) + + def execute(self, args): + super().execute(args) + InstallMarker(self.data_folder, args.prod).write(self._resolved_mode) + + +class TestgenStartAction(Action, ComposeActionMixin): + """Start a previously-installed TestGen app. + + Companion to the auto-start at the end of ``tg install``. For pip mode, + runs ``testgen run-app`` and blocks until Ctrl+C. For docker mode, runs + ``docker compose up --wait`` (detached) so the user can bring containers + back up after a reboot or a manual stop. + """ + + args_cmd = "start" + _per_invocation_attrs = (*Action._per_invocation_attrs, "_resolved_mode") + + def check_requirements(self, args): + if not hasattr(self, "_resolved_mode"): + self._resolve_install_mode(args) + super().check_requirements(args) + + def get_requirements(self, args): + if self._resolved_mode == INSTALL_MODE_DOCKER: + return [REQ_DOCKER, REQ_DOCKER_DAEMON] + return [] + + def _resolve_install_mode(self, args): + mode = InstallMarker(self.data_folder, args.prod, args.compose_file_name).read() + if mode is None: + CONSOLE.msg(f"No TestGen installation found in {self.data_folder}.") + CONSOLE.msg(f"To install TestGen, {command_hint(args.prod, 'install', 'Install TestGen')}.") + CONSOLE.space() + raise AbortAction + self._resolved_mode = mode + self.analytics.additional_properties["install_mode"] = mode + + def execute(self, args): + if self._resolved_mode == INSTALL_MODE_DOCKER: + CONSOLE.title("Start TestGen") + compose_path = self.get_compose_file_path(args) + self.run_cmd("docker", "compose", "-f", compose_path, "up", "--wait") + CONSOLE.msg("TestGen containers are running.") + CONSOLE.msg( + f"For the URL and credentials, {command_hint(args.prod, 'access-info', 'Access Installed App')}." + ) + # Match pip-mode parity: open the browser to the configured UI URL. + # Best-effort — if the compose file is malformed or missing the env, + # we just skip the browser launch (the URL is still in the creds file). + try: + for line in compose_path.read_text().splitlines(): + stripped = line.strip() + if stripped.startswith("TG_UI_BASE_URL:"): + open_app_in_browser(stripped.split(":", 1)[1].strip()) + break + except OSError: + LOG.exception("Could not read TG_UI_BASE_URL from %s", compose_path) + else: + start_testgen_app(self, args) + + +class TestgenDeleteAction(Action, ComposeActionMixin): + """Delete an existing TestGen install — pip or Docker — based on the marker. + + Reuses the ``delete_compose_*`` helpers on ``ComposeActionMixin`` for the + Docker path. The marker is removed at the end so a subsequent + ``tg install`` starts fresh. + """ + + args_cmd = "delete" + _per_invocation_attrs = (*Action._per_invocation_attrs, "_resolved_mode") + def get_parser(self, sub_parsers): parser = super().get_parser(sub_parsers) parser.add_argument( - "--skip-verify", - dest="skip_verify", + "--keep-images", action="store_true", - help="Whether to skip the version check before upgrading.", + help="(Docker mode only) Does not delete the images when deleting the installation", ) parser.add_argument( - "--pull-timeout", - type=int, - action="store", - default=TESTGEN_PULL_TIMEOUT, - help=( - "Maximum amount of time in minutes that Docker will be allowed to pull the images. " - "Defaults to '%(default)s'" - ), + "--keep-config", + action="store_true", + help="(Docker mode only) Does not delete the compose config file when deleting the installation", + ) + parser.add_argument( + "--keep-data", + action="store_true", + help="(Pip mode only) Keep the embedded Postgres data directory (~/.testgen by default).", ) + return parser + + def check_requirements(self, args): + if not hasattr(self, "_resolved_mode"): + self._resolve_install_mode(args) + super().check_requirements(args) + + def get_requirements(self, args): + if self._resolved_mode == INSTALL_MODE_DOCKER: + return [REQ_DOCKER, REQ_DOCKER_DAEMON] + return [] + + def _resolve_install_mode(self, args): + # Unlike install/upgrade, "no install found" is not an abort here — + # ``tg delete`` is idempotent. execute() handles the None case. + mode = InstallMarker(self.data_folder, args.prod, args.compose_file_name).read() + self._resolved_mode = mode + if mode is not None: + self.analytics.additional_properties["install_mode"] = mode + + def execute(self, args): + if self._resolved_mode is None: + CONSOLE.msg(f"No TestGen installation found in {self.data_folder}.") + CONSOLE.msg("Nothing to delete.") + CONSOLE.space() + return + + if self._resolved_mode == INSTALL_MODE_DOCKER: + self._delete_docker(args) + else: + self._delete_pip(args) + InstallMarker(self.data_folder, args.prod, args.compose_file_name).unlink() + + def _delete_docker(self, args): + if self.get_compose_file_path(args).exists(): + self.delete_compose_containers(args) + self.delete_compose_network() + else: + self.delete_compose_network() + self.delete_compose_volumes(args) + + def _delete_pip(self, args): + CONSOLE.title("Delete TestGen instance") + + uv_path = resolve_uv_path(self.data_folder) + if uv_path: + try: + self.run_cmd(uv_path, "tool", "uninstall", TESTGEN_PIP_PACKAGE) + except CommandFailed: + LOG.exception("Failed to uninstall testgen via uv") + CONSOLE.msg( + "Note: 'uv tool uninstall testgen' reported an error " + "(it may already be uninstalled); see session logs." + ) + else: + LOG.info("uv not found; skipping uv tool uninstall") + CONSOLE.msg("uv not found; skipping 'uv tool uninstall testgen'.") + + if not getattr(args, "keep_data", False): + tg_home = pathlib.Path(os.environ.get("TG_TESTGEN_HOME", pathlib.Path.home() / ".testgen")) + remove_path(tg_home, label="TestGen data directory") + + # Don't touch ~/.streamlit — Streamlit is widely used and the user + # may have other Streamlit projects on this machine. The config dir + # is tiny and harmless if left behind. + + # Remove the installer-local uv binary if we downloaded one. A + # pre-existing uv on PATH is left alone. + local_uv = self.data_folder / UV_BIN_SUBDIR / ("uv.exe" if platform.system() == "Windows" else "uv") + if remove_path(local_uv, label="installer-local uv"): + with contextlib.suppress(OSError): + local_uv.parent.rmdir() + + remove_path(self.data_folder / CREDENTIALS_FILE.format(args.prod)) + CONSOLE.space() + CONSOLE.msg("TestGen uninstalled.") + CONSOLE.space() class TestgenRunDemoAction(DemoContainerAction, ComposeActionMixin): + """Generate TestGen demo data — Docker-exec or pip-direct based on the marker.""" + args_cmd = "run-demo" + _per_invocation_attrs = (*Action._per_invocation_attrs, "_resolved_mode") def get_parser(self, sub_parsers): parser = super().get_parser(sub_parsers) @@ -2160,103 +3110,120 @@ def get_parser(self, sub_parsers): ) return parser + def check_requirements(self, args): + if not hasattr(self, "_resolved_mode"): + self._resolve_install_mode(args) + super().check_requirements(args) + + def get_requirements(self, args): + # Docker mode requires Docker. For pip mode, Docker is only needed when + # the user asked to export to Observability (the dk-demo container + # generates the export payload). + if self._resolved_mode == INSTALL_MODE_DOCKER or getattr(args, "obs_export", False): + return [REQ_DOCKER, REQ_DOCKER_DAEMON] + return [] + + def _resolve_install_mode(self, args): + mode = InstallMarker(self.data_folder, args.prod, args.compose_file_name).read() + if mode is None: + CONSOLE.msg(f"No TestGen installation found in {self.data_folder}.") + CONSOLE.msg(f"To install TestGen, {command_hint(args.prod, 'install', 'Install TestGen')}.") + CONSOLE.space() + raise AbortAction + self._resolved_mode = mode + self.analytics.additional_properties["install_mode"] = mode + def execute(self, args): self.analytics.additional_properties["obs_export"] = args.obs_export CONSOLE.title("Run TestGen demo") - tg_status = self.get_status(args) - if not tg_status or not re.match(".*running.*", tg_status["Status"], re.I): - CONSOLE.msg("Running the TestGen demo requires the platform to be running.") - raise AbortAction - if args.obs_export and not (self.data_folder / DEMO_CONFIG_FILE).exists(): CONSOLE.msg("Observability demo configuration missing.") raise AbortAction + if self._resolved_mode == INSTALL_MODE_DOCKER: + tg_status = self.get_status(args) + if not tg_status or not re.match(".*running.*", tg_status["Status"], re.I): + CONSOLE.msg("Running the TestGen demo requires the application to be running.") + raise AbortAction + CONSOLE.msg("This process may take up to 3 minutes depending on your system resources and network speed.") CONSOLE.space() + export_args = [] if args.obs_export: self.run_dk_demo_container("tg-run-demo") - - quick_start_command = [ - "testgen", - "quick-start", - ] - if args.obs_export: with open(self.data_folder / DEMO_CONFIG_FILE, "r") as file: json_config = json.load(file) - - quick_start_command.extend( - [ - "--observability-api-url", - json_config["api_host"], - "--observability-api-key", - json_config["api_key"], - ] - ) - - cli_commands = [ - quick_start_command, - ] + export_args = [ + "--observability-api-url", + json_config["api_host"], + "--observability-api-key", + json_config["api_key"], + ] + + run_testgen_cli(self, args, "quick-start", *export_args) if args.obs_export: - cli_commands.append( - [ - "testgen", - "export-observability", - "--project-key", - "DEFAULT", - "--test-suite-key", - "default-suite-1", - ] - ) - - for command in cli_commands: - CONSOLE.msg(f"Running command : docker compose exec engine {' '.join(command)}") - self.run_cmd( - "docker", - "compose", - "-f", - self.get_compose_file_path(args), - "exec", - "engine", - *command, + run_testgen_cli( + self, + args, + "export-observability", + "--project-key", + "DEFAULT", + "--test-suite-key", + "default-suite-1", ) CONSOLE.title("Demo SUCCEEDED") class TestgenDeleteDemoAction(DemoContainerAction, ComposeActionMixin): + """Delete TestGen demo data — Docker-exec or pip-direct based on the marker.""" + args_cmd = "delete-demo" + _per_invocation_attrs = (*Action._per_invocation_attrs, "_resolved_mode") + + def check_requirements(self, args): + if not hasattr(self, "_resolved_mode"): + self._resolve_install_mode(args) + super().check_requirements(args) + + def get_requirements(self, args): + # Docker mode requires Docker. For pip mode, the dk-demo container + # call below is wrapped in try/except so Docker absence is non-fatal. + return [REQ_DOCKER, REQ_DOCKER_DAEMON] if self._resolved_mode == INSTALL_MODE_DOCKER else [] + + def _resolve_install_mode(self, args): + # Like delete: idempotent, so "no install" returns rather than aborts. + mode = InstallMarker(self.data_folder, args.prod, args.compose_file_name).read() + self._resolved_mode = mode + if mode is not None: + self.analytics.additional_properties["install_mode"] = mode def execute(self, args): + if self._resolved_mode is None: + CONSOLE.msg(f"No TestGen installation found in {self.data_folder}.") + CONSOLE.msg("Nothing to delete.") + CONSOLE.space() + return + CONSOLE.title("Delete TestGen demo") try: self.run_dk_demo_container("tg-delete-demo") except Exception: pass - CONSOLE.msg("Cleaning up system database..") - tg_status = self.get_status(args) - if tg_status: - self.run_cmd( - "docker", - "compose", - "-f", - self.get_compose_file_path(args), - "exec", - "engine", - "testgen", - "setup-system-db", - "--delete-db", - "--yes", - ) + CONSOLE.msg("Cleaning up application database..") + if self._resolved_mode == INSTALL_MODE_DOCKER: + tg_status = self.get_status(args) + if not tg_status: + CONSOLE.msg("TestGen must be running for its demo data to be cleaned.") + raise AbortAction - CONSOLE.title("Demo data DELETED") - else: - CONSOLE.msg("TestGen must be running for its demo data to be cleaned.") - raise AbortAction + run_testgen_cli(self, args, "setup-system-db", "--delete-db", "--yes") + + CONSOLE.title("Demo data DELETED") class AccessInstructionsAction(Action): @@ -2297,6 +3264,7 @@ def run_installer(args): tg_menu = Menu(run_installer, "TestGen") tg_menu.add_option("Install TestGen", ["tg", "install"]) + tg_menu.add_option("Start TestGen", ["tg", "start"]) tg_menu.add_option("Upgrade TestGen", ["tg", "upgrade"]) tg_menu.add_option("Access Installed App", ["tg", "access-info"]) tg_menu.add_option("Install TestGen demo data", ["tg", "run-demo"]) @@ -2361,15 +3329,16 @@ def get_installer_instance(): "tg", [ TestgenInstallAction(), + TestgenStartAction(), TestgenUpgradeAction(), AccessInstructionsAction(), - ComposeDeleteAction(), + TestgenDeleteAction(), TestgenRunDemoAction(), TestgenDeleteDemoAction(), ], defaults={ "prod_name": "TestGen", - "compose_file_name": "docker-compose.yml", + "compose_file_name": TESTGEN_COMPOSE_FILE, "compose_project_name": "dataops-testgen", }, ) @@ -2380,9 +3349,6 @@ def get_installer_instance(): installer = get_installer_instance() if platform.system() == "Windows": - if platform.win32_edition() == "Core": - print("\nWARNING: Your Windows edition is not compatible with Docker.") - # For backward compatibility - move the old folder to the new path try: new_folder = pathlib.Path(os.environ["LOCALAPPDATA"], "DataKitchenApps") diff --git a/tests/conftest.py b/tests/conftest.py index 92484de..c57139c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -10,6 +10,13 @@ from tests.installer import CONSOLE, Action, TESTGEN_DEFAULT_IMAGE +@pytest.fixture(autouse=True) +def _no_real_browser_launch(): + """Stop the installer from actually opening a browser tab during tests.""" + with patch("tests.installer.webbrowser.open", return_value=False) as mock: + yield mock + + @pytest.fixture def stdout_mock(): return Mock(return_value=[]) @@ -134,6 +141,7 @@ def action_cls(analytics_mock): with ( patch.object(Action, "session_zip", create=True), patch.object(Action, "session_folder", create=True), + patch.object(Action, "logs_folder", create=True), patch.object(Action, "_cmd_idx", create=True, new=0), patch.object(Action, "configure_logging", create=True), patch.object(Action, "init_session_folder", create=True), @@ -173,14 +181,18 @@ def args_mock(): ns.api_port = 8530 ns.keep_images = False ns.keep_config = False + ns.keep_data = False ns.skip_verify = False # TestGen defaults - ns.pull_timeout = 10 + ns.pull_timeout = 5 ns.ssl_key_file = None ns.ssl_cert_file = None ns.image = TESTGEN_DEFAULT_IMAGE ns.obs_export = False + ns.install_mode = None + ns.generate_demo = True + ns.api_port = 8530 # Observability defaults ns.ui_image = "datakitchen/dataops-observability-ui:v2" diff --git a/tests/test_install_type.py b/tests/test_install_type.py new file mode 100644 index 0000000..a2215ee --- /dev/null +++ b/tests/test_install_type.py @@ -0,0 +1,123 @@ +import json +from pathlib import Path + +import pytest + +from tests.installer import ( + CREDENTIALS_FILE, + INSTALL_MARKER_FILE, + INSTALL_MODE_DOCKER, + INSTALL_MODE_PIP, + InstallMarker, + TESTGEN_COMPOSE_FILE, +) + + +@pytest.fixture +def data_folder(tmp_data_folder): + return Path(tmp_data_folder) + + +@pytest.mark.unit +def test_read_install_mode_returns_none_when_empty(data_folder): + assert InstallMarker(data_folder, "tg", "docker-compose.yml").read() is None + + +@pytest.mark.unit +@pytest.mark.parametrize("install_mode", [INSTALL_MODE_DOCKER, INSTALL_MODE_PIP]) +def test_read_install_mode_from_marker(data_folder, install_mode): + (data_folder / INSTALL_MARKER_FILE.format("tg")).write_text(json.dumps({"install_mode": install_mode})) + + assert InstallMarker(data_folder, "tg", "docker-compose.yml").read() == install_mode + + +@pytest.mark.unit +def test_read_install_mode_legacy_docker_backfill(data_folder): + (data_folder / TESTGEN_COMPOSE_FILE).write_text("version: '3'") + (data_folder / CREDENTIALS_FILE.format("tg")).write_text("admin\n") + + assert InstallMarker(data_folder, "tg", "docker-compose.yml").read() == INSTALL_MODE_DOCKER + + +@pytest.mark.unit +def test_read_install_mode_legacy_backfill_honors_compose_file_name(data_folder): + # Verifies the marker isn't TestGen-specific: pass a different product's + # compose file name and it should detect that product's legacy install. + (data_folder / "obs-docker-compose.yml").write_text("version: '3'") + (data_folder / CREDENTIALS_FILE.format("obs")).write_text("admin\n") + + assert InstallMarker(data_folder, "obs", "obs-docker-compose.yml").read() == INSTALL_MODE_DOCKER + # And does NOT match if we point at the wrong compose file name. + assert InstallMarker(data_folder, "obs", "docker-compose.yml").read() is None + + +@pytest.mark.unit +def test_read_install_mode_legacy_requires_both_files(data_folder): + # Only compose file, missing credentials → not a legacy install + (data_folder / TESTGEN_COMPOSE_FILE).write_text("version: '3'") + + assert InstallMarker(data_folder, "tg", "docker-compose.yml").read() is None + + +@pytest.mark.unit +def test_read_install_mode_malformed_marker_falls_back_to_legacy(data_folder): + (data_folder / INSTALL_MARKER_FILE.format("tg")).write_text("{not valid json") + (data_folder / TESTGEN_COMPOSE_FILE).write_text("version: '3'") + (data_folder / CREDENTIALS_FILE.format("tg")).write_text("admin\n") + + assert InstallMarker(data_folder, "tg", "docker-compose.yml").read() == INSTALL_MODE_DOCKER + + +@pytest.mark.unit +def test_read_install_mode_unknown_value_falls_back(data_folder): + (data_folder / INSTALL_MARKER_FILE.format("tg")).write_text(json.dumps({"install_mode": "bogus"})) + + assert InstallMarker(data_folder, "tg", "docker-compose.yml").read() is None + + +@pytest.mark.unit +def test_write_install_marker_round_trip(data_folder): + InstallMarker(data_folder, "tg", "docker-compose.yml").write( + INSTALL_MODE_PIP, version="5.9.4", python_version="3.13.1" + ) + + data = json.loads((data_folder / INSTALL_MARKER_FILE.format("tg")).read_text()) + assert data["install_mode"] == INSTALL_MODE_PIP + assert data["version"] == "5.9.4" + assert data["python_version"] == "3.13.1" + assert "created_on" in data + assert "last_updated_on" in data + assert InstallMarker(data_folder, "tg", "docker-compose.yml").read() == INSTALL_MODE_PIP + + +@pytest.mark.unit +def test_write_install_marker_preserves_created_on_across_writes(data_folder): + InstallMarker(data_folder, "tg").write(INSTALL_MODE_PIP, version="5.9.4") + initial = json.loads((data_folder / INSTALL_MARKER_FILE.format("tg")).read_text()) + + InstallMarker(data_folder, "tg").write(INSTALL_MODE_PIP, version="5.10.0") + after = json.loads((data_folder / INSTALL_MARKER_FILE.format("tg")).read_text()) + + assert after["created_on"] == initial["created_on"] + assert after["version"] == "5.10.0" + + +@pytest.mark.unit +def test_write_install_marker_rejects_unknown_type(data_folder): + with pytest.raises(ValueError, match="Unknown install_mode"): + InstallMarker(data_folder, "tg").write("sideways") + + +@pytest.mark.unit +def test_marker_unlink_removes_file(data_folder): + InstallMarker(data_folder, "tg").write(INSTALL_MODE_PIP) + assert (data_folder / INSTALL_MARKER_FILE.format("tg")).exists() + + InstallMarker(data_folder, "tg").unlink() + assert not (data_folder / INSTALL_MARKER_FILE.format("tg")).exists() + + +@pytest.mark.unit +def test_marker_unlink_is_idempotent(data_folder): + # No marker present — unlink should be a no-op, not raise. + InstallMarker(data_folder, "tg").unlink() diff --git a/tests/test_tg_install.py b/tests/test_tg_install.py index 8922ec4..ce2cf5f 100644 --- a/tests/test_tg_install.py +++ b/tests/test_tg_install.py @@ -5,6 +5,7 @@ import pytest from tests.installer import ( + INSTALL_MODE_DOCKER, TestgenInstallAction, AbortAction, TestGenCreateDockerComposeFileStep, @@ -17,6 +18,11 @@ def tg_install_action(action_cls, args_mock, tmp_data_folder, start_cmd_mock): action = TestgenInstallAction() args_mock.prod = "tg" args_mock.action = "install" + args_mock.install_mode = INSTALL_MODE_DOCKER + # Bypass check_requirements: pre-resolve mode to Docker and seed the step + # list so execute() runs the Docker pipeline directly. + action._resolved_mode = INSTALL_MODE_DOCKER + action.steps = action.docker_steps with patch.object(action, "execute", new=partial(action.execute, args_mock)): yield action @@ -41,6 +47,11 @@ def test_tg_install(tg_install_action, start_cmd_mock, stdout_mock, tmp_data_fol assert Path(tmp_data_folder).joinpath("test-compose.yml").stat().st_size > 0 assert Path(tmp_data_folder).joinpath("dk-tg-credentials.txt").stat().st_size > 0 + marker = Path(tmp_data_folder).joinpath("dk-tg-install.json") + assert marker.exists() + import json as _json + + assert _json.loads(marker.read_text())["install_mode"] == "docker" @pytest.mark.integration @@ -110,3 +121,33 @@ def test_tg_compose_base_url_ssl(tg_install_action, start_cmd_mock, stdout_mock, tg_install_action.execute() contents = compose_path.read_text() assert "TG_UI_BASE_URL: https://localhost:8501" in contents + + +@pytest.mark.integration +def test_tg_docker_install_auto_runs_demo(tg_install_action, start_cmd_mock, compose_path): + """Docker install also generates demo data so the user has something on first launch.""" + tg_install_action.execute() + + start_cmd_mock.assert_any_call( + "docker", + "compose", + "-f", + compose_path, + "exec", + "engine", + "testgen", + "quick-start", + raise_on_non_zero=True, + env=None, + ) + + +@pytest.mark.integration +def test_tg_docker_install_no_demo_flag_skips_quick_start(tg_install_action, args_mock, start_cmd_mock): + """--no-demo opts out of the auto-demo step in Docker mode.""" + args_mock.generate_demo = False + + tg_install_action.execute() + + for invocation in start_cmd_mock.call_args_list: + assert "quick-start" not in invocation.args, f"quick-start should be skipped, got: {invocation}" diff --git a/tests/test_tg_pip_delete.py b/tests/test_tg_pip_delete.py new file mode 100644 index 0000000..9a1fd43 --- /dev/null +++ b/tests/test_tg_pip_delete.py @@ -0,0 +1,188 @@ +from functools import partial +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +from tests.installer import ( + INSTALL_MARKER_FILE, + INSTALL_MODE_DOCKER, + INSTALL_MODE_PIP, + TestgenDeleteAction, + InstallMarker, +) + + +@pytest.fixture +def pip_delete_action(action_cls, args_mock, tmp_data_folder, start_cmd_mock, tmp_path): + """Drive the pip path of TestgenDeleteAction directly via _delete_pip.""" + action = TestgenDeleteAction() + args_mock.prod = "tg" + args_mock.action = "delete" + InstallMarker(action.data_folder, args_mock.prod).write(INSTALL_MODE_PIP) + # Bypass check_requirements: pre-resolve mode so execute() runs the + # delete branch directly. + action._resolved_mode = INSTALL_MODE_PIP + # Pin Path.home() to tmp_path so .streamlit/ cleanup doesn't touch the + # real home directory of whoever runs the suite. + with ( + patch.object(action, "execute", new=partial(action.execute, args_mock)), + patch("tests.installer.pathlib.Path.home", return_value=tmp_path), + ): + yield action + + +@pytest.mark.integration +def test_pip_delete_removes_uv_tool_and_home(pip_delete_action, start_cmd_mock, tmp_data_folder, tmp_path): + fake_tg_home = tmp_path / ".testgen" + fake_tg_home.mkdir() + (fake_tg_home / "config.env").write_text("TESTGEN_USERNAME=admin\n") + + # Streamlit config dir (if any) must be left alone — user may have other + # Streamlit projects on this machine. + fake_streamlit_dir = tmp_path / ".streamlit" + fake_streamlit_dir.mkdir() + (fake_streamlit_dir / "credentials.toml").write_text('[general]\nemail = ""\n') + + (Path(tmp_data_folder) / "dk-tg-credentials.txt").write_text("Username: admin\n") + + with ( + patch("tests.installer.shutil.which", return_value="/usr/local/bin/uv"), + patch.dict("tests.installer.os.environ", {"TG_TESTGEN_HOME": str(fake_tg_home)}), + ): + pip_delete_action.execute() + + start_cmd_mock.assert_any_call( + "/usr/local/bin/uv", + "tool", + "uninstall", + "dataops-testgen", + raise_on_non_zero=True, + env=None, + ) + assert not fake_tg_home.exists() + assert fake_streamlit_dir.exists() + assert (fake_streamlit_dir / "credentials.toml").exists() + assert not (Path(tmp_data_folder) / "dk-tg-credentials.txt").exists() + # Marker is removed too — a subsequent install should start clean. + assert not (Path(tmp_data_folder) / INSTALL_MARKER_FILE.format("tg")).exists() + + +@pytest.mark.integration +def test_pip_delete_removes_installer_local_uv(pip_delete_action, start_cmd_mock, tmp_data_folder, tmp_path): + local_bin = Path(tmp_data_folder) / "bin" + local_bin.mkdir() + local_uv = local_bin / "uv" + local_uv.write_bytes(b"#!/bin/sh\n") + + fake_tg_home = tmp_path / ".testgen" + fake_tg_home.mkdir() + + with ( + patch("tests.installer.shutil.which", return_value=None), + patch.dict("tests.installer.os.environ", {"TG_TESTGEN_HOME": str(fake_tg_home)}), + ): + pip_delete_action.execute() + + assert not local_uv.exists() + assert not local_bin.exists() + + +@pytest.mark.integration +def test_pip_delete_leaves_path_uv_alone(pip_delete_action, start_cmd_mock, tmp_data_folder, tmp_path): + fake_tg_home = tmp_path / ".testgen" + fake_tg_home.mkdir() + + with ( + patch("tests.installer.shutil.which", return_value="/usr/local/bin/uv"), + patch.dict("tests.installer.os.environ", {"TG_TESTGEN_HOME": str(fake_tg_home)}), + ): + pip_delete_action.execute() + + assert not (Path(tmp_data_folder) / "bin").exists() + + +@pytest.mark.integration +def test_pip_delete_respects_keep_data(pip_delete_action, start_cmd_mock, tmp_data_folder, args_mock, tmp_path): + args_mock.keep_data = True + fake_tg_home = tmp_path / ".testgen" + fake_tg_home.mkdir() + (fake_tg_home / "config.env").write_text("TESTGEN_USERNAME=admin\n") + + with ( + patch("tests.installer.shutil.which", return_value="/usr/local/bin/uv"), + patch.dict("tests.installer.os.environ", {"TG_TESTGEN_HOME": str(fake_tg_home)}), + ): + pip_delete_action.execute() + + assert fake_tg_home.exists() + assert (fake_tg_home / "config.env").exists() + + +@pytest.mark.integration +def test_pip_delete_handles_missing_uv(pip_delete_action, start_cmd_mock, tmp_data_folder, tmp_path, console_msg_mock): + fake_tg_home = tmp_path / ".testgen" + fake_tg_home.mkdir() + + with ( + patch("tests.installer.shutil.which", return_value=None), + patch.dict("tests.installer.os.environ", {"TG_TESTGEN_HOME": str(fake_tg_home)}), + ): + pip_delete_action.execute() + + for invocation in start_cmd_mock.call_args_list: + assert "tool" not in invocation.args, f"Unexpected uv invocation: {invocation}" + console_msg_mock.assert_any_msg_contains("uv not found") + assert not fake_tg_home.exists() + + +@pytest.fixture +def delete_action(action_cls, args_mock, tmp_data_folder): + """A bare TestgenDeleteAction for testing the marker-driven dispatch layer.""" + action = TestgenDeleteAction() + args_mock.prod = "tg" + args_mock.action = "delete" + action.analytics = MagicMock() + action.analytics.additional_properties = {} + return action + + +@pytest.mark.integration +def test_delete_nothing_to_delete(delete_action, args_mock, console_msg_mock): + delete_action._resolve_install_mode(args_mock) + delete_action.execute(args_mock) + console_msg_mock.assert_any_msg_contains("Nothing to delete") + # No analytics property recorded since there was no install. + assert "install_mode" not in delete_action.analytics.additional_properties + + +@pytest.mark.integration +def test_delete_routes_to_pip_and_removes_marker(delete_action, args_mock, tmp_data_folder, tmp_path): + InstallMarker(Path(tmp_data_folder), "tg").write(INSTALL_MODE_PIP) + assert (Path(tmp_data_folder) / INSTALL_MARKER_FILE.format("tg")).exists() + + delete_action._resolve_install_mode(args_mock) + with ( + patch.object(delete_action, "_delete_pip") as pip_branch, + patch("tests.installer.pathlib.Path.home", return_value=tmp_path), + ): + delete_action.execute(args_mock) + + pip_branch.assert_called_once_with(args_mock) + assert not (Path(tmp_data_folder) / INSTALL_MARKER_FILE.format("tg")).exists() + assert delete_action.analytics.additional_properties["install_mode"] == INSTALL_MODE_PIP + + +@pytest.mark.integration +def test_delete_routes_to_docker_legacy(delete_action, args_mock, tmp_data_folder): + # No marker but legacy Docker files present — read_install_mode falls back + # to Docker, and the unified action takes the docker branch. + (Path(tmp_data_folder) / args_mock.compose_file_name).write_text("version: '3'") + (Path(tmp_data_folder) / "dk-tg-credentials.txt").write_text("admin\n") + + delete_action._resolve_install_mode(args_mock) + with patch.object(delete_action, "_delete_docker") as docker_branch: + delete_action.execute(args_mock) + + docker_branch.assert_called_once_with(args_mock) + assert delete_action.analytics.additional_properties["install_mode"] == INSTALL_MODE_DOCKER diff --git a/tests/test_tg_pip_demo.py b/tests/test_tg_pip_demo.py new file mode 100644 index 0000000..7bc7605 --- /dev/null +++ b/tests/test_tg_pip_demo.py @@ -0,0 +1,249 @@ +from functools import partial +from pathlib import Path +from unittest.mock import MagicMock, call, patch + +import pytest + +from tests.installer import ( + AbortAction, + INSTALL_MODE_DOCKER, + INSTALL_MODE_PIP, + TestgenDeleteDemoAction, + TestgenRunDemoAction, + InstallMarker, +) + + +UV_PATH = "/usr/local/bin/uv" +TESTGEN_PATH = "/Users/test/.local/bin/testgen" + + +@pytest.fixture +def pip_run_demo_action(action_cls, args_mock, tmp_data_folder, start_cmd_mock): + """Drive the pip path of TestgenRunDemoAction.""" + action = TestgenRunDemoAction() + args_mock.prod = "tg" + args_mock.action = "run-demo" + InstallMarker(action.data_folder, args_mock.prod).write(INSTALL_MODE_PIP) + # Bypass check_requirements: pre-resolve mode so execute() runs directly. + action._resolved_mode = INSTALL_MODE_PIP + with ( + patch("tests.installer.shutil.which", return_value=UV_PATH), + patch("tests.installer.resolve_testgen_path", return_value=TESTGEN_PATH), + patch.object(action, "execute", new=partial(action.execute, args_mock)), + ): + yield action + + +@pytest.mark.integration +def test_pip_run_demo_without_export(pip_run_demo_action, start_cmd_mock, args_mock): + args_mock.obs_export = False + + pip_run_demo_action.execute() + + start_cmd_mock.assert_any_call( + TESTGEN_PATH, + "quick-start", + raise_on_non_zero=True, + env=None, + ) + for invocation in start_cmd_mock.call_args_list: + assert "export-observability" not in invocation.args + assert "datakitchen/data-observability-demo:latest" not in invocation.args + + +@pytest.mark.integration +def test_pip_run_demo_with_export(pip_run_demo_action, start_cmd_mock, args_mock, demo_config_path, tmp_data_folder): + args_mock.obs_export = True + + pip_run_demo_action.execute() + + start_cmd_mock.assert_has_calls( + [ + call( + TESTGEN_PATH, + "quick-start", + "--observability-api-url", + "demo-api-host", + "--observability-api-key", + "demo-api-key", + raise_on_non_zero=True, + env=None, + ), + call( + TESTGEN_PATH, + "export-observability", + "--project-key", + "DEFAULT", + "--test-suite-key", + "default-suite-1", + raise_on_non_zero=True, + env=None, + ), + call( + "docker", + "run", + "--rm", + "--mount", + f"type=bind,source={str(demo_config_path)},target=/dk/demo-config.json", + "--name", + "dk-demo", + "--network", + "datakitchen-network", + "--add-host", + "host.docker.internal:host-gateway", + "datakitchen/data-observability-demo:latest", + "tg-run-demo", + raise_on_non_zero=True, + env=None, + ), + ], + any_order=True, + ) + + +@pytest.mark.integration +def test_pip_run_demo_export_aborts_when_demo_config_missing(pip_run_demo_action, args_mock, console_msg_mock): + args_mock.obs_export = True + + with pytest.raises(AbortAction): + pip_run_demo_action.execute() + + console_msg_mock.assert_any_msg_contains("Observability demo configuration missing") + + +@pytest.mark.integration +def test_pip_run_demo_aborts_without_uv(action_cls, args_mock, tmp_data_folder, start_cmd_mock, console_msg_mock): + action = TestgenRunDemoAction() + args_mock.prod = "tg" + args_mock.action = "run-demo" + args_mock.obs_export = False + InstallMarker(action.data_folder, args_mock.prod).write(INSTALL_MODE_PIP) + action._resolved_mode = INSTALL_MODE_PIP + + with ( + patch("tests.installer.shutil.which", return_value=None), + patch.object(action, "execute", new=partial(action.execute, args_mock)), + pytest.raises(AbortAction), + ): + action.execute() + + console_msg_mock.assert_any_msg_contains("uv not found") + + +@pytest.fixture +def pip_delete_demo_action(action_cls, args_mock, tmp_data_folder, start_cmd_mock): + action = TestgenDeleteDemoAction() + args_mock.prod = "tg" + args_mock.action = "delete-demo" + InstallMarker(action.data_folder, args_mock.prod).write(INSTALL_MODE_PIP) + action._resolved_mode = INSTALL_MODE_PIP + with ( + patch("tests.installer.shutil.which", return_value=UV_PATH), + patch("tests.installer.resolve_testgen_path", return_value=TESTGEN_PATH), + patch.object(action, "execute", new=partial(action.execute, args_mock)), + ): + yield action + + +@pytest.mark.integration +def test_pip_delete_demo(pip_delete_demo_action, start_cmd_mock): + pip_delete_demo_action.execute() + + start_cmd_mock.assert_any_call( + TESTGEN_PATH, + "setup-system-db", + "--delete-db", + "--yes", + raise_on_non_zero=True, + env=None, + ) + + +# Marker-driven dispatch tests -------------------------------------------- + + +@pytest.fixture +def run_demo_action(action_cls, args_mock, tmp_data_folder): + action = TestgenRunDemoAction() + args_mock.prod = "tg" + args_mock.action = "run-demo" + args_mock.obs_export = False + action.analytics = MagicMock() + action.analytics.additional_properties = {} + return action + + +@pytest.mark.integration +def test_run_demo_aborts_without_install(run_demo_action, args_mock, console_msg_mock): + with pytest.raises(AbortAction): + run_demo_action._resolve_install_mode(args_mock) + + console_msg_mock.assert_any_msg_contains("tg install") + + +@pytest.mark.integration +@pytest.mark.parametrize("install_mode", [INSTALL_MODE_PIP, INSTALL_MODE_DOCKER]) +def test_run_demo_routes_by_marker(run_demo_action, args_mock, tmp_data_folder, install_mode, start_cmd_mock): + InstallMarker(Path(tmp_data_folder), "tg").write(install_mode) + run_demo_action._resolve_install_mode(args_mock) + + with ( + patch("tests.installer.shutil.which", return_value=UV_PATH), + patch("tests.installer.resolve_testgen_path", return_value=TESTGEN_PATH), + patch.object(run_demo_action, "get_status", return_value={"Status": "running(2)"}), + ): + run_demo_action.execute(args_mock) + + # Inspect the actual quick-start invocation to confirm dispatch. + quick_start_calls = [c for c in start_cmd_mock.call_args_list if "quick-start" in c.args] + assert len(quick_start_calls) == 1 + if install_mode == INSTALL_MODE_PIP: + assert quick_start_calls[0].args[0] == TESTGEN_PATH + else: + assert quick_start_calls[0].args[:2] == ("docker", "compose") + assert "exec" in quick_start_calls[0].args and "engine" in quick_start_calls[0].args + assert run_demo_action.analytics.additional_properties["install_mode"] == install_mode + + +@pytest.fixture +def delete_demo_action(action_cls, args_mock, tmp_data_folder): + action = TestgenDeleteDemoAction() + args_mock.prod = "tg" + args_mock.action = "delete-demo" + action.analytics = MagicMock() + action.analytics.additional_properties = {} + return action + + +@pytest.mark.integration +@pytest.mark.parametrize("install_mode", [INSTALL_MODE_PIP, INSTALL_MODE_DOCKER]) +def test_delete_demo_routes_by_marker(delete_demo_action, args_mock, tmp_data_folder, install_mode, start_cmd_mock): + InstallMarker(Path(tmp_data_folder), "tg").write(install_mode) + + delete_demo_action._resolve_install_mode(args_mock) + with ( + patch("tests.installer.shutil.which", return_value=UV_PATH), + patch("tests.installer.resolve_testgen_path", return_value=TESTGEN_PATH), + patch.object(delete_demo_action, "get_status", return_value={"Status": "running"}), + ): + delete_demo_action.execute(args_mock) + + assert delete_demo_action.analytics.additional_properties["install_mode"] == install_mode + + if install_mode == INSTALL_MODE_PIP: + # Pip branch invokes testgen directly. + start_cmd_mock.assert_any_call( + TESTGEN_PATH, + "setup-system-db", + "--delete-db", + "--yes", + raise_on_non_zero=True, + env=None, + ) + else: + # Docker branch invokes via docker compose exec. + assert any( + "docker" in c.args and "exec" in c.args and "setup-system-db" in c.args + for c in start_cmd_mock.call_args_list + ) diff --git a/tests/test_tg_pip_install.py b/tests/test_tg_pip_install.py new file mode 100644 index 0000000..399026f --- /dev/null +++ b/tests/test_tg_pip_install.py @@ -0,0 +1,310 @@ +import json +from functools import partial +from pathlib import Path +from unittest.mock import MagicMock, call, patch + +import pytest + +from tests.installer import ( + AbortAction, + INSTALL_MODE_DOCKER, + INSTALL_MODE_PIP, + TESTGEN_MAJOR_VERSION, + TESTGEN_PYTHON_VERSION, + TestgenInstallAction, + InstallMarker, +) + + +@pytest.fixture +def pip_install_action(action_cls, args_mock, tmp_data_folder, start_cmd_mock): + """Drive the pip path of TestgenInstallAction end-to-end (steps run).""" + action = TestgenInstallAction() + args_mock.prod = "tg" + args_mock.action = "install" + args_mock.install_mode = INSTALL_MODE_PIP + # Bypass mode resolution: pretend check_requirements already ran. + action._resolved_mode = INSTALL_MODE_PIP + action.steps = action.pip_steps + # uv "found" on PATH so UvBootstrapStep doesn't try to download. + # resolve_testgen_path is mocked to skip the on-disk existence check. + # start_testgen_app is mocked so the test doesn't actually spawn a + # subprocess and block waiting for a port. + with ( + patch("tests.installer.shutil.which", return_value="/usr/local/bin/uv"), + patch( + "tests.installer.resolve_testgen_path", + return_value="/Users/test/.local/bin/testgen", + ), + patch("tests.installer.start_testgen_app"), + patch.object(action, "execute", new=partial(action.execute, args_mock)), + ): + yield action + + +@pytest.fixture +def install_action(action_cls, args_mock, tmp_data_folder): + """A bare TestgenInstallAction for testing the mode-resolution layer.""" + action = TestgenInstallAction() + args_mock.prod = "tg" + args_mock.action = "install" + args_mock.install_mode = None + # Replace the action_cls-shared analytics mock with an instance-level one + # whose additional_properties is a real dict so tests can assert on it. + action.analytics = MagicMock() + action.analytics.additional_properties = {} + return action + + +@pytest.mark.integration +def test_tg_pip_install_happy_path(pip_install_action, start_cmd_mock, tmp_data_folder): + pip_install_action.execute() + + expected_constraint = f"dataops-testgen[standalone]>={TESTGEN_MAJOR_VERSION},<{int(TESTGEN_MAJOR_VERSION) + 1}" + + start_cmd_mock.assert_has_calls( + [ + call( + "/usr/local/bin/uv", + "tool", + "install", + "--force", + "--python", + TESTGEN_PYTHON_VERSION, + expected_constraint, + raise_on_non_zero=True, + env=None, + ), + call( + "/usr/local/bin/uv", + "tool", + "update-shell", + raise_on_non_zero=False, + env=None, + ), + ], + any_order=True, + ) + + cred_path = Path(tmp_data_folder) / "dk-tg-credentials.txt" + assert cred_path.exists() + content = cred_path.read_text() + assert "Username: admin" in content + assert "Password:" in content + assert "http://localhost:8501" in content + + marker_path = Path(tmp_data_folder) / "dk-tg-install.json" + assert marker_path.exists() + data = json.loads(marker_path.read_text()) + assert data["install_mode"] == INSTALL_MODE_PIP + assert "created_on" in data + assert "last_updated_on" in data + + +@pytest.mark.integration +def test_tg_pip_install_threads_user_ports_through_standalone_setup( + pip_install_action, args_mock, start_cmd_mock, tmp_data_folder +): + from tests.installer import TESTGEN_LOG_FILE_PATH + + args_mock.port = 9000 + args_mock.api_port = 9530 + args_mock.ssl_cert_file = "/etc/ssl/certs/tg.crt" + args_mock.ssl_key_file = "/etc/ssl/private/tg.key" + + pip_install_action.execute() + + setup_call = next(c for c in start_cmd_mock.call_args_list if "standalone-setup" in c.args) + assert setup_call.kwargs["env"] == { + "TG_UI_PORT": "9000", + "TG_API_PORT": "9530", + "TESTGEN_LOG_FILE_PATH": str(TESTGEN_LOG_FILE_PATH), + "SSL_CERT_FILE": "/etc/ssl/certs/tg.crt", + "SSL_KEY_FILE": "/etc/ssl/private/tg.key", + } + + cred_path = Path(tmp_data_folder) / "dk-tg-credentials.txt" + content = cred_path.read_text() + + assert "User Interface: https://localhost:9000" in content + assert "API & MCP: https://localhost:9530" in content + + +@pytest.mark.integration +def test_tg_pip_install_auto_starts_app(pip_install_action, args_mock): + """After install completes, the app is auto-started so user has a single-command experience.""" + from tests.installer import start_testgen_app as patched_start + + pip_install_action.execute() + + # The fixture patches tests.installer.start_testgen_app — confirm it was called. + assert patched_start.call_count == 1 + assert patched_start.call_args.args[0] is pip_install_action + assert patched_start.call_args.args[1] is args_mock + + +@pytest.mark.integration +def test_tg_pip_install_auto_runs_demo(pip_install_action, start_cmd_mock): + """A successful pip install also generates demo data so users see something on first launch.""" + pip_install_action.execute() + + start_cmd_mock.assert_any_call( + "/Users/test/.local/bin/testgen", + "quick-start", + raise_on_non_zero=True, + env=None, + ) + + +@pytest.mark.integration +def test_tg_pip_install_no_demo_flag_skips_quick_start(pip_install_action, args_mock, start_cmd_mock): + """--no-demo opts out of the auto-demo step.""" + args_mock.generate_demo = False + + pip_install_action.execute() + + for invocation in start_cmd_mock.call_args_list: + assert "quick-start" not in invocation.args, f"quick-start should be skipped, got: {invocation}" + + +@pytest.mark.integration +def test_tg_pip_install_password_redacted_in_logs(pip_install_action, start_cmd_mock): + """The autogenerated admin password must not appear in the logged cmd_str.""" + pip_install_action.execute() + + # The standalone-setup call should pass redact=(,) so cmd_str gets + # censored. Verify by inspecting the call kwargs. + setup_call = next(c for c in start_cmd_mock.call_args_list if "standalone-setup" in c.args) + redact = setup_call.kwargs.get("redact") + assert redact and len(redact) == 1 + password = redact[0] + assert isinstance(password, str) and len(password) >= 8 + + # The actual --password argument is still present on the command line — it + # only gets censored at log/filename time inside start_cmd. + assert "--password" in setup_call.args + pw_idx = setup_call.args.index("--password") + assert setup_call.args[pw_idx + 1] == password + + +@pytest.mark.integration +@pytest.mark.parametrize("user_input", ["", "p", "pip", "P"]) +def test_auto_mode_picks_pip_when_docker_unavailable(install_action, args_mock, console_msg_mock, user_input): + """Docker probe fails → user accepts the recommended pip default (or types pip explicitly) → resolve to pip.""" + with ( + patch("tests.installer.Requirement.check_availability", return_value=False), + patch("builtins.input", return_value=user_input), + ): + install_action._resolve_install_mode(args_mock) + + assert install_action._resolved_mode == INSTALL_MODE_PIP + assert install_action.analytics.additional_properties["install_mode"] == INSTALL_MODE_PIP + + +@pytest.mark.integration +def test_auto_mode_displays_prereq_status_when_docker_unavailable(install_action, args_mock, console_msg_mock): + """Docker probe fails → the prereq display lists each requirement with a marker and (for failures) a fix hint.""" + # Only the first prereq passes — exercises the mixed pass/fail rendering. + def selective_check(req_self, *_, **__): + return req_self.key == "DOCKER" + + with ( + patch("tests.installer.Requirement.check_availability", autospec=True, side_effect=selective_check), + patch("builtins.input", return_value="p"), + ): + install_action._resolve_install_mode(args_mock) + + console_msg_mock.assert_any_msg_contains("two installation modes") + console_msg_mock.assert_any_msg_contains("Prerequisites:") + console_msg_mock.assert_any_msg_contains("(✓) Docker installed") + console_msg_mock.assert_any_msg_contains("(X) Docker engine running") + + +@pytest.mark.integration +def test_auto_mode_pip_only_prompt_when_docker_unavailable(install_action, args_mock, console_msg_mock): + """When Docker prereqs fail, the prompt only offers pip and tells the user how to retry with Docker.""" + with ( + patch("tests.installer.Requirement.check_availability", return_value=False), + patch("builtins.input", return_value="") as input_mock, + ): + install_action._resolve_install_mode(args_mock) + + assert install_action._resolved_mode == INSTALL_MODE_PIP + # The prompt itself shouldn't offer [d] anymore in this branch. + prompt = input_mock.call_args.args[0] + assert "[d]" not in prompt + assert "[p]" in prompt + console_msg_mock.assert_any_msg_contains("To install with Docker, fix the prerequisites and run the install again.") + + +@pytest.mark.integration +@pytest.mark.parametrize("interrupt", [KeyboardInterrupt, EOFError]) +def test_auto_mode_aborts_on_user_interrupt_when_docker_unavailable(install_action, args_mock, interrupt): + """Docker probe fails → user interrupts the prompt → abort.""" + with ( + patch("tests.installer.Requirement.check_availability", return_value=False), + patch("builtins.input", side_effect=interrupt), + pytest.raises(AbortAction), + ): + install_action._resolve_install_mode(args_mock) + + +@pytest.mark.integration +@pytest.mark.parametrize("interrupt", [KeyboardInterrupt, EOFError]) +def test_auto_mode_prompt_aborts_on_user_interrupt(install_action, args_mock, interrupt): + with ( + patch("tests.installer.Requirement.check_availability", return_value=True), + patch("builtins.input", side_effect=interrupt), + pytest.raises(AbortAction), + ): + install_action._resolve_install_mode(args_mock) + + +@pytest.mark.integration +@pytest.mark.parametrize( + "user_input,expected", + [("", INSTALL_MODE_DOCKER), ("d", INSTALL_MODE_DOCKER), ("p", INSTALL_MODE_PIP), ("pip", INSTALL_MODE_PIP)], +) +def test_auto_mode_prompts_when_docker_available_tty(install_action, args_mock, user_input, expected): + with ( + patch("tests.installer.Requirement.check_availability", return_value=True), + patch("builtins.input", return_value=user_input), + ): + install_action._resolve_install_mode(args_mock) + + assert install_action._resolved_mode == expected + + +@pytest.mark.integration +def test_dispatcher_routes_to_pip_with_flag(install_action, args_mock): + args_mock.install_mode = INSTALL_MODE_PIP + + install_action._resolve_install_mode(args_mock) + + assert install_action._resolved_mode == INSTALL_MODE_PIP + assert install_action.steps == TestgenInstallAction.pip_steps + assert install_action.analytics.additional_properties["install_mode"] == INSTALL_MODE_PIP + + +@pytest.mark.integration +def test_dispatcher_routes_to_docker_with_flag(install_action, args_mock): + args_mock.install_mode = INSTALL_MODE_DOCKER + + install_action._resolve_install_mode(args_mock) + + assert install_action._resolved_mode == INSTALL_MODE_DOCKER + assert install_action.steps == TestgenInstallAction.docker_steps + assert install_action.analytics.additional_properties["install_mode"] == INSTALL_MODE_DOCKER + + +@pytest.mark.integration +@pytest.mark.parametrize("existing", [INSTALL_MODE_DOCKER, INSTALL_MODE_PIP]) +def test_dispatcher_aborts_on_existing_install(install_action, args_mock, tmp_data_folder, existing, console_msg_mock): + InstallMarker(Path(tmp_data_folder), "tg").write(existing) + + with pytest.raises(AbortAction): + install_action._resolve_install_mode(args_mock) + + console_msg_mock.assert_any_msg_contains("tg upgrade") + console_msg_mock.assert_any_msg_contains("tg delete") + console_msg_mock.assert_any_msg_contains(existing) diff --git a/tests/test_tg_pip_upgrade.py b/tests/test_tg_pip_upgrade.py new file mode 100644 index 0000000..aa25dd7 --- /dev/null +++ b/tests/test_tg_pip_upgrade.py @@ -0,0 +1,157 @@ +import json +from functools import partial +from pathlib import Path +from unittest.mock import MagicMock, call, patch + +import pytest + +from tests.installer import ( + AbortAction, + INSTALL_MODE_DOCKER, + INSTALL_MODE_PIP, + TestgenUpgradeAction, + InstallMarker, +) + + +@pytest.fixture +def pip_upgrade_action(action_cls, args_mock, tmp_data_folder, start_cmd_mock): + """Drive the pip path of TestgenUpgradeAction end-to-end.""" + action = TestgenUpgradeAction() + args_mock.prod = "tg" + args_mock.action = "upgrade" + InstallMarker(action.data_folder, args_mock.prod).write(INSTALL_MODE_PIP) + action._resolved_mode = INSTALL_MODE_PIP + action.steps = action.pip_steps + with ( + patch("tests.installer.shutil.which", return_value="/usr/local/bin/uv"), + patch( + "tests.installer.resolve_testgen_path", + return_value="/Users/test/.local/bin/testgen", + ), + patch.object(action, "execute", new=partial(action.execute, args_mock)), + ): + yield action + + +@pytest.mark.integration +def test_tg_pip_upgrade_happy_path(pip_upgrade_action, start_cmd_mock, stdout_mock, tmp_data_folder, console_msg_mock): + # Step pipeline: pre_execute uv tool list → execute uv --version + # (UvBootstrapStep records uv version) → uv tool upgrade → + # testgen upgrade-system-version → on_action_success uv tool list. + stdout_mock.side_effect = [ + ["dataops-testgen v5.10.0", "- testgen"], + ["uv 0.11.7"], + [], + [], + ["dataops-testgen v5.10.0", "- testgen"], + ] + + pip_upgrade_action.execute() + + start_cmd_mock.assert_has_calls( + [ + call( + "/usr/local/bin/uv", + "--no-cache", + "tool", + "upgrade", + "dataops-testgen", + raise_on_non_zero=True, + env=None, + ), + call( + "/Users/test/.local/bin/testgen", + "upgrade-system-version", + raise_on_non_zero=True, + env=None, + ), + ], + any_order=True, + ) + + console_msg_mock.assert_any_msg_contains("Current version: v5.10.0") + console_msg_mock.assert_any_msg_contains("already up-to-date (v5.10.0)") + + marker_path = Path(tmp_data_folder) / "dk-tg-install.json" + data = json.loads(marker_path.read_text()) + assert data["install_mode"] == INSTALL_MODE_PIP + + +@pytest.mark.integration +def test_tg_pip_upgrade_reports_version_change(pip_upgrade_action, start_cmd_mock, stdout_mock, console_msg_mock): + stdout_mock.side_effect = [ + ["dataops-testgen v5.10.0", "- testgen"], + ["uv 0.11.7"], + [], + [], + ["dataops-testgen v5.10.1", "- testgen"], + ] + + pip_upgrade_action.execute() + + console_msg_mock.assert_any_msg_contains("Current version: v5.10.0") + console_msg_mock.assert_any_msg_contains("Updated to v5.10.1") + + +@pytest.mark.integration +def test_tg_pip_upgrade_marker_preserves_created_on(pip_upgrade_action, stdout_mock, tmp_data_folder): + """Upgrading must not reset the original install's created_on timestamp.""" + marker_path = Path(tmp_data_folder) / "dk-tg-install.json" + initial = json.loads(marker_path.read_text()) + original_created_on = initial["created_on"] + # Force a different "now" via filesystem rewrite by burning a tick on disk. + stdout_mock.side_effect = [ + ["dataops-testgen v5.10.0", "- testgen"], + ["uv 0.11.7"], + [], + [], + ["dataops-testgen v5.10.0", "- testgen"], + ] + + pip_upgrade_action.execute() + + after = json.loads(marker_path.read_text()) + assert after["created_on"] == original_created_on + assert "last_updated_on" in after + + +@pytest.fixture +def upgrade_action(action_cls, args_mock, tmp_data_folder): + """A bare TestgenUpgradeAction for testing the mode-resolution layer.""" + action = TestgenUpgradeAction() + args_mock.prod = "tg" + args_mock.action = "upgrade" + action.analytics = MagicMock() + action.analytics.additional_properties = {} + return action + + +@pytest.mark.integration +def test_upgrade_aborts_with_no_install(upgrade_action, args_mock, console_msg_mock): + with pytest.raises(AbortAction): + upgrade_action._resolve_install_mode(args_mock) + + console_msg_mock.assert_any_msg_contains("tg install") + + +@pytest.mark.integration +def test_upgrade_resolves_to_pip_when_marker_says_pip(upgrade_action, args_mock, tmp_data_folder): + InstallMarker(Path(tmp_data_folder), "tg").write(INSTALL_MODE_PIP) + + upgrade_action._resolve_install_mode(args_mock) + + assert upgrade_action._resolved_mode == INSTALL_MODE_PIP + assert upgrade_action.steps == TestgenUpgradeAction.pip_steps + assert upgrade_action.analytics.additional_properties["install_mode"] == INSTALL_MODE_PIP + + +@pytest.mark.integration +def test_upgrade_resolves_to_docker_when_marker_says_docker(upgrade_action, args_mock, tmp_data_folder): + InstallMarker(Path(tmp_data_folder), "tg").write(INSTALL_MODE_DOCKER) + + upgrade_action._resolve_install_mode(args_mock) + + assert upgrade_action._resolved_mode == INSTALL_MODE_DOCKER + assert upgrade_action.steps == TestgenUpgradeAction.docker_steps + assert upgrade_action.analytics.additional_properties["install_mode"] == INSTALL_MODE_DOCKER diff --git a/tests/test_tg_run_demo.py b/tests/test_tg_run_demo.py index 40310bb..459a70e 100644 --- a/tests/test_tg_run_demo.py +++ b/tests/test_tg_run_demo.py @@ -3,7 +3,12 @@ import pytest -from tests.installer import AbortAction, TestgenRunDemoAction +from tests.installer import ( + INSTALL_MODE_DOCKER, + AbortAction, + TestgenRunDemoAction, + InstallMarker, +) @pytest.fixture @@ -11,6 +16,10 @@ def tg_run_demo_action(action_cls, args_mock, tmp_data_folder, start_cmd_mock): action = TestgenRunDemoAction() args_mock.prod = "tg" args_mock.action = "run-demo" + # Seed a Docker install marker so the unified action picks the Docker path. + InstallMarker(action.data_folder, args_mock.prod).write(INSTALL_MODE_DOCKER) + # Bypass check_requirements: pre-resolve mode so execute() runs directly. + action._resolved_mode = INSTALL_MODE_DOCKER with patch.object(action, "execute", new=partial(action.execute, args_mock)): yield action @@ -78,7 +87,7 @@ def test_tg_run_demo_abort_not_running(tg_run_demo_action, start_cmd_mock, conso with pytest.raises(AbortAction): tg_run_demo_action.execute() - console_msg_mock.assert_any_msg_contains("Running the TestGen demo requires the platform to be running.") + console_msg_mock.assert_any_msg_contains("Running the TestGen demo requires the application to be running.") @pytest.mark.integration diff --git a/tests/test_tg_start.py b/tests/test_tg_start.py new file mode 100644 index 0000000..f6d2315 --- /dev/null +++ b/tests/test_tg_start.py @@ -0,0 +1,200 @@ +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +from tests.installer import ( + AbortAction, + INSTALL_MODE_DOCKER, + INSTALL_MODE_PIP, + InstallerError, + TestgenStartAction, + start_testgen_app, + InstallMarker, +) + + +# --- start_testgen_app helper ------------------------------------------------- + + +@pytest.fixture +def app_action(action, tmp_data_folder, tmp_path): + """Bare Action with a real session_folder so file redirection works.""" + action.session_folder = tmp_path / "session" + action.session_folder.mkdir() + action.analytics = MagicMock() + action.analytics.additional_properties = {} + action.ctx = {} + return action + + +@pytest.fixture +def proc_running_then_stops(): + """A Popen-like mock that pretends to be alive, then exits cleanly.""" + proc = MagicMock() + proc.poll.return_value = None # still alive while running + proc.wait.return_value = 0 + return proc + + +@pytest.fixture +def empty_tg_config(monkeypatch): + """Pretend ~/.testgen/config.env doesn't exist so port/SSL fall to defaults.""" + monkeypatch.setattr("tests.installer.read_testgen_config_env", dict) + + +@pytest.mark.unit +def test_start_testgen_app_happy_path(app_action, args_mock, proc_running_then_stops, empty_tg_config): + args_mock.prod = "tg" + + with ( + patch("tests.installer.resolve_testgen_path", return_value="/bin/testgen"), + patch("tests.installer.subprocess.Popen", return_value=proc_running_then_stops) as popen_mock, + patch("tests.installer.wait_for_tcp_port", return_value=True) as port_mock, + ): + start_testgen_app(app_action, args_mock) + + popen_mock.assert_called_once() + invocation = popen_mock.call_args + assert invocation.args[0] == ["/bin/testgen", "run-app"] + # Output is discarded (DEVNULL). + assert "stdout" in invocation.kwargs + port_mock.assert_called_once() + proc_running_then_stops.wait.assert_called() + + +@pytest.mark.unit +def test_start_testgen_app_uses_port_from_config_env(app_action, args_mock, proc_running_then_stops, monkeypatch): + """Port + SSL come from ~/.testgen/config.env (the source of truth post-setup) + — not from args, which doesn't carry these flags on `tg start`.""" + args_mock.prod = "tg" + monkeypatch.setattr( + "tests.installer.read_testgen_config_env", + lambda: {"TG_UI_PORT": "9000", "SSL_CERT_FILE": "/etc/cert", "SSL_KEY_FILE": "/etc/key"}, + ) + + with ( + patch("tests.installer.resolve_testgen_path", return_value="/bin/testgen"), + patch("tests.installer.subprocess.Popen", return_value=proc_running_then_stops), + patch("tests.installer.wait_for_tcp_port", return_value=True) as port_mock, + ): + start_testgen_app(app_action, args_mock) + + # The port we wait for is what config.env says, not args defaults. + port_mock.assert_called_once() + assert port_mock.call_args.args[0] == 9000 + + +@pytest.mark.unit +def test_start_testgen_app_aborts_on_port_timeout(app_action, args_mock, proc_running_then_stops, empty_tg_config): + args_mock.prod = "tg" + + with ( + patch("tests.installer.resolve_testgen_path", return_value="/bin/testgen"), + patch("tests.installer.subprocess.Popen", return_value=proc_running_then_stops), + patch("tests.installer.wait_for_tcp_port", return_value=False), + pytest.raises(InstallerError, match="did not start within"), + ): + start_testgen_app(app_action, args_mock) + + proc_running_then_stops.terminate.assert_called() + + +@pytest.mark.unit +def test_start_testgen_app_handles_keyboard_interrupt(app_action, args_mock, console_msg_mock, empty_tg_config): + """User Ctrl+C during run is the expected stop signal — terminate cleanly, + don't propagate the exception, and hint at the start command for next time.""" + args_mock.prod = "tg" + + proc = MagicMock() + # poll() returns None while alive; after terminate() it transitions to 0. + proc.poll.return_value = None + + def _on_terminate(): + proc.poll.return_value = 0 + + proc.terminate.side_effect = _on_terminate + proc.wait.side_effect = [KeyboardInterrupt(), 0] + + with ( + patch("tests.installer.resolve_testgen_path", return_value="/bin/testgen"), + patch("tests.installer.subprocess.Popen", return_value=proc), + patch("tests.installer.wait_for_tcp_port", return_value=True), + ): + start_testgen_app(app_action, args_mock) + + proc.terminate.assert_called() + console_msg_mock.assert_any_msg_contains("TestGen stopped") + console_msg_mock.assert_any_msg_contains("tg start") + + +# --- TestgenStartAction ------------------------------------------------------- + + +@pytest.fixture +def start_action(action_cls, args_mock, tmp_data_folder): + action = TestgenStartAction() + args_mock.prod = "tg" + args_mock.action = "start" + action.analytics = MagicMock() + action.analytics.additional_properties = {} + return action + + +@pytest.mark.integration +def test_start_action_aborts_with_no_install(start_action, args_mock, console_msg_mock): + with pytest.raises(AbortAction): + start_action._resolve_install_mode(args_mock) + + console_msg_mock.assert_any_msg_contains("No TestGen installation found") + console_msg_mock.assert_any_msg_contains("tg install") + + +@pytest.mark.integration +def test_start_action_runs_compose_up_in_docker_mode( + start_action, args_mock, tmp_data_folder, start_cmd_mock, compose_path +): + InstallMarker(Path(tmp_data_folder), "tg").write(INSTALL_MODE_DOCKER) + + start_action._resolve_install_mode(args_mock) + start_action.execute(args_mock) + + start_cmd_mock.assert_any_call( + "docker", + "compose", + "-f", + compose_path, + "up", + "--wait", + raise_on_non_zero=True, + env=None, + ) + assert start_action.analytics.additional_properties["install_mode"] == INSTALL_MODE_DOCKER + + +@pytest.mark.integration +def test_start_action_routes_to_helper_in_pip_mode(start_action, args_mock, tmp_data_folder): + InstallMarker(Path(tmp_data_folder), "tg").write(INSTALL_MODE_PIP) + + start_action._resolve_install_mode(args_mock) + with patch("tests.installer.start_testgen_app") as start_helper: + start_action.execute(args_mock) + + start_helper.assert_called_once_with(start_action, args_mock) + assert start_action.analytics.additional_properties["install_mode"] == INSTALL_MODE_PIP + + +@pytest.fixture +def console_msg_mock(): + """Local override of the project-wide fixture so this file is self-contained.""" + from tests.installer import CONSOLE + + with patch.object(CONSOLE, "msg") as mock: + + def _assert_any_msg_contains(text: str): + assert any(c for c in mock.call_args_list if text in c.args[0]), ( + f"The text '{text}' wasn't found in any of the {len(mock.call_args_list)} message(s) printed." + ) + + mock.assert_any_msg_contains = _assert_any_msg_contains + yield mock diff --git a/tests/test_tg_upgrade.py b/tests/test_tg_upgrade.py index be94d92..67cae08 100644 --- a/tests/test_tg_upgrade.py +++ b/tests/test_tg_upgrade.py @@ -4,7 +4,14 @@ import pytest -from tests.installer import AbortAction, CommandFailed, TestgenUpgradeAction, TESTGEN_LATEST_TAG +from tests.installer import ( + INSTALL_MODE_DOCKER, + AbortAction, + CommandFailed, + TESTGEN_MAJOR_VERSION, + TestgenUpgradeAction, + InstallMarker, +) @pytest.fixture @@ -12,6 +19,10 @@ def tg_upgrade_action(action_cls, args_mock, tmp_data_folder, start_cmd_mock, re action = TestgenUpgradeAction() args_mock.prod = "tg" args_mock.action = "upgrade" + # Seed a Docker install marker so the unified upgrade resolves to Docker mode. + InstallMarker(action.data_folder, args_mock.prod).write(INSTALL_MODE_DOCKER) + action._resolved_mode = INSTALL_MODE_DOCKER + action.steps = action.docker_steps yield action @@ -69,7 +80,7 @@ def test_tg_upgrade_compose_missing(tg_upgrade_action, args_mock, start_cmd_mock start_cmd_mock.__exit__.side_effect = [None, None, CommandFailed] with pytest.raises(AbortAction, match=""): - tg_upgrade_action._check_requirements(args_mock) + tg_upgrade_action.check_requirements(args_mock) console_msg_mock.assert_any_msg_contains("TestGen's Docker configuration file is not available") @@ -110,7 +121,7 @@ def test_tg_upgrade( compose_content = compose_path.read_text() - assert f"image: datakitchen/dataops-testgen:{TESTGEN_LATEST_TAG}" in compose_content + assert f"image: datakitchen/dataops-testgen:v{TESTGEN_MAJOR_VERSION}" in compose_content assert "TG_INSTANCE_ID:" in compose_content diff --git a/tests/test_uv_bootstrap.py b/tests/test_uv_bootstrap.py new file mode 100644 index 0000000..6b4adb1 --- /dev/null +++ b/tests/test_uv_bootstrap.py @@ -0,0 +1,253 @@ +import hashlib +import io +import tarfile +import zipfile +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +from tests.installer import ( + AbortAction, + InstallerError, + UV_ASSETS, + UV_VERSION, + UvBootstrapStep, + get_uv_asset, +) + + +def _make_tarball(inner_path: str, payload: bytes) -> bytes: + buf = io.BytesIO() + with tarfile.open(fileobj=buf, mode="w:gz") as tf: + info = tarfile.TarInfo(inner_path) + info.size = len(payload) + tf.addfile(info, io.BytesIO(payload)) + return buf.getvalue() + + +def _make_zip(inner_path: str, payload: bytes) -> bytes: + buf = io.BytesIO() + with zipfile.ZipFile(buf, "w") as zf: + zf.writestr(inner_path, payload) + return buf.getvalue() + + +@pytest.fixture +def uv_step(): + return UvBootstrapStep() + + +@pytest.fixture +def bootstrap_action(action, tmp_data_folder): + action.ctx = {} + action.analytics = MagicMock() + action.analytics.additional_properties = {} + return action + + +@pytest.fixture +def linux_x86_64(monkeypatch): + monkeypatch.setattr("tests.installer.platform.system", lambda: "Linux") + monkeypatch.setattr("tests.installer.platform.machine", lambda: "x86_64") + + +@pytest.fixture +def windows_amd64(monkeypatch): + monkeypatch.setattr("tests.installer.platform.system", lambda: "Windows") + monkeypatch.setattr("tests.installer.platform.machine", lambda: "AMD64") + + +@pytest.mark.unit +def test_get_uv_asset_known_platform(monkeypatch): + monkeypatch.setattr("tests.installer.platform.system", lambda: "Darwin") + monkeypatch.setattr("tests.installer.platform.machine", lambda: "arm64") + + asset, sha256 = get_uv_asset("tg") + + assert asset == "uv-aarch64-apple-darwin.tar.gz" + assert sha256 == UV_ASSETS[("Darwin", "arm64")][1] + + +@pytest.mark.unit +def test_get_uv_asset_unknown_platform_raises(monkeypatch, console_msg_mock): + monkeypatch.setattr("tests.installer.platform.system", lambda: "SunOS") + monkeypatch.setattr("tests.installer.platform.machine", lambda: "sparc64") + + with pytest.raises(AbortAction): + get_uv_asset("tg") + + console_msg_mock.assert_any_msg_contains("No prebuilt uv binary available for platform SunOS/sparc64") + + +@pytest.mark.unit +def test_skips_when_uv_on_path(uv_step, bootstrap_action, args_mock, linux_x86_64): + # `uv --version` returns a different version than UV_VERSION on purpose: + # the existing-uv branch must record the *actual* installed version, not + # what we'd ship if we'd downloaded. + with ( + patch("tests.installer.shutil.which", return_value="/usr/local/bin/uv"), + patch("tests.installer.urllib.request.urlopen") as urlopen_mock, + patch.object(bootstrap_action, "run_cmd", return_value="uv 0.4.30 (deadbeef 2024-08-01)"), + ): + # MultiStepAction always invokes pre_execute before execute; mirror that. + uv_step.pre_execute(bootstrap_action, args_mock) + uv_step.execute(bootstrap_action, args_mock) + + urlopen_mock.assert_not_called() + assert bootstrap_action.ctx["uv_path"] == "/usr/local/bin/uv" + assert bootstrap_action.analytics.additional_properties["uv_source"] == "existing" + assert bootstrap_action.analytics.additional_properties["uv_version"] == "0.4.30" + + +@pytest.mark.unit +def test_skips_when_uv_local(uv_step, bootstrap_action, args_mock, linux_x86_64): + local_uv = Path(bootstrap_action.data_folder) / "bin" / "uv" + local_uv.parent.mkdir(parents=True, exist_ok=True) + local_uv.write_bytes(b"existing-uv") + + with ( + patch("tests.installer.shutil.which", return_value=None), + patch("tests.installer.urllib.request.urlopen") as urlopen_mock, + patch.object(bootstrap_action, "run_cmd", return_value=f"uv {UV_VERSION} (abcd 2024-09-15)"), + ): + uv_step.pre_execute(bootstrap_action, args_mock) + uv_step.execute(bootstrap_action, args_mock) + + urlopen_mock.assert_not_called() + assert bootstrap_action.ctx["uv_path"] == str(local_uv) + assert bootstrap_action.analytics.additional_properties["uv_source"] == "existing" + assert bootstrap_action.analytics.additional_properties["uv_version"] == UV_VERSION + + +@pytest.mark.unit +def test_downloads_and_extracts_on_linux(uv_step, bootstrap_action, args_mock, linux_x86_64, monkeypatch): + payload = b"fake-uv-binary-linux" + tarball = _make_tarball("uv-x86_64-unknown-linux-gnu/uv", payload) + # Include uvx too to ensure the step picks the right binary. + buf = io.BytesIO() + with tarfile.open(fileobj=buf, mode="w:gz") as tf: + for inner, data in ( + ("uv-x86_64-unknown-linux-gnu/uv", payload), + ("uv-x86_64-unknown-linux-gnu/uvx", b"fake-uvx"), + ): + info = tarfile.TarInfo(inner) + info.size = len(data) + tf.addfile(info, io.BytesIO(data)) + tarball = buf.getvalue() + + sha256 = hashlib.sha256(tarball).hexdigest() + monkeypatch.setitem(UV_ASSETS, ("Linux", "x86_64"), ("uv-x86_64-unknown-linux-gnu.tar.gz", sha256)) + + resp = MagicMock() + resp.read.return_value = tarball + resp.__enter__.return_value = resp + resp.__exit__.return_value = False + + with ( + patch("tests.installer.shutil.which", return_value=None), + patch("tests.installer.urllib.request.urlopen", return_value=resp) as urlopen_mock, + patch.object(bootstrap_action, "run_cmd", return_value=f"uv {UV_VERSION} (abcd 2024-09-15)"), + ): + uv_step.execute(bootstrap_action, args_mock) + + urlopen_mock.assert_called_once() + called_url = urlopen_mock.call_args.args[0] + assert UV_VERSION in called_url + assert called_url.endswith("uv-x86_64-unknown-linux-gnu.tar.gz") + + installed = Path(bootstrap_action.data_folder) / "bin" / "uv" + assert installed.exists() + assert installed.read_bytes() == payload + # Archive cleaned up + assert not (Path(bootstrap_action.data_folder) / "bin" / "uv-x86_64-unknown-linux-gnu.tar.gz").exists() + + assert bootstrap_action.ctx["uv_path"] == str(installed) + assert bootstrap_action.analytics.additional_properties["uv_source"] == "download" + # uv_version comes from parsing `uv --version`, not from the UV_VERSION constant. + assert bootstrap_action.analytics.additional_properties["uv_version"] == UV_VERSION + + +@pytest.mark.unit +def test_downloads_and_extracts_on_windows(uv_step, bootstrap_action, args_mock, windows_amd64, monkeypatch): + payload = b"fake-uv-binary-windows" + zip_bytes = _make_zip("uv-x86_64-pc-windows-msvc/uv.exe", payload) + sha256 = hashlib.sha256(zip_bytes).hexdigest() + monkeypatch.setitem(UV_ASSETS, ("Windows", "AMD64"), ("uv-x86_64-pc-windows-msvc.zip", sha256)) + + resp = MagicMock() + resp.read.return_value = zip_bytes + resp.__enter__.return_value = resp + resp.__exit__.return_value = False + + with ( + patch("tests.installer.shutil.which", return_value=None), + patch("tests.installer.urllib.request.urlopen", return_value=resp), + patch.object(bootstrap_action, "run_cmd", return_value=f"uv {UV_VERSION} (abcd 2024-09-15)"), + ): + uv_step.execute(bootstrap_action, args_mock) + + installed = Path(bootstrap_action.data_folder) / "bin" / "uv.exe" + assert installed.exists() + assert installed.read_bytes() == payload + + +@pytest.mark.unit +def test_sha256_mismatch_fails_fast_without_retry(uv_step, bootstrap_action, args_mock, linux_x86_64, monkeypatch): + """SHA256 mismatch is deterministic — a corp proxy serving the wrong file + won't fix itself on retry. Fail fast so the user sees the real error.""" + wrong_bytes = _make_tarball("uv-x86_64-unknown-linux-gnu/uv", b"garbage") + # UV_ASSETS has the real hash, so the wrong-bytes payload won't match. + + resp = MagicMock() + resp.read.return_value = wrong_bytes + resp.__enter__.return_value = resp + resp.__exit__.return_value = False + + with ( + patch("tests.installer.shutil.which", return_value=None), + patch("tests.installer.urllib.request.urlopen", return_value=resp) as urlopen_mock, + patch.object(bootstrap_action, "run_cmd"), + pytest.raises(InstallerError, match="SHA256 mismatch"), + ): + uv_step.execute(bootstrap_action, args_mock) + + assert urlopen_mock.call_count == 1 + assert not (Path(bootstrap_action.data_folder) / "bin" / "uv").exists() + + +@pytest.mark.unit +def test_download_failure_retries_then_fails(uv_step, bootstrap_action, args_mock, linux_x86_64): + from tests.installer import UV_DOWNLOAD_RETRIES + + with ( + patch("tests.installer.shutil.which", return_value=None), + patch( + "tests.installer.urllib.request.urlopen", + side_effect=OSError("network down"), + ) as urlopen_mock, + patch.object(bootstrap_action, "run_cmd"), + # Skip the exponential backoff sleeps so the test runs instantly. + patch("tests.installer.time.sleep") as sleep_mock, + pytest.raises(InstallerError, match="Failed to bootstrap uv"), + ): + uv_step.execute(bootstrap_action, args_mock) + + assert urlopen_mock.call_count == UV_DOWNLOAD_RETRIES + # Backoff sleeps fire between retries, not after the final failure. + assert sleep_mock.call_count == UV_DOWNLOAD_RETRIES - 1 + + +@pytest.mark.unit +def test_unsupported_platform_raises(uv_step, bootstrap_action, args_mock, monkeypatch, console_msg_mock): + monkeypatch.setattr("tests.installer.platform.system", lambda: "SunOS") + monkeypatch.setattr("tests.installer.platform.machine", lambda: "sparc64") + + with ( + patch("tests.installer.shutil.which", return_value=None), + patch.object(bootstrap_action, "run_cmd"), + pytest.raises(AbortAction), + ): + uv_step.execute(bootstrap_action, args_mock) + + console_msg_mock.assert_any_msg_contains("No prebuilt uv binary available")