From eb7a60506aee1b780a08d75bb0b9f5452807c3fe Mon Sep 17 00:00:00 2001 From: Aarthy Adityan Date: Wed, 22 Apr 2026 16:19:23 -0400 Subject: [PATCH 1/6] feat(testgen): support pip install --- .gitignore | 8 + CLAUDE.md | 101 +++ README.md | 48 +- dk-installer.py | 1252 ++++++++++++++++++++++++++++++---- tests/conftest.py | 7 +- tests/test_install_type.py | 109 +++ tests/test_tg_install.py | 11 + tests/test_tg_pip_delete.py | 188 +++++ tests/test_tg_pip_demo.py | 247 +++++++ tests/test_tg_pip_install.py | 292 ++++++++ tests/test_tg_pip_upgrade.py | 157 +++++ tests/test_tg_run_demo.py | 13 +- tests/test_tg_start.py | 200 ++++++ tests/test_tg_upgrade.py | 17 +- tests/test_uv_bootstrap.py | 248 +++++++ 15 files changed, 2757 insertions(+), 141 deletions(-) create mode 100644 CLAUDE.md create mode 100644 tests/test_install_type.py create mode 100644 tests/test_tg_pip_delete.py create mode 100644 tests/test_tg_pip_demo.py create mode 100644 tests/test_tg_pip_install.py create mode 100644 tests/test_tg_pip_upgrade.py create mode 100644 tests/test_tg_start.py create mode 100644 tests/test_uv_bootstrap.py 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..6a0e7be 100644 --- a/README.md +++ b/README.md @@ -47,12 +47,19 @@ 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` | +- **pip** (default, recommended for evaluation) — no Docker required. The installer downloads `uv`, installs Python 3.13 if needed, and installs TestGen with an embedded Postgres database. +- **Docker** — deploys TestGen as a Docker Compose application. Use for team eval on a shared VM, or if you already standardize on Docker. + +Observability is always installed via Docker Compose. + +| Install path | Required software | +|---|---| +| TestGen (pip, default) | [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/) 26+, [Docker Compose](https://docs.docker.com/compose/install/) 2.38+. | + +Check versions with `python3 --version`, `docker -v`, `docker compose version`. ### Download the installer @@ -75,18 +82,27 @@ The [Data Observability quickstart](https://docs.datakitchen.io/tutorials/quicks Before going through the quickstart, complete the prequisites above and then the following steps to install the two products and setup the demo data. For any of the commands, you can view additional options by appending `--help` at the end. -### Install the TestGen application +### Install the TestGen application (pip, default) -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. +`tg install` defaults to a pip-based install with an embedded Postgres database and no Docker requirement. The installer downloads `uv` (if it isn't already on your PATH), uses it to install Python 3.13 (if needed) and TestGen in an isolated environment, then prints credentials plus the command to start the app. ```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. +The process typically takes 2-5 minutes. On completion the installer writes credentials to `dk-tg-credentials.txt` and prints the `testgen run-app` command to start the UI in a separate terminal. + +#### Install the TestGen application (Docker) + +If you prefer the Docker Compose install — for team evaluations on a shared VM, or if you already standardize on Docker — use the `--docker` flag: + +```shell +python3 dk-installer.py tg install --docker +``` -Once the installation completes, verify that you can login to the UI with the URL and credentials provided in the output. +The Docker install takes 5-10 minutes. `--port` sets a custom localhost port (default 8501). `--ssl-cert-file` / `--ssl-key-file` enable HTTPS. + +Either install path can later be upgraded with `python3 dk-installer.py tg upgrade` — the installer detects which flavor is present and upgrades accordingly. ### Install the Observability application @@ -136,7 +152,17 @@ Leave this process running, and continue with the [quickstart guide](https://doc ## Useful Commands -### DataOps TestGen +### DataOps TestGen (pip install, default) + +Start the app: `TG_STANDALONE_MODE=yes uv tool run testgen run-app` (reachable at `http://localhost:8501`) + +Stop the app: `Ctrl+C` in the terminal running `testgen run-app` + +Access the `testgen` CLI: `TG_STANDALONE_MODE=yes uv tool run testgen ` + +Upgrade the app to latest version: `python3 dk-installer.py tg upgrade` + +### DataOps TestGen (Docker install) 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. diff --git a/dk-installer.py b/dk-installer.py index ce14beb..51cf421 100755 --- a/dk-installer.py +++ b/dk-installer.py @@ -19,10 +19,14 @@ 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 @@ -46,8 +50,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 +60,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 +133,16 @@ 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 collect_images_digest(action, images, env=None): @@ -132,9 +182,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 +226,93 @@ 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}`" + + +def read_install_mode(data_folder: pathlib.Path, prod: str, compose_file_name: str) -> typing.Optional[str]: + """Return 'docker', 'pip', or None for the ``prod`` install in data_folder. + + Reads the install marker file. Falls back to detecting a legacy Docker install + (the product's compose file + credentials) from before the marker was introduced. + """ + marker_path = data_folder / INSTALL_MARKER_FILE.format(prod) + if marker_path.exists(): + try: + data = json.loads(marker_path.read_text()) + except Exception: + LOG.exception("Failed to read install marker at %s", marker_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 (data_folder / compose_file_name).exists() and (data_folder / CREDENTIALS_FILE.format(prod)).exists(): + LOG.info("No marker present; detected legacy Docker install in %s", data_folder) + return INSTALL_MODE_DOCKER + + return None + + +def write_install_marker(data_folder: pathlib.Path, prod: str, install_mode: str, **extra) -> None: + if install_mode not in (INSTALL_MODE_DOCKER, INSTALL_MODE_PIP): + raise ValueError(f"Unknown install_mode: {install_mode}") + marker_path = data_folder / INSTALL_MARKER_FILE.format(prod) + now = datetime.datetime.now(datetime.timezone.utc).isoformat() + created_on = now + if marker_path.exists(): + try: + existing = json.loads(marker_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", marker_path) + data = { + "install_mode": install_mode, + "created_on": created_on, + "last_updated_on": now, + **extra, + } + marker_path.write_text(json.dumps(data, indent=2)) + + @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 +378,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 @@ -299,7 +460,7 @@ class Requirement: cmd: tuple[typing.Union[str, pathlib.Path], ...] fail_msg: tuple[str, ...] - 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 +468,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 +713,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 +736,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 +778,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 +791,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 +866,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 +892,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 +915,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 +958,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: @@ -976,6 +1167,146 @@ def run(self, parent=None): ) +def get_uv_asset(prod: str) -> tuple[str, str]: + """Return (asset_name, sha256) for the current platform, or raise InstallerError.""" + key = (platform.system(), platform.machine()) + try: + return UV_ASSETS[key] + except KeyError: + supported = ", ".join(f"{s}/{m}" for s, m in UV_ASSETS) + raise InstallerError( + f"No prebuilt uv binary available for platform {key[0]}/{key[1]}. " + f"Supported: {supported}. " + f"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." + ) + + +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 = "Setting up 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 +1339,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 +1358,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 +1370,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 +1385,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 +1426,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 +1592,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 +1706,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 +1966,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 +1990,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 +2003,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): @@ -1765,7 +2101,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 +2170,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 +2229,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,7 +2304,7 @@ 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( @@ -1985,7 +2321,7 @@ def execute(self, action, args): 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" @@ -2022,8 +2358,289 @@ def on_action_success(self, action, args): 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 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: + raise InstallerError("uv not found.") + 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(): + raise InstallerError( + f"testgen script not found at {testgen_path}. " + f"Try {command_hint(args.prod, 'delete', 'Uninstall TestGen')}, " + f"then {command_hint(args.prod, 'install', 'Install TestGen')}." + ) + 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}.") + 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 + self.testgen_path = None + + def pre_execute(self, action, args): + self.username = DEFAULT_USER_DATA["username"] + self.password = generate_password() + + def execute(self, action, args): + self.testgen_path = resolve_testgen_path(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 + action.run_cmd( + self.testgen_path, + "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. Non-blocking — failure here logs and continues; the user can + run ``tg run-demo`` later if they want to retry. + """ + + label = "Generating demo data" + required = False + + def execute(self, action, args): + if getattr(args, "no_demo", False): + raise SkipStep + testgen_path = resolve_testgen_path(action, args) + action.run_cmd(testgen_path, "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, @@ -2032,16 +2649,40 @@ class TestgenInstallAction(ComposeActionMixin, AnalyticsMultiStepAction): TestGenSetupDatabaseStep, TestGenUpgradeDatabaseStep, ] + pip_intro = [ + "Installing TestGen with pip.", + "The process may take 2~5 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,6 +2697,28 @@ 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.", + ) + # Pip-only args + parser.add_argument( + "--no-demo", + dest="no_demo", + action="store_true", + help="(Pip mode only) Skip generating demo data after install.", + ) + # Docker-only args parser.add_argument( "--image", dest="image", @@ -2073,25 +2736,117 @@ def get_parser(self, sub_parsers): "Defaults to '%(default)s'" ), ) - parser.add_argument( - "--ssl-cert-file", - dest="ssl_cert_file", - action="store", - default=None, - help="Path to SSL certificate 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 = read_install_mode(self.data_folder, args.prod, args.compose_file_name) + 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 with the same requirement list a real Docker install would check + # (including REQ_TESTGEN_IMAGE, since some networks may block docker.io image pulls). + if not all(req.check_availability(self, args, quiet=True) for req in self.docker_requirements): + CONSOLE.space() + CONSOLE.msg("Docker is not fully available on this machine.") + CONSOLE.msg("TestGen can be installed with pip instead, which uses an embedded Postgres database.") + CONSOLE.space() + try: + choice = input(f"{CONSOLE.MARGIN}Install TestGen with pip? [Y/n]: ").strip().lower() + except (KeyboardInterrupt, EOFError): + print("") + raise AbortAction + CONSOLE.space() + if choice in ("", "y", "yes"): + return INSTALL_MODE_PIP + if getattr(sys, "frozen", False): + CONSOLE.msg("Aborted. Fix the Docker prerequisites and select 'Install TestGen' from the menu again.") + else: + CONSOLE.msg( + f"Aborted. To retry with Docker, fix the prerequisites and run " + f"`python3 {INSTALLER_NAME} {args.prod} install --docker`." + ) + CONSOLE.msg( + f"To install with pip explicitly, run `python3 {INSTALLER_NAME} {args.prod} install --pip`." + ) + raise AbortAction + + CONSOLE.space() + CONSOLE.msg("Two installation modes are available:") + CONSOLE.space() + CONSOLE.msg("[d] Docker Compose (Recommended)") + CONSOLE.msg( + "The most stable TestGen experience for persistent use. Provides a fully managed " + "environment with an isolated PostgreSQL container." ) - parser.add_argument( - "--ssl-key-file", - dest="ssl_key_file", - action="store", - default=None, - help="Path to SSL key file.", + CONSOLE.space() + CONSOLE.msg("[p] Pip + embedded PostgreSQL") + CONSOLE.msg( + "A light-weight Python installation suited for evaluation. Manages the PostgreSQL " + "database on the file system." ) - return parser + CONSOLE.space() + while True: + try: + choice = input(f"{CONSOLE.MARGIN}Install with Docker [d] or pip [p]? (default: d): ").strip().lower() + except (KeyboardInterrupt, EOFError): + print("") + raise AbortAction + if 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) + write_install_marker(self.data_folder, args.prod, 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. + if self._resolved_mode == INSTALL_MODE_PIP: + start_testgen_app(self, args) + + +class TestgenStandaloneUpgradeStep(Step): + label = "Upgrading the application database" + + def execute(self, action, args): + testgen_path = resolve_testgen_path(action, args) + action.run_cmd(testgen_path, "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 +2855,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="Whether to skip the version check before upgrading.", + ) + 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'" + ), + ) + 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 +2903,186 @@ 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 = read_install_mode(self.data_folder, args.prod, args.compose_file_name) + 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) + write_install_marker(self.data_folder, args.prod, 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 = read_install_mode(self.data_folder, args.prod, args.compose_file_name) + 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") + self.run_cmd("docker", "compose", "-f", self.get_compose_file_path(args), "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')}." + ) + 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 = read_install_mode(self.data_folder, args.prod, args.compose_file_name) + 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) + remove_path(self.data_folder / INSTALL_MARKER_FILE.format(args.prod)) + + 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,46 +3095,70 @@ 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 = read_install_mode(self.data_folder, args.prod, args.compose_file_name) + 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() if args.obs_export: self.run_dk_demo_container("tg-run-demo") - quick_start_command = [ - "testgen", - "quick-start", - ] + export_args = [] if args.obs_export: with open(self.data_folder / DEMO_CONFIG_FILE, "r") as file: json_config = json.load(file) + export_args = [ + "--observability-api-url", + json_config["api_host"], + "--observability-api-key", + json_config["api_key"], + ] + + if self._resolved_mode == INSTALL_MODE_DOCKER: + self._run_docker_demo(args, export_args) + else: + self._run_pip_demo(args, export_args) - quick_start_command.extend( - [ - "--observability-api-url", - json_config["api_host"], - "--observability-api-key", - json_config["api_key"], - ] - ) + CONSOLE.title("Demo SUCCEEDED") - cli_commands = [ - quick_start_command, - ] + def _run_docker_demo(self, args, export_args): + cli_commands = [["testgen", "quick-start", *export_args]] if args.obs_export: cli_commands.append( [ @@ -2211,7 +3170,6 @@ def execute(self, args): "default-suite-1", ] ) - for command in cli_commands: CONSOLE.msg(f"Running command : docker compose exec engine {' '.join(command)}") self.run_cmd( @@ -2224,13 +3182,54 @@ def execute(self, args): *command, ) - CONSOLE.title("Demo SUCCEEDED") + def _run_pip_demo(self, args, export_args): + if resolve_uv_path(self.data_folder) is None: + CONSOLE.msg(f"uv not found. To install TestGen, {command_hint(args.prod, 'install', 'Install TestGen')}.") + raise AbortAction + + testgen_path = resolve_testgen_path(self, args) + self.run_cmd(testgen_path, "quick-start", *export_args) + if args.obs_export: + self.run_cmd( + testgen_path, + "export-observability", + "--project-key", + "DEFAULT", + "--test-suite-key", + "default-suite-1", + ) 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 = read_install_mode(self.data_folder, args.prod, args.compose_file_name) + 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") @@ -2238,8 +3237,11 @@ def execute(self, args): pass CONSOLE.msg("Cleaning up system database..") - tg_status = self.get_status(args) - if tg_status: + 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 self.run_cmd( "docker", "compose", @@ -2252,11 +3254,14 @@ def execute(self, args): "--delete-db", "--yes", ) - - CONSOLE.title("Demo data DELETED") else: - CONSOLE.msg("TestGen must be running for its demo data to be cleaned.") - raise AbortAction + if resolve_uv_path(self.data_folder) is None: + CONSOLE.msg("uv not found. Cannot clean TestGen standalone database.") + raise AbortAction + testgen_path = resolve_testgen_path(self, args) + self.run_cmd(testgen_path, "setup-system-db", "--delete-db", "--yes") + + CONSOLE.title("Demo data DELETED") class AccessInstructionsAction(Action): @@ -2297,6 +3302,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 +3367,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 +3387,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..5b989d1 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -134,6 +134,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 +174,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.no_demo = False + 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..e5fa919 --- /dev/null +++ b/tests/test_install_type.py @@ -0,0 +1,109 @@ +import json +from pathlib import Path + +import pytest + +from tests.installer import ( + CREDENTIALS_FILE, + INSTALL_MARKER_FILE, + INSTALL_MODE_DOCKER, + INSTALL_MODE_PIP, + TESTGEN_COMPOSE_FILE, + read_install_mode, + write_install_marker, +) + + +@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 read_install_mode(data_folder, "tg", "docker-compose.yml") 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 read_install_mode(data_folder, "tg", "docker-compose.yml") == 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 read_install_mode(data_folder, "tg", "docker-compose.yml") == INSTALL_MODE_DOCKER + + +@pytest.mark.unit +def test_read_install_mode_legacy_backfill_honors_compose_file_name(data_folder): + # Verifies the function 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 read_install_mode(data_folder, "obs", "obs-docker-compose.yml") == INSTALL_MODE_DOCKER + # And does NOT match if we point at the wrong compose file name. + assert read_install_mode(data_folder, "obs", "docker-compose.yml") 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 read_install_mode(data_folder, "tg", "docker-compose.yml") 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 read_install_mode(data_folder, "tg", "docker-compose.yml") == 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 read_install_mode(data_folder, "tg", "docker-compose.yml") is None + + +@pytest.mark.unit +def test_write_install_marker_round_trip(data_folder): + write_install_marker(data_folder, "tg", 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 read_install_mode(data_folder, "tg", "docker-compose.yml") == INSTALL_MODE_PIP + + +@pytest.mark.unit +def test_write_install_marker_preserves_created_on_across_writes(data_folder): + write_install_marker(data_folder, "tg", INSTALL_MODE_PIP, version="5.9.4") + initial = json.loads((data_folder / INSTALL_MARKER_FILE.format("tg")).read_text()) + + write_install_marker(data_folder, "tg", 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"): + write_install_marker(data_folder, "tg", "sideways") + + diff --git a/tests/test_tg_install.py b/tests/test_tg_install.py index 8922ec4..645f83f 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 diff --git a/tests/test_tg_pip_delete.py b/tests/test_tg_pip_delete.py new file mode 100644 index 0000000..84b86b0 --- /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, + write_install_marker, +) + + +@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" + write_install_marker(action.data_folder, args_mock.prod, 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): + write_install_marker(Path(tmp_data_folder), "tg", 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..ba63c65 --- /dev/null +++ b/tests/test_tg_pip_demo.py @@ -0,0 +1,247 @@ +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, + write_install_marker, +) + + +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" + write_install_marker(action.data_folder, args_mock.prod, 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 + write_install_marker(action.data_folder, args_mock.prod, 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" + write_install_marker(action.data_folder, args_mock.prod, 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): + write_install_marker(Path(tmp_data_folder), "tg", install_mode) + + run_demo_action._resolve_install_mode(args_mock) + with ( + patch.object(run_demo_action, "_run_pip_demo") as pip_branch, + patch.object(run_demo_action, "_run_docker_demo") as docker_branch, + patch.object(run_demo_action, "get_status", return_value={"Status": "running(2)"}), + ): + run_demo_action.execute(args_mock) + + if install_mode == INSTALL_MODE_PIP: + pip_branch.assert_called_once() + docker_branch.assert_not_called() + else: + docker_branch.assert_called_once() + pip_branch.assert_not_called() + 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): + write_install_marker(Path(tmp_data_folder), "tg", 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..591974f --- /dev/null +++ b/tests/test_tg_pip_install.py @@ -0,0 +1,292 @@ +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, + write_install_marker, +) + + +@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.no_demo = True + + 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", ["", "y", "yes", "Y"]) +def test_auto_mode_picks_pip_when_user_confirms_pip_fallback( + install_action, args_mock, console_msg_mock, user_input +): + """Docker probe fails → user confirms pip → 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 + console_msg_mock.assert_any_msg_contains("Docker is not fully available") + + +@pytest.mark.integration +@pytest.mark.parametrize("user_input", ["n", "no", "N", "anything-else"]) +def test_auto_mode_aborts_when_user_declines_pip_fallback(install_action, args_mock, console_msg_mock, user_input): + """Docker probe fails → user declines pip → abort with hint.""" + with ( + patch("tests.installer.Requirement.check_availability", return_value=False), + patch("builtins.input", return_value=user_input), + pytest.raises(AbortAction), + ): + install_action._resolve_install_mode(args_mock) + + console_msg_mock.assert_any_msg_contains("Aborted") + console_msg_mock.assert_any_msg_contains("install --docker") + + +@pytest.mark.integration +@pytest.mark.parametrize("interrupt", [KeyboardInterrupt, EOFError]) +def test_auto_mode_pip_fallback_prompt_aborts_on_user_interrupt(install_action, args_mock, interrupt): + """Docker probe fails → user interrupts the pip-fallback 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): + write_install_marker(Path(tmp_data_folder), "tg", 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..f90b06d --- /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, + write_install_marker, +) + + +@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" + write_install_marker(action.data_folder, args_mock.prod, 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): + write_install_marker(Path(tmp_data_folder), "tg", 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): + write_install_marker(Path(tmp_data_folder), "tg", 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..a507ebd 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, + write_install_marker, +) @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. + write_install_marker(action.data_folder, args_mock.prod, 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..da33bb7 --- /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, + write_install_marker, +) + + +# --- 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 +): + write_install_marker(Path(tmp_data_folder), "tg", 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): + write_install_marker(Path(tmp_data_folder), "tg", 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..e491e7d 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, + write_install_marker, +) @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. + write_install_marker(action.data_folder, args_mock.prod, 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..0f88fa5 --- /dev/null +++ b/tests/test_uv_bootstrap.py @@ -0,0 +1,248 @@ +import hashlib +import io +import tarfile +import zipfile +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +from tests.installer import ( + 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): + monkeypatch.setattr("tests.installer.platform.system", lambda: "SunOS") + monkeypatch.setattr("tests.installer.platform.machine", lambda: "sparc64") + + with pytest.raises(InstallerError, match="No prebuilt uv binary available for platform SunOS/sparc64"): + get_uv_asset("tg") + + +@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): + 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(InstallerError, match="No prebuilt uv binary available"), + ): + uv_step.execute(bootstrap_action, args_mock) From 09319b4d930ffd7a8c9ac52f03fb422cf05a3a99 Mon Sep 17 00:00:00 2001 From: Aarthy Adityan Date: Tue, 5 May 2026 14:33:00 -0400 Subject: [PATCH 2/6] fix(testgen): improve UX and refactor to reuse code --- dk-installer.py | 251 +++++++++++++---------------------- tests/conftest.py | 2 +- tests/test_tg_install.py | 30 +++++ tests/test_tg_pip_demo.py | 18 +-- tests/test_tg_pip_install.py | 50 ++++--- 5 files changed, 165 insertions(+), 186 deletions(-) diff --git a/dk-installer.py b/dk-installer.py index 51cf421..42a0547 100755 --- a/dk-installer.py +++ b/dk-installer.py @@ -459,6 +459,7 @@ 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, quiet=False): try: @@ -1151,11 +1152,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", @@ -1164,6 +1167,7 @@ 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", ) @@ -2016,17 +2020,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) @@ -2307,17 +2301,7 @@ class TestGenSetupDatabaseStep(Step): 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): @@ -2329,30 +2313,10 @@ 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)}") @@ -2400,6 +2364,29 @@ def read_installed_testgen_version(action) -> typing.Optional[str]: 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 @@ -2570,14 +2557,12 @@ class TestgenStandaloneSetupStep(Step): def __init__(self): self.username = None self.password = None - self.testgen_path = None def pre_execute(self, action, args): self.username = DEFAULT_USER_DATA["username"] self.password = generate_password() def execute(self, action, args): - self.testgen_path = resolve_testgen_path(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. @@ -2590,8 +2575,9 @@ def execute(self, action, args): env["SSL_CERT_FILE"] = args.ssl_cert_file if args.ssl_key_file: env["SSL_KEY_FILE"] = args.ssl_key_file - action.run_cmd( - self.testgen_path, + run_testgen_cli( + action, + args, "standalone-setup", "--username", self.username, @@ -2617,18 +2603,17 @@ def on_action_success(self, action, args): class TestgenQuickStartStep(Step): """Generate demo data so the user has something to look at right after - install. Non-blocking — failure here logs and continues; the user can - run ``tg run-demo`` later if they want to retry. + 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 getattr(args, "no_demo", False): + if not getattr(args, "generate_demo", True): raise SkipStep - testgen_path = resolve_testgen_path(action, args) - action.run_cmd(testgen_path, "quick-start") + run_testgen_cli(action, args, "quick-start") class TestgenInstallAction(ComposeActionMixin, AnalyticsMultiStepAction): @@ -2648,6 +2633,7 @@ class TestgenInstallAction(ComposeActionMixin, AnalyticsMultiStepAction): ComposeStartStep, TestGenSetupDatabaseStep, TestGenUpgradeDatabaseStep, + TestgenQuickStartStep, ] pip_intro = [ "Installing TestGen with pip.", @@ -2711,12 +2697,11 @@ def get_parser(self, sub_parsers): default=None, help="Path to SSL key file.", ) - # Pip-only args parser.add_argument( "--no-demo", - dest="no_demo", - action="store_true", - help="(Pip mode only) Skip generating demo data after install.", + dest="generate_demo", + action="store_false", + help="Skip generating demo data after install. Default is to generate.", ) # Docker-only args parser.add_argument( @@ -2724,7 +2709,7 @@ def get_parser(self, sub_parsers): 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", @@ -2732,7 +2717,7 @@ 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'" ), ) @@ -2769,55 +2754,46 @@ def _resolve_install_mode(self, args): LOG.info("tg install resolved to %s mode", mode) def _auto_select_mode(self, args): - # Probe with the same requirement list a real Docker install would check - # (including REQ_TESTGEN_IMAGE, since some networks may block docker.io image pulls). - if not all(req.check_availability(self, args, quiet=True) for req in self.docker_requirements): - CONSOLE.space() - CONSOLE.msg("Docker is not fully available on this machine.") - CONSOLE.msg("TestGen can be installed with pip instead, which uses an embedded Postgres database.") - CONSOLE.space() - try: - choice = input(f"{CONSOLE.MARGIN}Install TestGen with pip? [Y/n]: ").strip().lower() - except (KeyboardInterrupt, EOFError): - print("") - raise AbortAction - CONSOLE.space() - if choice in ("", "y", "yes"): - return INSTALL_MODE_PIP - if getattr(sys, "frozen", False): - CONSOLE.msg("Aborted. Fix the Docker prerequisites and select 'Install TestGen' from the menu again.") - else: - CONSOLE.msg( - f"Aborted. To retry with Docker, fix the prerequisites and run " - f"`python3 {INSTALLER_NAME} {args.prod} install --docker`." - ) - CONSOLE.msg( - f"To install with pip explicitly, run `python3 {INSTALLER_NAME} {args.prod} install --pip`." - ) - raise AbortAction + # 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("Two installation modes are available:") + 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. Provides a fully managed " - "environment with an isolated PostgreSQL container." + 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 light-weight Python installation suited for evaluation. Manages the PostgreSQL " - "database on the file system." - ) + 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(f"{CONSOLE.MARGIN}Install with Docker [d] or pip [p]? (default: d): ").strip().lower() + choice = input(prompt).strip().lower() except (KeyboardInterrupt, EOFError): print("") raise AbortAction - if choice in ("", "d", "docker"): + 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 @@ -2838,8 +2814,7 @@ class TestgenStandaloneUpgradeStep(Step): label = "Upgrading the application database" def execute(self, action, args): - testgen_path = resolve_testgen_path(action, args) - action.run_cmd(testgen_path, "upgrade-system-version") + run_testgen_cli(action, args, "upgrade-system-version") class TestgenUpgradeAction(ComposeActionMixin, AnalyticsMultiStepAction): @@ -2868,7 +2843,7 @@ def get_parser(self, sub_parsers): "--skip-verify", dest="skip_verify", action="store_true", - help="Whether to skip the version check before upgrading.", + help="(Docker mode only) Whether to skip the version check before upgrading.", ) parser.add_argument( "--pull-timeout", @@ -2876,7 +2851,7 @@ 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'" ), ) @@ -3127,6 +3102,10 @@ def execute(self, args): CONSOLE.msg("Observability demo configuration missing.") raise AbortAction + if self._resolved_mode == INSTALL_MODE_PIP and resolve_uv_path(self.data_folder) is None: + CONSOLE.msg(f"uv not found. To install TestGen, {command_hint(args.prod, 'install', 'Install TestGen')}.") + 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): @@ -3136,11 +3115,9 @@ def execute(self, args): CONSOLE.msg("This process may take up to 3 minutes depending on your system resources and network speed.") CONSOLE.space() - if args.obs_export: - self.run_dk_demo_container("tg-run-demo") - export_args = [] if args.obs_export: + self.run_dk_demo_container("tg-run-demo") with open(self.data_folder / DEMO_CONFIG_FILE, "r") as file: json_config = json.load(file) export_args = [ @@ -3150,48 +3127,11 @@ def execute(self, args): json_config["api_key"], ] - if self._resolved_mode == INSTALL_MODE_DOCKER: - self._run_docker_demo(args, export_args) - else: - self._run_pip_demo(args, export_args) - - CONSOLE.title("Demo SUCCEEDED") - - def _run_docker_demo(self, args, export_args): - cli_commands = [["testgen", "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, - ) - - def _run_pip_demo(self, args, export_args): - if resolve_uv_path(self.data_folder) is None: - CONSOLE.msg(f"uv not found. To install TestGen, {command_hint(args.prod, 'install', 'Install TestGen')}.") - raise AbortAction - - testgen_path = resolve_testgen_path(self, args) - self.run_cmd(testgen_path, "quick-start", *export_args) + run_testgen_cli(self, args, "quick-start", *export_args) if args.obs_export: - self.run_cmd( - testgen_path, + run_testgen_cli( + self, + args, "export-observability", "--project-key", "DEFAULT", @@ -3199,6 +3139,8 @@ def _run_pip_demo(self, args, export_args): "default-suite-1", ) + CONSOLE.title("Demo SUCCEEDED") + class TestgenDeleteDemoAction(DemoContainerAction, ComposeActionMixin): """Delete TestGen demo data — Docker-exec or pip-direct based on the marker.""" @@ -3242,24 +3184,11 @@ def execute(self, args): if not tg_status: CONSOLE.msg("TestGen must be running for its demo data to be cleaned.") raise AbortAction - self.run_cmd( - "docker", - "compose", - "-f", - self.get_compose_file_path(args), - "exec", - "engine", - "testgen", - "setup-system-db", - "--delete-db", - "--yes", - ) - else: - if resolve_uv_path(self.data_folder) is None: - CONSOLE.msg("uv not found. Cannot clean TestGen standalone database.") - raise AbortAction - testgen_path = resolve_testgen_path(self, args) - self.run_cmd(testgen_path, "setup-system-db", "--delete-db", "--yes") + elif resolve_uv_path(self.data_folder) is None: + CONSOLE.msg("uv not found. Cannot clean TestGen standalone database.") + raise AbortAction + + run_testgen_cli(self, args, "setup-system-db", "--delete-db", "--yes") CONSOLE.title("Demo data DELETED") diff --git a/tests/conftest.py b/tests/conftest.py index 5b989d1..f87b1c1 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -184,7 +184,7 @@ def args_mock(): ns.image = TESTGEN_DEFAULT_IMAGE ns.obs_export = False ns.install_mode = None - ns.no_demo = False + ns.generate_demo = True ns.api_port = 8530 # Observability defaults diff --git a/tests/test_tg_install.py b/tests/test_tg_install.py index 645f83f..ce2cf5f 100644 --- a/tests/test_tg_install.py +++ b/tests/test_tg_install.py @@ -121,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_demo.py b/tests/test_tg_pip_demo.py index ba63c65..0ba0e81 100644 --- a/tests/test_tg_pip_demo.py +++ b/tests/test_tg_pip_demo.py @@ -184,23 +184,25 @@ def test_run_demo_aborts_without_install(run_demo_action, args_mock, console_msg @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): +def test_run_demo_routes_by_marker(run_demo_action, args_mock, tmp_data_folder, install_mode, start_cmd_mock): write_install_marker(Path(tmp_data_folder), "tg", install_mode) - run_demo_action._resolve_install_mode(args_mock) + with ( - patch.object(run_demo_action, "_run_pip_demo") as pip_branch, - patch.object(run_demo_action, "_run_docker_demo") as docker_branch, + 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: - pip_branch.assert_called_once() - docker_branch.assert_not_called() + assert quick_start_calls[0].args[0] == TESTGEN_PATH else: - docker_branch.assert_called_once() - pip_branch.assert_not_called() + 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 diff --git a/tests/test_tg_pip_install.py b/tests/test_tg_pip_install.py index 591974f..e6bfb27 100644 --- a/tests/test_tg_pip_install.py +++ b/tests/test_tg_pip_install.py @@ -159,7 +159,7 @@ def test_tg_pip_install_auto_runs_demo(pip_install_action, start_cmd_mock): @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.no_demo = True + args_mock.generate_demo = False pip_install_action.execute() @@ -188,11 +188,9 @@ def test_tg_pip_install_password_redacted_in_logs(pip_install_action, start_cmd_ @pytest.mark.integration -@pytest.mark.parametrize("user_input", ["", "y", "yes", "Y"]) -def test_auto_mode_picks_pip_when_user_confirms_pip_fallback( - install_action, args_mock, console_msg_mock, user_input -): - """Docker probe fails → user confirms pip → resolve to pip.""" +@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), @@ -201,28 +199,48 @@ def test_auto_mode_picks_pip_when_user_confirms_pip_fallback( assert install_action._resolved_mode == INSTALL_MODE_PIP assert install_action.analytics.additional_properties["install_mode"] == INSTALL_MODE_PIP - console_msg_mock.assert_any_msg_contains("Docker is not fully available") @pytest.mark.integration -@pytest.mark.parametrize("user_input", ["n", "no", "N", "anything-else"]) -def test_auto_mode_aborts_when_user_declines_pip_fallback(install_action, args_mock, console_msg_mock, user_input): - """Docker probe fails → user declines pip → abort with hint.""" +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=user_input), - pytest.raises(AbortAction), + patch("builtins.input", return_value="") as input_mock, ): install_action._resolve_install_mode(args_mock) - console_msg_mock.assert_any_msg_contains("Aborted") - console_msg_mock.assert_any_msg_contains("install --docker") + 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_pip_fallback_prompt_aborts_on_user_interrupt(install_action, args_mock, interrupt): - """Docker probe fails → user interrupts the pip-fallback prompt → abort.""" +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), From 555c8485ab16b380e078c108c363f22c7f9c84c8 Mon Sep 17 00:00:00 2001 From: Aarthy Adityan Date: Tue, 5 May 2026 14:40:56 -0400 Subject: [PATCH 3/6] fix: address review feedback --- dk-installer.py | 26 +++++++++++--------------- tests/test_uv_bootstrap.py | 13 +++++++++---- 2 files changed, 20 insertions(+), 19 deletions(-) diff --git a/dk-installer.py b/dk-installer.py index 42a0547..8bf41da 100755 --- a/dk-installer.py +++ b/dk-installer.py @@ -1172,18 +1172,19 @@ def run(self, parent=None): def get_uv_asset(prod: str) -> tuple[str, str]: - """Return (asset_name, sha256) for the current platform, or raise InstallerError.""" + """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) - raise InstallerError( - f"No prebuilt uv binary available for platform {key[0]}/{key[1]}. " - f"Supported: {supported}. " - f"Install uv manually (https://docs.astral.sh/uv/getting-started/installation/) and re-run, " + 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]: @@ -2397,18 +2398,20 @@ def resolve_testgen_path(action, args) -> str: ctx = getattr(action, "ctx", None) or {} uv_path = ctx.get("uv_path") or resolve_uv_path(action.data_folder) if uv_path is None: - raise InstallerError("uv not found.") + 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(): - raise InstallerError( - f"testgen script not found at {testgen_path}. " + 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) @@ -3102,10 +3105,6 @@ def execute(self, args): CONSOLE.msg("Observability demo configuration missing.") raise AbortAction - if self._resolved_mode == INSTALL_MODE_PIP and resolve_uv_path(self.data_folder) is None: - CONSOLE.msg(f"uv not found. To install TestGen, {command_hint(args.prod, 'install', 'Install TestGen')}.") - 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): @@ -3184,9 +3183,6 @@ def execute(self, args): if not tg_status: CONSOLE.msg("TestGen must be running for its demo data to be cleaned.") raise AbortAction - elif resolve_uv_path(self.data_folder) is None: - CONSOLE.msg("uv not found. Cannot clean TestGen standalone database.") - raise AbortAction run_testgen_cli(self, args, "setup-system-db", "--delete-db", "--yes") diff --git a/tests/test_uv_bootstrap.py b/tests/test_uv_bootstrap.py index 0f88fa5..6b4adb1 100644 --- a/tests/test_uv_bootstrap.py +++ b/tests/test_uv_bootstrap.py @@ -8,6 +8,7 @@ import pytest from tests.installer import ( + AbortAction, InstallerError, UV_ASSETS, UV_VERSION, @@ -69,13 +70,15 @@ def test_get_uv_asset_known_platform(monkeypatch): @pytest.mark.unit -def test_get_uv_asset_unknown_platform_raises(monkeypatch): +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(InstallerError, match="No prebuilt uv binary available for platform SunOS/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): @@ -236,13 +239,15 @@ def test_download_failure_retries_then_fails(uv_step, bootstrap_action, args_moc @pytest.mark.unit -def test_unsupported_platform_raises(uv_step, bootstrap_action, args_mock, monkeypatch): +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(InstallerError, match="No prebuilt uv binary available"), + pytest.raises(AbortAction), ): uv_step.execute(bootstrap_action, args_mock) + + console_msg_mock.assert_any_msg_contains("No prebuilt uv binary available") From 04feb2ebbb707353f0bb2e2dfbf0df0fb1bd8983 Mon Sep 17 00:00:00 2001 From: Aarthy Adityan Date: Tue, 5 May 2026 14:55:36 -0400 Subject: [PATCH 4/6] refactor: move install marker logic to a class --- dk-installer.py | 115 +++++++++++++++++++---------------- tests/test_install_type.py | 46 +++++++++----- tests/test_tg_pip_delete.py | 6 +- tests/test_tg_pip_demo.py | 12 ++-- tests/test_tg_pip_install.py | 4 +- tests/test_tg_pip_upgrade.py | 8 +-- tests/test_tg_run_demo.py | 4 +- tests/test_tg_start.py | 6 +- tests/test_tg_upgrade.py | 4 +- 9 files changed, 114 insertions(+), 91 deletions(-) diff --git a/dk-installer.py b/dk-installer.py index 8bf41da..52a5259 100755 --- a/dk-installer.py +++ b/dk-installer.py @@ -266,51 +266,60 @@ def command_hint(prod: str, subcmd: str, menu_label: str) -> str: return f"run `python3 {INSTALLER_NAME} {prod} {subcmd}`" -def read_install_mode(data_folder: pathlib.Path, prod: str, compose_file_name: str) -> typing.Optional[str]: - """Return 'docker', 'pip', or None for the ``prod`` install in data_folder. - - Reads the install marker file. Falls back to detecting a legacy Docker install - (the product's compose file + credentials) from before the marker was introduced. +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. """ - marker_path = data_folder / INSTALL_MARKER_FILE.format(prod) - if marker_path.exists(): - try: - data = json.loads(marker_path.read_text()) - except Exception: - LOG.exception("Failed to read install marker at %s", marker_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 (data_folder / compose_file_name).exists() and (data_folder / CREDENTIALS_FILE.format(prod)).exists(): - LOG.info("No marker present; detected legacy Docker install in %s", data_folder) - return INSTALL_MODE_DOCKER + 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) - return None + 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 write_install_marker(data_folder: pathlib.Path, prod: str, install_mode: str, **extra) -> None: - if install_mode not in (INSTALL_MODE_DOCKER, INSTALL_MODE_PIP): - raise ValueError(f"Unknown install_mode: {install_mode}") - marker_path = data_folder / INSTALL_MARKER_FILE.format(prod) - now = datetime.datetime.now(datetime.timezone.utc).isoformat() - created_on = now - if marker_path.exists(): - try: - existing = json.loads(marker_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", marker_path) - data = { - "install_mode": install_mode, - "created_on": created_on, - "last_updated_on": now, - **extra, - } - marker_path.write_text(json.dumps(data, indent=2)) + def unlink(self) -> None: + if self.path.exists(): + self.path.unlink() @contextlib.contextmanager @@ -2737,7 +2746,7 @@ def get_requirements(self, args): return [] def _resolve_install_mode(self, args): - existing = read_install_mode(self.data_folder, args.prod, args.compose_file_name) + 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() @@ -2769,14 +2778,14 @@ def _auto_select_mode(self, args): 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 - ) + 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.msg( + " Sets up an isolated Python environment and manages the PostgreSQL database on the file system." + ) CONSOLE.space() if docker_ready: @@ -2805,7 +2814,7 @@ def _auto_select_mode(self, args): def execute(self, args): self.intro_text = self.pip_intro if self._resolved_mode == INSTALL_MODE_PIP else self.docker_intro super().execute(args) - write_install_marker(self.data_folder, args.prod, self._resolved_mode) + 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. @@ -2889,7 +2898,7 @@ def get_requirements(self, args): ] def _resolve_install_mode(self, args): - mode = read_install_mode(self.data_folder, args.prod, args.compose_file_name) + 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')}.") @@ -2902,7 +2911,7 @@ def _resolve_install_mode(self, args): def execute(self, args): super().execute(args) - write_install_marker(self.data_folder, args.prod, self._resolved_mode) + InstallMarker(self.data_folder, args.prod).write(self._resolved_mode) class TestgenStartAction(Action, ComposeActionMixin): @@ -2928,7 +2937,7 @@ def get_requirements(self, args): return [] def _resolve_install_mode(self, args): - mode = read_install_mode(self.data_folder, args.prod, args.compose_file_name) + 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')}.") @@ -2992,7 +3001,7 @@ def get_requirements(self, args): 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 = read_install_mode(self.data_folder, args.prod, args.compose_file_name) + 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 @@ -3008,7 +3017,7 @@ def execute(self, args): self._delete_docker(args) else: self._delete_pip(args) - remove_path(self.data_folder / INSTALL_MARKER_FILE.format(args.prod)) + InstallMarker(self.data_folder, args.prod, args.compose_file_name).unlink() def _delete_docker(self, args): if self.get_compose_file_path(args).exists(): @@ -3087,7 +3096,7 @@ def get_requirements(self, args): return [] def _resolve_install_mode(self, args): - mode = read_install_mode(self.data_folder, args.prod, args.compose_file_name) + 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')}.") @@ -3159,7 +3168,7 @@ def get_requirements(self, args): def _resolve_install_mode(self, args): # Like delete: idempotent, so "no install" returns rather than aborts. - mode = read_install_mode(self.data_folder, args.prod, args.compose_file_name) + 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 diff --git a/tests/test_install_type.py b/tests/test_install_type.py index e5fa919..a2215ee 100644 --- a/tests/test_install_type.py +++ b/tests/test_install_type.py @@ -8,9 +8,8 @@ INSTALL_MARKER_FILE, INSTALL_MODE_DOCKER, INSTALL_MODE_PIP, + InstallMarker, TESTGEN_COMPOSE_FILE, - read_install_mode, - write_install_marker, ) @@ -21,7 +20,7 @@ def data_folder(tmp_data_folder): @pytest.mark.unit def test_read_install_mode_returns_none_when_empty(data_folder): - assert read_install_mode(data_folder, "tg", "docker-compose.yml") is None + assert InstallMarker(data_folder, "tg", "docker-compose.yml").read() is None @pytest.mark.unit @@ -29,7 +28,7 @@ def test_read_install_mode_returns_none_when_empty(data_folder): 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 read_install_mode(data_folder, "tg", "docker-compose.yml") == install_mode + assert InstallMarker(data_folder, "tg", "docker-compose.yml").read() == install_mode @pytest.mark.unit @@ -37,19 +36,19 @@ 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 read_install_mode(data_folder, "tg", "docker-compose.yml") == INSTALL_MODE_DOCKER + 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 function isn't TestGen-specific: pass a different product's + # 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 read_install_mode(data_folder, "obs", "obs-docker-compose.yml") == INSTALL_MODE_DOCKER + 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 read_install_mode(data_folder, "obs", "docker-compose.yml") is None + assert InstallMarker(data_folder, "obs", "docker-compose.yml").read() is None @pytest.mark.unit @@ -57,7 +56,7 @@ 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 read_install_mode(data_folder, "tg", "docker-compose.yml") is None + assert InstallMarker(data_folder, "tg", "docker-compose.yml").read() is None @pytest.mark.unit @@ -66,19 +65,21 @@ def test_read_install_mode_malformed_marker_falls_back_to_legacy(data_folder): (data_folder / TESTGEN_COMPOSE_FILE).write_text("version: '3'") (data_folder / CREDENTIALS_FILE.format("tg")).write_text("admin\n") - assert read_install_mode(data_folder, "tg", "docker-compose.yml") == INSTALL_MODE_DOCKER + 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 read_install_mode(data_folder, "tg", "docker-compose.yml") is None + assert InstallMarker(data_folder, "tg", "docker-compose.yml").read() is None @pytest.mark.unit def test_write_install_marker_round_trip(data_folder): - write_install_marker(data_folder, "tg", INSTALL_MODE_PIP, version="5.9.4", python_version="3.13.1") + 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 @@ -86,15 +87,15 @@ def test_write_install_marker_round_trip(data_folder): assert data["python_version"] == "3.13.1" assert "created_on" in data assert "last_updated_on" in data - assert read_install_mode(data_folder, "tg", "docker-compose.yml") == INSTALL_MODE_PIP + 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): - write_install_marker(data_folder, "tg", INSTALL_MODE_PIP, version="5.9.4") + InstallMarker(data_folder, "tg").write(INSTALL_MODE_PIP, version="5.9.4") initial = json.loads((data_folder / INSTALL_MARKER_FILE.format("tg")).read_text()) - write_install_marker(data_folder, "tg", INSTALL_MODE_PIP, version="5.10.0") + 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"] @@ -104,6 +105,19 @@ def test_write_install_marker_preserves_created_on_across_writes(data_folder): @pytest.mark.unit def test_write_install_marker_rejects_unknown_type(data_folder): with pytest.raises(ValueError, match="Unknown install_mode"): - write_install_marker(data_folder, "tg", "sideways") + 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_pip_delete.py b/tests/test_tg_pip_delete.py index 84b86b0..9a1fd43 100644 --- a/tests/test_tg_pip_delete.py +++ b/tests/test_tg_pip_delete.py @@ -9,7 +9,7 @@ INSTALL_MODE_DOCKER, INSTALL_MODE_PIP, TestgenDeleteAction, - write_install_marker, + InstallMarker, ) @@ -19,7 +19,7 @@ def pip_delete_action(action_cls, args_mock, tmp_data_folder, start_cmd_mock, tm action = TestgenDeleteAction() args_mock.prod = "tg" args_mock.action = "delete" - write_install_marker(action.data_folder, args_mock.prod, INSTALL_MODE_PIP) + 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 @@ -158,7 +158,7 @@ def test_delete_nothing_to_delete(delete_action, args_mock, console_msg_mock): @pytest.mark.integration def test_delete_routes_to_pip_and_removes_marker(delete_action, args_mock, tmp_data_folder, tmp_path): - write_install_marker(Path(tmp_data_folder), "tg", INSTALL_MODE_PIP) + 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) diff --git a/tests/test_tg_pip_demo.py b/tests/test_tg_pip_demo.py index 0ba0e81..7bc7605 100644 --- a/tests/test_tg_pip_demo.py +++ b/tests/test_tg_pip_demo.py @@ -10,7 +10,7 @@ INSTALL_MODE_PIP, TestgenDeleteDemoAction, TestgenRunDemoAction, - write_install_marker, + InstallMarker, ) @@ -24,7 +24,7 @@ def pip_run_demo_action(action_cls, args_mock, tmp_data_folder, start_cmd_mock): action = TestgenRunDemoAction() args_mock.prod = "tg" args_mock.action = "run-demo" - write_install_marker(action.data_folder, args_mock.prod, INSTALL_MODE_PIP) + 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 ( @@ -118,7 +118,7 @@ def test_pip_run_demo_aborts_without_uv(action_cls, args_mock, tmp_data_folder, args_mock.prod = "tg" args_mock.action = "run-demo" args_mock.obs_export = False - write_install_marker(action.data_folder, args_mock.prod, INSTALL_MODE_PIP) + InstallMarker(action.data_folder, args_mock.prod).write(INSTALL_MODE_PIP) action._resolved_mode = INSTALL_MODE_PIP with ( @@ -136,7 +136,7 @@ def pip_delete_demo_action(action_cls, args_mock, tmp_data_folder, start_cmd_moc action = TestgenDeleteDemoAction() args_mock.prod = "tg" args_mock.action = "delete-demo" - write_install_marker(action.data_folder, args_mock.prod, INSTALL_MODE_PIP) + 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), @@ -185,7 +185,7 @@ def test_run_demo_aborts_without_install(run_demo_action, args_mock, console_msg @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): - write_install_marker(Path(tmp_data_folder), "tg", install_mode) + InstallMarker(Path(tmp_data_folder), "tg").write(install_mode) run_demo_action._resolve_install_mode(args_mock) with ( @@ -219,7 +219,7 @@ def delete_demo_action(action_cls, args_mock, tmp_data_folder): @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): - write_install_marker(Path(tmp_data_folder), "tg", install_mode) + InstallMarker(Path(tmp_data_folder), "tg").write(install_mode) delete_demo_action._resolve_install_mode(args_mock) with ( diff --git a/tests/test_tg_pip_install.py b/tests/test_tg_pip_install.py index e6bfb27..399026f 100644 --- a/tests/test_tg_pip_install.py +++ b/tests/test_tg_pip_install.py @@ -12,7 +12,7 @@ TESTGEN_MAJOR_VERSION, TESTGEN_PYTHON_VERSION, TestgenInstallAction, - write_install_marker, + InstallMarker, ) @@ -300,7 +300,7 @@ def test_dispatcher_routes_to_docker_with_flag(install_action, args_mock): @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): - write_install_marker(Path(tmp_data_folder), "tg", existing) + InstallMarker(Path(tmp_data_folder), "tg").write(existing) with pytest.raises(AbortAction): install_action._resolve_install_mode(args_mock) diff --git a/tests/test_tg_pip_upgrade.py b/tests/test_tg_pip_upgrade.py index f90b06d..aa25dd7 100644 --- a/tests/test_tg_pip_upgrade.py +++ b/tests/test_tg_pip_upgrade.py @@ -10,7 +10,7 @@ INSTALL_MODE_DOCKER, INSTALL_MODE_PIP, TestgenUpgradeAction, - write_install_marker, + InstallMarker, ) @@ -20,7 +20,7 @@ def pip_upgrade_action(action_cls, args_mock, tmp_data_folder, start_cmd_mock): action = TestgenUpgradeAction() args_mock.prod = "tg" args_mock.action = "upgrade" - write_install_marker(action.data_folder, args_mock.prod, INSTALL_MODE_PIP) + InstallMarker(action.data_folder, args_mock.prod).write(INSTALL_MODE_PIP) action._resolved_mode = INSTALL_MODE_PIP action.steps = action.pip_steps with ( @@ -137,7 +137,7 @@ def test_upgrade_aborts_with_no_install(upgrade_action, args_mock, console_msg_m @pytest.mark.integration def test_upgrade_resolves_to_pip_when_marker_says_pip(upgrade_action, args_mock, tmp_data_folder): - write_install_marker(Path(tmp_data_folder), "tg", INSTALL_MODE_PIP) + InstallMarker(Path(tmp_data_folder), "tg").write(INSTALL_MODE_PIP) upgrade_action._resolve_install_mode(args_mock) @@ -148,7 +148,7 @@ def test_upgrade_resolves_to_pip_when_marker_says_pip(upgrade_action, args_mock, @pytest.mark.integration def test_upgrade_resolves_to_docker_when_marker_says_docker(upgrade_action, args_mock, tmp_data_folder): - write_install_marker(Path(tmp_data_folder), "tg", INSTALL_MODE_DOCKER) + InstallMarker(Path(tmp_data_folder), "tg").write(INSTALL_MODE_DOCKER) upgrade_action._resolve_install_mode(args_mock) diff --git a/tests/test_tg_run_demo.py b/tests/test_tg_run_demo.py index a507ebd..459a70e 100644 --- a/tests/test_tg_run_demo.py +++ b/tests/test_tg_run_demo.py @@ -7,7 +7,7 @@ INSTALL_MODE_DOCKER, AbortAction, TestgenRunDemoAction, - write_install_marker, + InstallMarker, ) @@ -17,7 +17,7 @@ def tg_run_demo_action(action_cls, args_mock, tmp_data_folder, start_cmd_mock): args_mock.prod = "tg" args_mock.action = "run-demo" # Seed a Docker install marker so the unified action picks the Docker path. - write_install_marker(action.data_folder, args_mock.prod, INSTALL_MODE_DOCKER) + 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)): diff --git a/tests/test_tg_start.py b/tests/test_tg_start.py index da33bb7..f6d2315 100644 --- a/tests/test_tg_start.py +++ b/tests/test_tg_start.py @@ -10,7 +10,7 @@ InstallerError, TestgenStartAction, start_testgen_app, - write_install_marker, + InstallMarker, ) @@ -154,7 +154,7 @@ def test_start_action_aborts_with_no_install(start_action, args_mock, console_ms def test_start_action_runs_compose_up_in_docker_mode( start_action, args_mock, tmp_data_folder, start_cmd_mock, compose_path ): - write_install_marker(Path(tmp_data_folder), "tg", INSTALL_MODE_DOCKER) + InstallMarker(Path(tmp_data_folder), "tg").write(INSTALL_MODE_DOCKER) start_action._resolve_install_mode(args_mock) start_action.execute(args_mock) @@ -174,7 +174,7 @@ def test_start_action_runs_compose_up_in_docker_mode( @pytest.mark.integration def test_start_action_routes_to_helper_in_pip_mode(start_action, args_mock, tmp_data_folder): - write_install_marker(Path(tmp_data_folder), "tg", INSTALL_MODE_PIP) + 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: diff --git a/tests/test_tg_upgrade.py b/tests/test_tg_upgrade.py index e491e7d..67cae08 100644 --- a/tests/test_tg_upgrade.py +++ b/tests/test_tg_upgrade.py @@ -10,7 +10,7 @@ CommandFailed, TESTGEN_MAJOR_VERSION, TestgenUpgradeAction, - write_install_marker, + InstallMarker, ) @@ -20,7 +20,7 @@ def tg_upgrade_action(action_cls, args_mock, tmp_data_folder, start_cmd_mock, re args_mock.prod = "tg" args_mock.action = "upgrade" # Seed a Docker install marker so the unified upgrade resolves to Docker mode. - write_install_marker(action.data_folder, args_mock.prod, INSTALL_MODE_DOCKER) + InstallMarker(action.data_folder, args_mock.prod).write(INSTALL_MODE_DOCKER) action._resolved_mode = INSTALL_MODE_DOCKER action.steps = action.docker_steps yield action From a4cf2e6d54f81d63f3f299b3ab0577e292fecff4 Mon Sep 17 00:00:00 2001 From: Aarthy Adityan Date: Tue, 5 May 2026 16:45:03 -0400 Subject: [PATCH 5/6] feat(testgen): auto-open browser after docker install/start Pip mode opens the browser via Streamlit. Docker mode now matches: install opens to the configured port, start reads TG_UI_BASE_URL from the compose file. Both fall back silently when no browser is available. Tests patch webbrowser.open globally so local pytest runs don't spawn tabs. Also: only print the log path on tg start (it's already in the credentials file at install time), tighten the UvBootstrapStep label, and reword a delete-demo console message. Co-Authored-By: Claude Opus 4.7 (1M context) --- dk-installer.py | 38 +++++++++++++++++++++++++++++++++----- tests/conftest.py | 7 +++++++ 2 files changed, 40 insertions(+), 5 deletions(-) diff --git a/dk-installer.py b/dk-installer.py index 52a5259..3c986db 100755 --- a/dk-installer.py +++ b/dk-installer.py @@ -31,6 +31,7 @@ import time import urllib.request import urllib.parse +import webbrowser import zipfile import typing @@ -145,6 +146,15 @@ def get_tg_url(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): if images: action.run_cmd( @@ -1211,7 +1221,7 @@ def resolve_uv_path(data_folder: pathlib.Path) -> typing.Optional[str]: class UvBootstrapStep(Step): - label = "Setting up the Python environment" + label = "Preparing the Python environment" def pre_execute(self, action, args): # Resolve uv eagerly so later steps' pre_execute hooks can use it. @@ -2501,7 +2511,10 @@ def start_testgen_app(action, args) -> None: ) CONSOLE.msg(f"TestGen is running at {url}.") - CONSOLE.msg(f"Logs: {simplify_path(TESTGEN_LOG_FILE_PATH)}") + # 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')}.") @@ -2817,9 +2830,12 @@ def execute(self, 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. + # ``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): @@ -2949,11 +2965,23 @@ def _resolve_install_mode(self, args): def execute(self, args): if self._resolved_mode == INSTALL_MODE_DOCKER: CONSOLE.title("Start TestGen") - self.run_cmd("docker", "compose", "-f", self.get_compose_file_path(args), "up", "--wait") + 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) @@ -3186,7 +3214,7 @@ def execute(self, args): except Exception: pass - CONSOLE.msg("Cleaning up system database..") + CONSOLE.msg("Cleaning up application database..") if self._resolved_mode == INSTALL_MODE_DOCKER: tg_status = self.get_status(args) if not tg_status: diff --git a/tests/conftest.py b/tests/conftest.py index f87b1c1..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=[]) From 2bdc707546701b2333e640d8171a443ad32fa799 Mon Sep 17 00:00:00 2001 From: Aarthy Adityan Date: Tue, 5 May 2026 17:12:30 -0400 Subject: [PATCH 6/6] docs: refresh README for pip-install path additions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Recommend Docker when available; document the [d/p] auto-detect prompt and the failed-prereq fallback path - Mention --no-demo, --api-port, tg start, and the auto-browser-open - Drop the TG_STANDALONE_MODE=yes uv tool run … references - Bump pip-install time estimate to 4-8 minutes (matching the intro_text) Co-Authored-By: Claude Opus 4.7 (1M context) --- README.md | 47 ++++++++++++++++++----------------------------- dk-installer.py | 2 +- 2 files changed, 19 insertions(+), 30 deletions(-) diff --git a/README.md b/README.md index 6a0e7be..ae50f41 100644 --- a/README.md +++ b/README.md @@ -49,15 +49,17 @@ And it allows you to make fast, safe development changes. TestGen can be installed two ways: -- **pip** (default, recommended for evaluation) — no Docker required. The installer downloads `uv`, installs Python 3.13 if needed, and installs TestGen with an embedded Postgres database. -- **Docker** — deploys TestGen as a Docker Compose application. Use for team eval on a shared VM, or if you already standardize on Docker. +- **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 path | Required software | +| Install mode | Required software | |---|---| -| TestGen (pip, default) | [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/) 26+, [Docker Compose](https://docs.docker.com/compose/install/) 2.38+. | +| 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`. @@ -82,27 +84,20 @@ The [Data Observability quickstart](https://docs.datakitchen.io/tutorials/quicks Before going through the quickstart, complete the prequisites above and then the following steps to install the two products and setup the demo data. For any of the commands, you can view additional options by appending `--help` at the end. -### Install the TestGen application (pip, default) - -`tg install` defaults to a pip-based install with an embedded Postgres database and no Docker requirement. The installer downloads `uv` (if it isn't already on your PATH), uses it to install Python 3.13 (if needed) and TestGen in an isolated environment, then prints credentials plus the command to start the app. +### Install the TestGen application ```shell python3 dk-installer.py tg install ``` -The process typically takes 2-5 minutes. On completion the installer writes credentials to `dk-tg-credentials.txt` and prints the `testgen run-app` command to start the UI in a separate terminal. - -#### Install the TestGen application (Docker) +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. -If you prefer the Docker Compose install — for team evaluations on a shared VM, or if you already standardize on Docker — use the `--docker` flag: +* **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. -```shell -python3 dk-installer.py tg install --docker -``` +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. -The Docker install takes 5-10 minutes. `--port` sets a custom localhost port (default 8501). `--ssl-cert-file` / `--ssl-key-file` enable HTTPS. - -Either install path can later be upgraded with `python3 dk-installer.py tg upgrade` — the installer detects which flavor is present and upgrades accordingly. +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 @@ -152,25 +147,19 @@ Leave this process running, and continue with the [quickstart guide](https://doc ## Useful Commands -### DataOps TestGen (pip install, default) - -Start the app: `TG_STANDALONE_MODE=yes uv tool run testgen run-app` (reachable at `http://localhost:8501`) +### DataOps TestGen (pip install) -Stop the app: `Ctrl+C` in the terminal running `testgen run-app` +Start the app: `python3 dk-installer.py tg start` (reachable at `http://localhost:8501`, blocks until Ctrl+C) -Access the `testgen` CLI: `TG_STANDALONE_MODE=yes uv tool run testgen ` +Stop the app: `Ctrl+C` in the terminal running `tg start` Upgrade the app to latest version: `python3 dk-installer.py tg upgrade` ### DataOps TestGen (Docker install) -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. - -Access the _testgen_ CLI: `docker compose exec engine bash` (use `exit` to return to the regular terminal) - -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 3c986db..9200142 100755 --- a/dk-installer.py +++ b/dk-installer.py @@ -2662,7 +2662,7 @@ class TestgenInstallAction(ComposeActionMixin, AnalyticsMultiStepAction): ] pip_intro = [ "Installing TestGen with pip.", - "The process may take 2~5 minutes depending on your system resources and network speed.", + "The process may take 4~8 minutes depending on your system resources and network speed.", ] docker_intro = [ "Installing TestGen with Docker Compose.",