diff --git a/integrations/agentbricks/pyproject.toml b/integrations/agentbricks/pyproject.toml index fad3b192d..efc755f99 100644 --- a/integrations/agentbricks/pyproject.toml +++ b/integrations/agentbricks/pyproject.toml @@ -64,6 +64,7 @@ tests = [ "httpx>=0.28.1", "pytest==9.0.2", "pytest-asyncio==1.3.0", + "pytest-xdist==3.8.0", "uvicorn>=0.20.0", ] @@ -95,6 +96,8 @@ extend = "../../pyproject.toml" [tool.pytest.ini_options] testpaths = ["tests/unit_tests"] +# Lets the suites under tests/suites import the shared tests/harness package as `harness`. +pythonpath = ["tests"] [tool.ty.environment] root = ["./src", "./tests"] diff --git a/integrations/agentbricks/tests/harness/__init__.py b/integrations/agentbricks/tests/harness/__init__.py new file mode 100644 index 000000000..47679ee61 --- /dev/null +++ b/integrations/agentbricks/tests/harness/__init__.py @@ -0,0 +1 @@ +"""Plain-Python building blocks of the E2E suites; ``suites/conftest.py`` wires them into pytest.""" diff --git a/integrations/agentbricks/tests/harness/agentbricks_cli.py b/integrations/agentbricks/tests/harness/agentbricks_cli.py new file mode 100644 index 000000000..a46b4a38b --- /dev/null +++ b/integrations/agentbricks/tests/harness/agentbricks_cli.py @@ -0,0 +1,308 @@ +"""The ``agentbricks`` CLI: the system under test. + +Every test step is a real ``agentbricks`` subprocess, written as the command a developer would type +(``cli.run("tools add ...")``), so a regression in the CLI fails the test that makes it. That +includes reaching a running agent: tests call it with ``agentbricks endpoint invoke``. Nothing in +this module knows about any particular suite or deletes anything: the ``project`` fixture does. +""" + +from __future__ import annotations + +import contextlib +import dataclasses +import json +import os +import pathlib +import re +import shlex +import shutil +import signal +import socket +import subprocess +import sys +import time +import urllib.error +import urllib.request +import uuid +from collections.abc import Iterator, Mapping +from typing import Any + +from harness.logs import last_lines, last_nonempty_line, log, log_command +from harness.run import E2EError, RunConfig, now +from harness.target import TargetWorkspace + +# `endpoint invoke`'s error for a 5xx answer: the App restarting after a redeploy, a new App's +# session store coming up, or a model call failing on the server. +_SERVER_ERROR = re.compile(r"HTTP 5\d\d\b") +_NOT_FOUND = re.compile( + r"NOT_FOUND|RESOURCE_DOES_NOT_EXIST|does not exist|not found", re.IGNORECASE +) + + +@dataclasses.dataclass +class BackgroundCommand: + """A CLI command left running by ``AgentbricksCli.start``, e.g. ``agentbricks dev``.""" + + label: str + process: subprocess.Popen[str] + log_path: pathlib.Path + + def wait_until_serving(self, port: int, timeout: float = 1200) -> None: + """Block until something answers HTTP on ``port``; fails if the command exits first.""" + started = time.monotonic() + next_tick = 60.0 + while True: + if self.process.poll() is not None: + raise E2EError( + f"{self.label} exited {self.process.returncode}: {last_lines(self.log_path, 30)}" + ) + try: + with urllib.request.urlopen(f"http://127.0.0.1:{port}/", timeout=5): + return + except urllib.error.HTTPError as exc: + if exc.code < 500: + return + except (urllib.error.URLError, TimeoutError): + pass + elapsed = time.monotonic() - started + if elapsed > timeout: + raise E2EError( + f"{self.label} did not become reachable: {last_lines(self.log_path, 30)}" + ) + if elapsed >= next_tick: + log( + f"tick {now():%H:%M} | {self.label} | starting | {last_nonempty_line(self.log_path)}" + ) + next_tick += 60 + time.sleep(5) + + +class AgentbricksCli: + """Runs ``agentbricks`` commands against a target workspace, logging each to its own files.""" + + def __init__(self, target_workspace: TargetWorkspace, run_config: RunConfig, test_name: str): + self.target_workspace = target_workspace + self.run_config = run_config + # One directory per test: tests run concurrently and each numbers its commands from 1. + self.logs_dir = ( + run_config.output / "logs" / re.sub(r"[^A-Za-z0-9_.-]+", "-", test_name).strip("-") + ) + # The CLI prints a hint when it detects a coding agent; the suite must behave the same + # whoever launches it, and the hint would corrupt --output json errors on stderr. Test runs + # are not product usage, so they send no CLI telemetry. + self._env = { + **os.environ, + **target_workspace.env, + "AGENTBRICKS_DISABLE_AGENT_HINT": "1", + "AGENTBRICKS_DISABLE_TELEMETRY": "1", + } + self._sequence = 0 + + def run( + self, + command: str, + *, + timeout: float = 300, + check: bool = True, + retry_server_errors: bool = False, + ) -> subprocess.CompletedProcess[str]: + """Run ``agentbricks `` to completion. + + ``retry_server_errors`` retries an ``endpoint invoke`` the agent answered with a 5xx; any + other failure, including a 4xx, is the CLI's or the agent's and fails at once. + """ + argv = self._argv(command) + for attempt in range(1, 4): + result = self._run_once(argv, self._label(command), timeout) + output = result.stderr or result.stdout + if ( + result.returncode == 0 + or not retry_server_errors + or attempt == 3 + or _SERVER_ERROR.search(output) is None + ): + break + log(f"attempt {attempt}/3 | {command.split()[0]} | server error; retrying") + time.sleep(15) + if check and result.returncode != 0: + raise E2EError(f"Command failed ({result.returncode}): {shlex.join(argv)}\n{output}") + return result + + def invoke( + self, target: str, prompt: str, *, session_id: str | None = None, actor: str | None = None + ) -> str: + """``agentbricks --output json endpoint invoke `` with ``prompt``; returns its output. + + ``target`` is an App name, or ``--url --no-auth`` for `agentbricks dev`. Each call is + a new session unless ``session_id`` is given. The agent keys long-term memory by ``actor``, + which defaults to the session. + """ + invocation_id = str(uuid.uuid4()) + agent_input: dict[str, Any] = { + "model": self.run_config.model, + "messages": [{"role": "user", "content": prompt}], + } + if actor: + agent_input["actor"] = actor + body = { + "id": invocation_id, + # session_id is a top-level invocation field; the adapter rejects it nested in input. + "session_id": session_id or invocation_id, + "input": agent_input, + } + return self.run( + f"--output json endpoint invoke {target} --path /api/invocations " + f"--json {shlex.quote(json.dumps(body))}", + timeout=360, + retry_server_errors=True, + ).stdout + + @contextlib.contextmanager + def start( + self, command: str, *, env: Mapping[str, str] | None = None + ) -> Iterator[BackgroundCommand]: + """Run ``agentbricks `` in the background for the block, then stop it.""" + argv = self._argv(command) + label = self._label(command) + log_path = self.logs_dir / f"{label}.log" + log_path.parent.mkdir(parents=True, exist_ok=True) + log_command(argv) + with log_path.open("w", encoding="utf-8") as log_file: + process = subprocess.Popen( + argv, + text=True, + stdout=log_file, + stderr=subprocess.STDOUT, + start_new_session=True, + env={**self._env, **(env or {})}, + ) + try: + yield BackgroundCommand(label, process, log_path) + finally: + if process.poll() is None: + terminate_process_group(process) + log(f"{label} | stopped") + + def _argv(self, command: str) -> list[str]: + return [_agentbricks_bin(), *self.target_workspace.cli_args, *shlex.split(command)] + + def _label(self, command: str) -> str: + """A log name from the command's first two words, e.g. ``007-tools-add``.""" + self._sequence += 1 + words = shlex.split(command) + while words and words[0].startswith("-"): + # A leading global option such as ``--output json``, and its value. + words = words[1:] if "=" in words[0] else words[2:] + # The second word only when it is a subcommand or name, never a flag or a path. + names = words[:1] + [w for w in words[1:2] if re.fullmatch(r"[a-z][a-z0-9_-]*", w)] + return "-".join([f"{self._sequence:03d}", *names]) + + def _run_once( + self, argv: list[str], label: str, timeout: float + ) -> subprocess.CompletedProcess[str]: + """Run with stdout and stderr in log files, printing progress while a command is slow.""" + stdout_path = self.logs_dir / f"{label}.out" + stderr_path = self.logs_dir / f"{label}.err" + stdout_path.parent.mkdir(parents=True, exist_ok=True) + log_command(argv) + started = time.monotonic() + with ( + stdout_path.open("w", encoding="utf-8") as out, + stderr_path.open("w", encoding="utf-8") as err, + ): + process = subprocess.Popen( + argv, + env=self._env, + text=True, + stdout=out, + stderr=err, + start_new_session=True, + ) + next_tick = 60.0 + while process.poll() is None: + elapsed = time.monotonic() - started + if elapsed >= timeout: + terminate_process_group(process) + raise E2EError(f"{label} timed out after {timeout:.0f}s; log: {stderr_path}") + if elapsed >= next_tick: + progress = last_nonempty_line(stderr_path) or last_nonempty_line(stdout_path) + log(f"tick {now():%H:%M} | {label} | running | {progress}") + next_tick += 60.0 + time.sleep(0.2 if elapsed < 5 else 2) + stdout = stdout_path.read_text(encoding="utf-8", errors="replace") + stderr = stderr_path.read_text(encoding="utf-8", errors="replace") + for text in (stdout, stderr): + if text.strip(): + log(text) + return subprocess.CompletedProcess(argv, process.returncode, stdout, stderr) + + +def _agentbricks_bin() -> str: + """The CLI of the interpreter running the tests, i.e. the source checkout's own install.""" + sibling = pathlib.Path(sys.executable).parent / "agentbricks" + found = str(sibling) if sibling.exists() else shutil.which("agentbricks") + if found is None: + raise E2EError("The agentbricks CLI is not installed in the test environment.") + return found + + +def terminate_process_group(process: subprocess.Popen[str], *, grace: float = 20) -> None: + """Stop the command and what it started, returning only once it has exited. + + A command still running during teardown could create what the cleanup has already deleted. + """ + try: + os.killpg(process.pid, signal.SIGTERM) + except ProcessLookupError: + pass + try: + process.wait(timeout=grace) + except subprocess.TimeoutExpired: + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass + process.wait() + + +def not_found(result: subprocess.CompletedProcess[str]) -> bool: + """Whether a command failed because what it names doesn't exist, e.g. an App already deleted.""" + return result.returncode != 0 and bool(_NOT_FOUND.search(result.stderr or result.stdout)) + + +# Reading `endpoint invoke --output json`: the LangGraph agent lists every step in +# ``body.output.output``, the model's messages (``type == "ai"``) and each tool's own result +# (``type == "tool"``). + + +def excerpt(serialized: str, head: int = 1000, tail: int = 3000) -> str: + """The response for a failure message: tool results come last, so keep the end, not just the start.""" + if len(serialized) <= head + tail: + return serialized + omitted = len(serialized) - head - tail + return f"{serialized[:head]} ...[{omitted} chars omitted]... {serialized[-tail:]}" + + +def _steps(serialized: str) -> list[dict[str, Any]]: + try: + steps = json.loads(serialized)["body"]["output"]["output"] + except (ValueError, KeyError, TypeError) as exc: + raise AssertionError(f"Unexpected endpoint invoke output: {excerpt(serialized)}") from exc + return [step for step in steps if isinstance(step, dict)] + + +def tool_results(serialized: str) -> list[dict[str, Any]]: + """Each tool's own result, separate from what the model says about it.""" + return [step for step in _steps(serialized) if step.get("type") == "tool"] + + +def reply(serialized: str) -> str: + """The model's last message.""" + messages = [step for step in _steps(serialized) if step.get("type") == "ai"] + return str(messages[-1].get("content")) if messages else "" + + +def free_port() -> int: + with socket.socket() as probe: + probe.bind(("127.0.0.1", 0)) + return probe.getsockname()[1] diff --git a/integrations/agentbricks/tests/harness/cleanup.py b/integrations/agentbricks/tests/harness/cleanup.py new file mode 100644 index 000000000..8df2acb02 --- /dev/null +++ b/integrations/agentbricks/tests/harness/cleanup.py @@ -0,0 +1,279 @@ +"""Deleting what the E2E suites create: per test when it ends, and per run afterwards. + +Each test registers the delete commands for what it creates with the ``cleanup`` fixture +(``Cleanup``), which runs them when the test ends, pass or fail. ``delete_run_leftovers`` deletes +whatever is left: everything carrying the run's name prefix (stores, Apps, synced source folders, +schemas), found through the CLI where it can list them, and the Apps' Lakebase roles, found from +the identities the ``project`` fixture records because they outlive their App. It is safe to repeat +(a missing resource counts as deleted), runs once after every live run, and runs on its own with +``pytest tests/suites --cleanup-only --run-id `` after a run that was killed before it could. +""" + +from __future__ import annotations + +import dataclasses +import json +import pathlib +import time +from collections.abc import Callable, Iterator +from typing import Any + +from harness.agentbricks_cli import AgentbricksCli, not_found +from harness.logs import log +from harness.run import APP_PREFIX, E2EError, now, project_prefix, scratch_schema_prefix +from harness.workspace import AppIdentity, Workspace + +_APPS_RECORD = "apps.jsonl" + + +# One test + + +class Cleanup: + """What a test registers runs when it ends, newest first, pass or fail, like Go's ``t.Cleanup``. + + Register right after the step that creates something, or before it when a failed step can + still leave it behind. Something already gone counts as deleted; any other failure only warns, + because ``delete_run_leftovers`` retries everything the run left. + """ + + def __init__(self, cli: AgentbricksCli): + self._cli = cli + self._steps: list[tuple[str, Callable[[], None]]] = [] + + def __call__(self, command: str) -> None: + """Run ``agentbricks `` when the test ends.""" + self._steps.append((command, lambda: self._run(command))) + + def run_all(self) -> None: + while self._steps: + label, step = self._steps.pop() + try: + step() + except Exception as exc: + log(f"cleanup warning | {label} | {exc}") + + def _run(self, command: str) -> None: + result = self._cli.run(command, timeout=1200, check=False) + if result.returncode != 0 and not not_found(result): + raise E2EError((result.stderr or result.stdout).strip()) + + +# Identities recorded by per-test cleanup + + +def app_identity(cli: AgentbricksCli, workspace: Workspace, name: str) -> AppIdentity: + """The App's principal and Lakebase branch as far as they still exist; never raises.""" + principal: str | None = None + result = cli.run(f"--output json deployments get {name}", check=False) + message = (result.stderr or result.stdout).strip() + if result.returncode == 0: + try: + principal = json.loads(result.stdout).get("service_principal_client_id") + except (ValueError, AttributeError): + log(f"cleanup warning | principal lookup for {name} | {message}") + elif not not_found(result): + log(f"cleanup warning | principal lookup for {name} | {message}") + return AppIdentity(principal, workspace.lakebase_branch(name)) + + +def record_app_identity(output: pathlib.Path, app: str, identity: AppIdentity) -> None: + """Append one line for leftover cleanup, which runs in another process and cannot see worker memory.""" + line = json.dumps({"app": app, **dataclasses.asdict(identity)}) + "\n" + # The leftover sweep records too, and with --cleanup-only it may be the first to write. + output.mkdir(parents=True, exist_ok=True) + # One write per line: appends this small do not interleave across processes on POSIX. + with (output / _APPS_RECORD).open("a", encoding="utf-8") as record: + record.write(line) + + +def recorded_app_identities(output: pathlib.Path) -> dict[str, AppIdentity]: + """Every App identity the run's workers recorded, by App name.""" + path = output / _APPS_RECORD + identities: dict[str, AppIdentity] = {} + if not path.exists(): + return identities + for line in path.read_text(encoding="utf-8").splitlines(): + try: + data = json.loads(line) + identity = AppIdentity(data.get("service_principal"), data.get("branch")) + identities[data["app"]] = identities.get(data["app"], AppIdentity()).merged(identity) + except (ValueError, KeyError, TypeError, AttributeError): + continue + return identities + + +# One run + + +def delete_run_leftovers( + cli: AgentbricksCli, workspace: Workspace, run_id: str, output: pathlib.Path, catalog: str +) -> list[str]: + """Delete what a whole run left behind; returns a description of each failure.""" + failures: list[str] = [] + steps = ( + ("leaked stores", lambda: _delete_leaked_stores(cli, run_id)), + # Before the roles: the App's role is only safe to delete once the App is gone. + ("leaked Apps", lambda: _delete_leaked_apps(cli, workspace, run_id, output)), + ("leaked Lakebase roles", lambda: _delete_leaked_roles(workspace, output)), + ("leaked source folders", lambda: _delete_leaked_deployment_sources(workspace, run_id)), + ("leaked experiments", lambda: _delete_leaked_experiments(workspace, run_id)), + ("leaked schemas", lambda: _delete_leaked_schemas(workspace, catalog, run_id)), + ) + for label, action in steps: + try: + failures.extend(action()) + except Exception as exc: + failures.append(f"{label}: {exc}") + return failures + + +def _app_prefix(run_id: str) -> str: + return f"{APP_PREFIX}{project_prefix(run_id)}" + + +def _delete_leaked_stores(cli: AgentbricksCli, run_id: str) -> list[str]: + """Delete this run's memory and session stores, found by name prefix. + + `agentbricks init` names a project's stores after its folder, which carries the run's prefix. + """ + prefix = project_prefix(run_id) + failures: list[str] = [] + for kind, key, name_field, address_field in ( + ("memory stores", "managed_memory_stores", "display_name", "name"), + ("sessions stores", "session_stores", "session_store_name", "session_store_name"), + ): + for store in _list_all(cli, kind, key): + if not str(store.get(name_field) or "").startswith(prefix): + continue + address = store[address_field] + log(f"deleting leaked store {address}") + result = cli.run(f"{kind} delete {address} --yes", timeout=1200, check=False) + if result.returncode != 0: + failures.append(f"{kind} {address}: {(result.stderr or result.stdout).strip()}") + return failures + + +def _list_all(cli: AgentbricksCli, kind: str, key: str) -> Iterator[dict[str, Any]]: + token = "" + while True: + page = json.loads( + cli.run( + f"--output json {kind} list --page-size 100{token and f' --page-token {token}'}" + ).stdout + ) + yield from page.get(key) or [] + token = page.get("next_page_token") or "" + if not token: + return + + +def _delete_leaked_apps( + cli: AgentbricksCli, workspace: Workspace, run_id: str, output: pathlib.Path +) -> list[str]: + """Delete Apps whose per-test cleanup failed or never ran, then wait until every run App is gone. + + `deployments delete` returns while the App is still deleting, so a teardown's App can still be + listed here; deleting it again only warns, and the wait is what decides success. + """ + prefix = _app_prefix(run_id) + leaked: list[str] = [] + for deployment in json.loads(cli.run("--output json deployments list").stdout or "[]"): + name = str(deployment.get("name") or "") + if not name.startswith(prefix): + continue + log(f"deleting leaked App {name}") + leaked.append(name) + # The principal stops resolving once the App is gone, and the later steps need it. + record_app_identity(output, name, app_identity(cli, workspace, name)) + result = cli.run(f"deployments delete {name} --yes", timeout=1200, check=False) + if result.returncode != 0: + log(f"cleanup warning | {name} | {(result.stderr or result.stdout).strip()}") + # Teardowns deleted the rest of the run's Apps, which apps.jsonl names. + failures: list[str] = [] + for name in sorted(set(leaked) | set(recorded_app_identities(output))): + try: + wait_until_app_gone(cli, name) + except Exception as exc: + failures.append(f"App {name}: {exc}") + return failures + + +def wait_until_app_gone(cli: AgentbricksCli, name: str, timeout: float = 1200) -> None: + """Block until the App is gone; `deployments delete` returns while it is still deleting.""" + started = time.monotonic() + next_tick = 0.0 + while not not_found(cli.run(f"--output json deployments get {name}", check=False)): + elapsed = time.monotonic() - started + if elapsed >= timeout: + raise E2EError(f"App {name!r} still existed {timeout:.0f}s after delete returned.") + if elapsed >= next_tick: + log(f"tick {now():%H:%M} | delete-{name} | deleting") + next_tick += 60 + time.sleep(15) + + +def _delete_leaked_roles(workspace: Workspace, output: pathlib.Path) -> list[str]: + """Delete the Lakebase roles of this run's Apps, which `deployments delete` leaves behind.""" + recorded = list(recorded_app_identities(output).values()) + failures: list[str] = [] + for branch in sorted({item.branch for item in recorded if item.branch}): + # A Runtime Store create that fails partway records no branch but can still leave the App's + # role behind, so look for branchless identities on every branch this run used. + principals = [ + item.service_principal + for item in recorded + if item.branch in (branch, None) and item.service_principal + ] + try: + roles = workspace.roles_of(branch, principals) + except Exception as exc: + failures.append(f"Lakebase role lookup on {branch}: {exc}") + continue + for role in roles: + try: + workspace.delete_role(role) + except Exception as exc: + failures.append(f"Lakebase role {role}: {exc}") + return failures + + +def _delete_leaked_deployment_sources(workspace: Workspace, run_id: str) -> list[str]: + """Delete this run's synced source folders, which `apps delete` leaves behind.""" + failures: list[str] = [] + for path in workspace.deployment_sources_with_prefix(_app_prefix(run_id)): + log(f"deleting leaked deployment source {path}") + try: + workspace.delete_workspace_path(path) + except Exception as exc: + failures.append(f"deployment source {path}: {exc}") + return failures + + +def _delete_leaked_experiments(workspace: Workspace, run_id: str) -> list[str]: + """Delete this run's tracing experiments, which deploy creates and nothing deletes. + + `agentbricks init` names a project's experiment after its folder, under this shared path. + """ + failures: list[str] = [] + for experiment_id in workspace.experiments_with_prefix( + f"/Shared/agentbricks_traces/{project_prefix(run_id)}" + ): + log(f"deleting leaked experiment {experiment_id}") + try: + workspace.delete_experiment(experiment_id) + except Exception as exc: + failures.append(f"experiment {experiment_id}: {exc}") + return failures + + +def _delete_leaked_schemas(workspace: Workspace, catalog: str, run_id: str) -> list[str]: + """Delete this run's temporary schemas whose test teardown did not run, found by name prefix.""" + failures: list[str] = [] + for schema in workspace.schemas_with_prefix(catalog, scratch_schema_prefix(run_id)): + log(f"deleting leaked schema {schema}") + try: + workspace.delete_schema(schema) + except Exception as exc: + failures.append(f"schema {schema}: {exc}") + return failures diff --git a/integrations/agentbricks/tests/harness/logs.py b/integrations/agentbricks/tests/harness/logs.py new file mode 100644 index 000000000..d196f41bf --- /dev/null +++ b/integrations/agentbricks/tests/harness/logs.py @@ -0,0 +1,31 @@ +"""Test-log output: every command the suites run is echoed, so a run's log replays it.""" + +from __future__ import annotations + +import pathlib +import shlex +import sys +from collections.abc import Sequence + + +def log(text: str) -> None: + sys.stdout.write(text.rstrip() + "\n") + sys.stdout.flush() + + +def log_command(argv: Sequence[str], cwd: pathlib.Path | None = None) -> None: + prefix = f"cd {shlex.quote(str(cwd))} && " if cwd else "" + log(f"$ {prefix}{shlex.join(list(argv))}") + + +def last_lines(path: pathlib.Path, count: int) -> str: + if not path.exists(): + return "" + return "\n".join(path.read_text(encoding="utf-8", errors="replace").splitlines()[-count:]) + + +def last_nonempty_line(path: pathlib.Path) -> str: + for line in reversed(last_lines(path, 20).splitlines()): + if line.strip(): + return line.strip()[:300] + return "no output yet" diff --git a/integrations/agentbricks/tests/harness/project.py b/integrations/agentbricks/tests/harness/project.py new file mode 100644 index 000000000..ebd42e135 --- /dev/null +++ b/integrations/agentbricks/tests/harness/project.py @@ -0,0 +1,77 @@ +"""One test's Agent Bricks project: its run-scoped name and path, its files, and its runtime pin. + +Tests create the project themselves with ``agentbricks init`` and then call ``pin_runtime``, so the +project runs the Agent Bricks runtime under test instead of the release on PyPI. +""" + +from __future__ import annotations + +import dataclasses +import pathlib +import shutil +import uuid +from collections.abc import Mapping + +import tomli +import tomlkit + +from harness.logs import log +from harness.run import APP_PREFIX, E2EError, RunConfig, project_prefix + + +@dataclasses.dataclass(frozen=True) +class Project: + path: pathlib.Path + app_name: str + # The runtime under test: the wheel built from this checkout, or the one --wheel gives. + wheel: pathlib.Path + + @classmethod + def allocate(cls, run_config: RunConfig) -> Project: + """A unique project path and App name under the run's prefix; nothing is created yet.""" + name = f"{project_prefix(run_config.run_id)}{uuid.uuid4().hex[:6]}" + path = run_config.output / "projects" / name + path.parent.mkdir(parents=True, exist_ok=True) + if run_config.wheel is None: + # Only the run_config fixture builds it; the bare run config is for cleanup. + raise E2EError("The run has no runtime wheel to pin projects to.") + return cls(path, f"{APP_PREFIX}{name}", run_config.wheel) + + def pin_runtime(self) -> None: + """Point the project's databricks-agentbricks dependency at the runtime under test.""" + _pin_project_wheel(self.path, self.wheel) + + def is_pinned(self) -> bool: + pyproject = self.path / "pyproject.toml" + document = tomli.loads(pyproject.read_text(encoding="utf-8")) + sources = document.get("tool", {}).get("uv", {}).get("sources", {}) + return "databricks-agentbricks" in sources + + def write_file(self, relative_path: str, content: str) -> None: + """Write a user-owned file in the project, as a developer would.""" + target = self.path / relative_path + target.parent.mkdir(parents=True, exist_ok=True) + log(f"# write {target}: user-owned file") + target.write_text(content, encoding="utf-8") + + def write_files(self, files: Mapping[str, str]) -> None: + for relative_path, content in files.items(): + self.write_file(relative_path, content) + + +def _pin_project_wheel(project: pathlib.Path, wheel: pathlib.Path) -> None: + vendor_dir = project / "agentbricks_e2e_wheels" + vendor_dir.mkdir() + vendored_wheel = vendor_dir / wheel.name + shutil.copy2(wheel, vendored_wheel) + pyproject = project / "pyproject.toml" + document = tomlkit.parse(pyproject.read_text(encoding="utf-8")) + tool = document.setdefault("tool", {}) + uv = tool.setdefault("uv", {}) + sources = uv.setdefault("sources", {}) + sources["databricks-agentbricks"] = {"path": vendored_wheel.relative_to(project).as_posix()} + log( + f"# write {pyproject}: " + f"pin databricks-agentbricks runtime to {vendored_wheel.relative_to(project)}" + ) + pyproject.write_text(tomlkit.dumps(document), encoding="utf-8") diff --git a/integrations/agentbricks/tests/harness/run.py b/integrations/agentbricks/tests/harness/run.py new file mode 100644 index 000000000..f5acbf7f9 --- /dev/null +++ b/integrations/agentbricks/tests/harness/run.py @@ -0,0 +1,58 @@ +"""One run's identity: ``RunConfig``, the names every resource of the run carries, and constants. + +``suites/conftest.py`` builds one ``RunConfig`` per pytest run; fixtures and tests read it through the +``run_config`` fixture. Resource names start with the run id, so leftover cleanup can find them. +""" + +from __future__ import annotations + +import dataclasses +import datetime as dt +import pathlib +import re +import tempfile + +E2E_MODEL = "system.ai.gpt-5-2" + + +APP_PREFIX = "agent-bricks-" + + +class E2EError(RuntimeError): + """A reproducible setup or execution failure.""" + + +def now() -> dt.datetime: + return dt.datetime.now(dt.timezone.utc) + + +def project_prefix(run_id: str) -> str: + """Names every project of a run, so leftover cleanup can find the Apps a crashed test left.""" + return f"t-{run_id}-" + + +# App names are capped at 30 characters: "agent-bricks-t-" + run id + "-" + 6 random characters. +RUN_ID_PATTERN = re.compile(r"[a-z0-9]{1,8}") + + +def run_output_dir(run_id: str) -> pathlib.Path: + """Where a run writes its logs, projects and ``apps.jsonl``; leftover cleanup finds it by the id alone.""" + return pathlib.Path(tempfile.gettempdir()) / f"agentbricks-e2e-{run_id}" + + +def scratch_schema_prefix(run_id: str) -> str: + """Names every temporary schema of a run, so leftover cleanup can find ones a test left behind.""" + return f"ab_e2e_{run_id}_" + + +@dataclasses.dataclass(frozen=True) +class RunConfig: + """Values shared by every fixture of one pytest run, including all xdist workers.""" + + run_id: str + output: pathlib.Path + # Operator override; None lets the workspace pick the first running warehouse. + warehouse_id: str | None + # What every invocation asks for; workspaces name their models differently. + model: str + wheel: pathlib.Path | None diff --git a/integrations/agentbricks/tests/harness/target.py b/integrations/agentbricks/tests/harness/target.py new file mode 100644 index 000000000..857b345ed --- /dev/null +++ b/integrations/agentbricks/tests/harness/target.py @@ -0,0 +1,102 @@ +"""The Databricks workspace the system under test talks to (``--workspace live|fake``). + +``live`` is a real workspace. ``fake`` is a local stand-in with no SDK access, so every requirement +that needs a real workspace is skipped rather than failed. +""" + +from __future__ import annotations + +import enum +from collections.abc import Mapping, Sequence +from typing import Protocol + +import pytest + +from harness.run import RunConfig +from harness.workspace import Workspace + + +class WorkspaceKind(str, enum.Enum): + LIVE = "live" + FAKE = "fake" + + def __str__(self) -> str: + return self.value + + +class TargetWorkspaceUnavailable(Exception): + """The selected workspace cannot run in this environment; the suite skips instead of failing.""" + + +class TargetWorkspace(Protocol): + name: WorkspaceKind + + @property + def env(self) -> Mapping[str, str]: + """Environment variables the ``agentbricks`` and ``databricks`` subprocesses need.""" + ... + + @property + def cli_args(self) -> Sequence[str]: + """Global arguments placed after the CLI binary, e.g. ``--profile``.""" + ... + + @property + def client(self) -> Workspace | None: + """SDK access for cleanup and UC setup; None for the fake workspace.""" + ... + + +class LiveWorkspace: + name = WorkspaceKind.LIVE + + def __init__(self, profile: str | None, run_config: RunConfig): + self.profile = profile + self._run_config = run_config + self._client: Workspace | None = None + + @property + def env(self) -> Mapping[str, str]: + # The CLI looks up OAuth tokens by host, so several profiles on one host need the pin. + return {"DATABRICKS_CONFIG_PROFILE": self.profile} if self.profile else {} + + @property + def cli_args(self) -> Sequence[str]: + return ["--profile", self.profile] if self.profile else [] + + @property + def client(self) -> Workspace: + if self._client is None: + self._client = Workspace( + self.profile, + warehouse_id=self._run_config.warehouse_id, + ) + return self._client + + +class FakeWorkspace: + """Seam for the in-process fake: DATABRICKS_HOST/DATABRICKS_TOKEN env, no CLI args, no Workspace.""" + + name = WorkspaceKind.FAKE + env: Mapping[str, str] = {} + cli_args: Sequence[str] = () + client: Workspace | None = None + + def __init__(self, run_config: RunConfig): + raise TargetWorkspaceUnavailable("fake workspace not implemented yet") + + +def require_workspace(target_workspace: TargetWorkspace, what: str) -> Workspace: + """The live workspace's SDK client; skips the calling test or fixture on any other target.""" + workspace = target_workspace.client + if workspace is None: + pytest.skip(f"{what} needs a live workspace; got --workspace {target_workspace.name}") + return workspace + + +def make_target_workspace( + kind: WorkspaceKind, profile: str | None, run_config: RunConfig +) -> TargetWorkspace: + if kind is WorkspaceKind.FAKE: + return FakeWorkspace(run_config) + return LiveWorkspace(profile, run_config) diff --git a/integrations/agentbricks/tests/harness/workspace.py b/integrations/agentbricks/tests/harness/workspace.py new file mode 100644 index 000000000..35c1c7f3e --- /dev/null +++ b/integrations/agentbricks/tests/harness/workspace.py @@ -0,0 +1,270 @@ +"""Databricks SDK helpers for what is not under test: the prerequisites the suites need, and cleanup. + +Tests drive and inspect the agent through the ``agentbricks`` CLI. This module is for the rest: +UC objects, the warehouse and Genie space a tool needs, and, under "Workarounds", the cleanup no +public command does yet. ``Workspace`` holds one operation per method; deciding what to delete and +in which order is ``harness.cleanup``. +""" + +from __future__ import annotations + +import dataclasses +import datetime as dt +import io +import pathlib +import re +import subprocess +import time +from collections.abc import Collection +from typing import Any, cast + +from databricks.sdk import WorkspaceClient +from databricks.sdk.errors import DatabricksError, NotFound +from databricks.sdk.service.catalog import VolumeType +from databricks.sdk.service.sql import ExecuteStatementRequestOnWaitTimeout, State, StatementState + +from harness.logs import log +from harness.run import E2EError + +_BRANCH = re.compile(r"projects/[^/]+/branches/[^/]+") + + +@dataclasses.dataclass(frozen=True) +class AppIdentity: + """What cleanup needs to find the App's Lakebase role once the App and its Runtime Store are gone.""" + + service_principal: str | None = None + branch: str | None = None + + def merged(self, other: AppIdentity) -> AppIdentity: + """This identity with any missing field taken from ``other``.""" + return AppIdentity( + self.service_principal or other.service_principal, + self.branch or other.branch, + ) + + +class Workspace: + def __init__( + self, + profile: str | None, + *, + warehouse_id: str | None = None, + ): + self.profile = profile + self.warehouse_id = warehouse_id + self.client = WorkspaceClient(profile=profile) + self._warehouse_started = False + self._user_name: str | None = None + + # Existing resources the suite depends on + + def require_catalog(self, name: str) -> None: + try: + self.client.catalogs.get(name) + except DatabricksError as exc: + raise E2EError(f"Catalog {name!r} is not accessible: {exc}") from exc + + def require_genie_space(self, space_id: str) -> None: + try: + self.client.genie.get_space(space_id) + except DatabricksError as exc: + raise E2EError(f"Genie space {space_id!r} is not accessible: {exc}") from exc + + # SQL, warehouse, files + + def start_warehouse(self, override: str | None = None) -> str: + if override: + self.warehouse_id = override + else: + warehouses = list(self.client.warehouses.list()) + if not warehouses: + raise E2EError("The workspace has no SQL warehouse available for UC setup.") + running = next( + (item for item in warehouses if item.state == State.RUNNING), warehouses[0] + ) + self.warehouse_id = str(running.id) + log(f"# start warehouse {self.warehouse_id}") + self.client.warehouses.start_and_wait(self.warehouse_id, timeout=dt.timedelta(minutes=20)) + return self.warehouse_id + + def _ensure_warehouse(self) -> str: + # Each xdist worker has its own Workspace, so the warehouse is started on first use. + if not self._warehouse_started: + self.start_warehouse(self.warehouse_id) + self._warehouse_started = True + assert self.warehouse_id is not None + return self.warehouse_id + + def sql(self, statement: str, *, timeout: float = 600) -> None: + warehouse_id = self._ensure_warehouse() + log(f"$ sql: {statement}") + response = self.client.statement_execution.execute_statement( + statement=statement, + warehouse_id=warehouse_id, + wait_timeout="30s", + on_wait_timeout=ExecuteStatementRequestOnWaitTimeout.CONTINUE, + ) + statement_id = response.statement_id + while response.status and response.status.state in { + StatementState.PENDING, + StatementState.RUNNING, + }: + if not statement_id: + raise E2EError(f"SQL response has no statement_id: {response}") + if timeout <= 0: + raise E2EError(f"SQL statement timed out: {statement_id}") + time.sleep(10) + timeout -= 10 + response = self.client.statement_execution.get_statement(statement_id) + if not response.status or response.status.state != StatementState.SUCCEEDED: + raise E2EError(f"SQL failed: {response.as_dict()}") + + def upload(self, path: str, data: bytes) -> None: + self.client.files.upload(path, io.BytesIO(data), overwrite=True) + log(f"# uploaded {path}") + + # Temporary UC objects + + def create_schema(self, catalog: str, name: str, *, remove_after: dt.datetime) -> str: + """Create a schema and return its full name; ``RemoveAfter`` lets a workspace janitor delete it later.""" + info = self.client.schemas.create( + name, + catalog, + comment="Temporary Agent Bricks E2E schema; safe to delete", + properties={"RemoveAfter": remove_after.strftime("%Y%m%d%H")}, + ) + log(f"# created schema {info.full_name}") + return f"{catalog}.{name}" + + def delete_schema(self, full_name: str) -> None: + """Delete a schema with everything in it, grants included; a missing schema is fine.""" + try: + self.client.schemas.delete(full_name, force=True) + except NotFound: + return + log(f"# deleted schema {full_name}") + + def schemas_with_prefix(self, catalog: str, prefix: str) -> list[str]: + return [ + info.full_name + for info in self.client.schemas.list(catalog) + if info.full_name and info.name and info.name.startswith(prefix) + ] + + def create_volume(self, catalog: str, schema: str, name: str) -> None: + self.client.volumes.create(catalog, schema, name, VolumeType.MANAGED) + log(f"# created volume {catalog}.{schema}.{name}") + + # Apps + + def app_logs(self, name: str, log_path: pathlib.Path) -> pathlib.Path | None: + """Write the App's recent logs to ``log_path``; the SDK has no Apps log API.""" + argv = ["databricks", "apps", "logs", name, "--tail-lines", "200"] + if self.profile: + argv += ["--profile", self.profile] + try: + result = subprocess.run(argv, text=True, capture_output=True, timeout=120, check=False) + content = result.stdout if result.returncode == 0 else result.stderr or result.stdout + except Exception as exc: + content = f"Could not retrieve App logs: {exc}\n" + try: + log_path.parent.mkdir(parents=True, exist_ok=True) + log_path.write_text(content, encoding="utf-8") + except OSError as exc: + log(f"App runtime log capture warning for {name}: {exc}") + return None + log(f"App runtime logs captured: {log_path}") + return log_path + + # Workarounds for the end-of-run cleanup: what `agentbricks deployments delete` leaves behind. + + def runtime_store(self, app_name: str) -> dict[str, Any]: + return cast( + dict[str, Any], + self.client.api_client.do("GET", f"/api/2.0/agents/runtime-stores/{app_name}"), + ) + + def lakebase_branch(self, app_name: str) -> str | None: + """The Lakebase branch of the App's Runtime Store, if it still exists; never raises.""" + try: + lakebase = self.runtime_store(app_name).get("storage_backend", {}).get("lakebase", {}) + except NotFound: + return None + except Exception as exc: + log(f"cleanup warning | Runtime Store lookup for {app_name} | {exc}") + return None + branch = lakebase.get("branch") + return branch if isinstance(branch, str) and _BRANCH.fullmatch(branch) else None + + def roles_of(self, branch: str, principals: Collection[str]) -> list[str]: + """Resource names of the branch's roles for these principals. + + Role ids are the principal's UUID, sometimes behind a prefix such as ``agents-``. + """ + wanted = [principal for principal in principals if principal] + return [ + role.name + for role in self.client.postgres.list_roles(parent=branch) + if role.name and any(role.name.rsplit("/", 1)[-1].endswith(item) for item in wanted) + ] + + def delete_role(self, name: str) -> None: + try: + _wait(self.client.postgres.delete_role(name=name)) + except NotFound: + pass + + # Deployment sources: `agentbricks deploy` syncs them into the user's workspace home and + # `apps delete` leaves them behind. + + @property + def user_name(self) -> str: + if self._user_name is None: + name = self.client.current_user.me().user_name + if not name: + raise E2EError("The current user has no user_name.") + self._user_name = name + return self._user_name + + def _deployments_dir(self) -> str: + return f"/Workspace/Users/{self.user_name}/agentbricks_deployments" + + def experiments_with_prefix(self, prefix: str) -> list[str]: + """Ids of the experiments whose full path starts with ``prefix``.""" + return [ + str(experiment.experiment_id) + for experiment in self.client.experiments.search_experiments( + filter=f"name LIKE '{prefix}%'" + ) + if experiment.experiment_id + ] + + def delete_experiment(self, experiment_id: str) -> None: + try: + self.client.experiments.delete_experiment(experiment_id) + except NotFound: + pass + + def deployment_sources_with_prefix(self, prefix: str) -> list[str]: + try: + entries = list(self.client.workspace.list(self._deployments_dir())) + except NotFound: + return [] + return [ + entry.path + for entry in entries + if entry.path and entry.path.rsplit("/", 1)[-1].startswith(prefix) + ] + + def delete_workspace_path(self, path: str) -> None: + try: + self.client.workspace.delete(path, recursive=True) + except NotFound: + pass + + +def _wait(operation: Any) -> None: + wait = getattr(operation, "wait", None) + if callable(wait): + wait() diff --git a/integrations/agentbricks/tests/suites/README.md b/integrations/agentbricks/tests/suites/README.md new file mode 100644 index 000000000..b66a0d1bb --- /dev/null +++ b/integrations/agentbricks/tests/suites/README.md @@ -0,0 +1,79 @@ +# Agent Bricks E2E suites + +Black-box tests of the `agentbricks` CLI, at the top of the test pyramid: one happy-path scenario +per command group, each running the real CLI and checking the agent it produces. Error cases and +edge cases belong in the unit tests. + +## Running + +From `integrations/agentbricks`: + +```bash +uv run pytest tests/suites -n 6 --dist loadgroup \ + --databricks-profile --catalog +``` + +- Narrow it to a file (`tests/suites/test_deploy.py`) or a test. +- `-n 6 --dist loadgroup` runs tests in parallel; `loadgroup` keeps the dev tests on one worker, + because local dev servers share a proxy port. +- Workspace resources the suite borrows but never creates are flags: `--catalog`, + `--warehouse-id`, and `--genie-space-id` (without one the Genie tool is skipped). +- To set flags once for a shell or a CI job, use pytest's own `PYTEST_ADDOPTS`, e.g. + `export PYTEST_ADDOPTS="--databricks-profile --catalog "`. +- `--model` sets the model every invocation asks for (workspaces name models differently, e.g. + `system.ai.databricks-gpt-5-2` on staging). +- `--workspace fake` runs against no workspace; everything that needs one is skipped. +- `--keep-resources` keeps leftovers for debugging; each test still deletes its App and stores. + +The run id and output directory (logs, generated projects, `apps.jsonl`) are printed at the start: +`Agent Bricks E2E output: /tmp/agentbricks-e2e- (run )`. + +## Cleanup + +Every resource a run creates carries its run id, and is deleted at three points: + +1. **When each test ends**, the test's own cleanup runs. Each test registers the delete command + for everything it creates with the `cleanup` fixture, like Go's `t.Cleanup`, e.g. + `cleanup(f"deployments delete {app} --yes")`; registered commands run newest first, pass or fail. + Fixtures delete what they create, such as `scratch_schema` and its schema. +2. **When the run ends**, `pytest_sessionfinish` deletes everything carrying the run's name prefix + that is still there (stores, Apps, synced source folders, tracing experiments, schemas) and the + Apps' Lakebase roles; `deployments delete` leaves the folders, experiments and roles behind (`harness.cleanup.delete_run_leftovers`). A failure + there fails the run. +3. **After a run that was killed** before step 2, for example by a CI timeout: + + ```bash + uv run pytest tests/suites --cleanup-only --run-id --databricks-profile --catalog + ``` + + CI passes its own `--run-id` and runs this in an `always()` step. + +## Layout + +| Path | What it is | +|---|---| +| `tests/suites/conftest.py` | All pytest wiring: options, hooks, and the fixtures every suite shares | +| `tests/harness/run.py` | `RunConfig`, the run id, and the name prefixes every resource carries | +| `tests/harness/target.py` | The `--workspace live\|fake` seam the CLI and fixtures talk to | +| `tests/harness/agentbricks_cli.py` | Runs `agentbricks` commands: `cli.run("tools add ...")`, and `cli.start("dev ...")` in the background; `cli.invoke(app, prompt)` runs `endpoint invoke` with the `--model` the run asks for | +| `tests/harness/project.py` | A test's project: run-scoped path and App name, its files, and `pin_runtime()` | +| `tests/harness/workspace.py` | Databricks SDK helpers for what is not under test: prerequisites (UC objects, warehouse, Genie) and cleanup no CLI command does yet | +| `tests/harness/cleanup.py` | Deleting what the suites create: per test through the CLI, and per run for the rest | +| `tests/harness/logs.py` | Test-log output; every command is echoed so a run's log replays it | +| `tests/suites/test_.py` | One file per CLI command group: `test_tools`, `test_dev`, `test_deploy`, `test_memory`, `test_sessions`, `test_tracing` | +| `tests/suites/tool_definitions.py`, `uc_objects.py` | The tools `test_tools`, `test_dev` and `test_deploy` add, and the UC objects they use | + +`conftest.py` files are the only code pytest calls by name. `tests/harness/` sits next to `tests/suites/` +rather than inside it: it is plain Python the suites import as `harness` (via `pythonpath` in +`pyproject.toml`), not a suite. + +## Adding things + +- **A tool:** add one `ToolDefinition(name, requires, build)` in `tool_definitions.py`. + `requires` names the fixtures it needs; tools whose fixtures skip are left out and logged. +- **A command group:** add `test_.py` with its one happy path. A test takes the `project` fixture, runs + `agentbricks_cli.run(f"init ... {project.path}")`, then `project.pin_runtime()` so the project + runs the checkout's runtime; teardown fails a test that skips the pin. Register a `cleanup(...)` + delete command for everything the test creates, before the step that creates it if a failed step + can leave it behind. Resources the CLI under test does not create itself are fixtures in + `conftest.py`, created per test and deleted in the fixture's teardown. diff --git a/integrations/agentbricks/tests/suites/conftest.py b/integrations/agentbricks/tests/suites/conftest.py new file mode 100644 index 000000000..7e6406401 --- /dev/null +++ b/integrations/agentbricks/tests/suites/conftest.py @@ -0,0 +1,393 @@ +"""All pytest wiring for the E2E suites: options, hooks, and the fixtures the tests share. + +pytest calls everything in this file by name; ``harness/`` is plain code, so following its imports +answers what calls what. From ``integrations/agentbricks``: + + uv run pytest tests/suites -n 6 --dist loadgroup --databricks-profile + +``--workspace fake`` runs against no workspace; whatever needs one is skipped. + +Each test registers the delete commands for what it creates with ``cleanup``; fixtures such as +``scratch_schema`` delete what they create. After every live run, the controller also deletes +whatever is left (``harness.cleanup``). A run killed before that point is cleaned up afterwards +with ``--cleanup-only --run-id ``. +""" + +from __future__ import annotations + +import dataclasses +import datetime as dt +import os +import pathlib +import re +import subprocess +import uuid +from collections.abc import Iterator +from typing import Any + +import pytest +from harness.agentbricks_cli import AgentbricksCli, not_found +from harness.cleanup import ( + Cleanup, + app_identity, + delete_run_leftovers, + record_app_identity, +) +from harness.logs import log, log_command +from harness.project import Project +from harness.run import ( + E2E_MODEL, + RUN_ID_PATTERN, + E2EError, + RunConfig, + now, + run_output_dir, + scratch_schema_prefix, +) +from harness.target import ( + TargetWorkspace, + TargetWorkspaceUnavailable, + WorkspaceKind, + make_target_workspace, + require_workspace, +) +from harness.workspace import Workspace +from tool_definitions import TOOL_DEFINITIONS, Tool +from uc_objects import UcFunction, create_uc_function + +PACKAGE_DIR = pathlib.Path(__file__).resolve().parents[2] +_RUN_KEY = pytest.StashKey[RunConfig]() +_HANDOFF = "agentbricks_run" +_failed_tests: set[str] = set() + + +# Options + + +def pytest_addoption(parser: pytest.Parser) -> None: + group = parser.getgroup("agentbricks", "Agent Bricks E2E suites") + group.addoption( + "--workspace", + type=WorkspaceKind, + choices=list(WorkspaceKind), + default=WorkspaceKind.LIVE, + help="live runs against a Databricks workspace; fake skips whatever needs one.", + ) + group.addoption( + "--databricks-profile", + help="Databricks CLI profile; default is ambient credentials (DATABRICKS_CONFIG_PROFILE).", + ) + group.addoption( + "--catalog", + default="main", + help="Existing catalog for temporary test schemas (default: main).", + ) + group.addoption( + "--model", + default=E2E_MODEL, + help=f"Model every invocation asks for (default: {E2E_MODEL}).", + ) + group.addoption( + "--wheel", + type=pathlib.Path, + help="Prebuilt databricks-agentbricks wheel (default: build one from this checkout).", + ) + group.addoption( + "--warehouse-id", + help="SQL warehouse for UC setup (default: first running).", + ) + group.addoption( + "--genie-space-id", + help="Existing Genie space for the Genie tool; without one that tool is skipped.", + ) + group.addoption( + "--run-id", + help="Names every resource of the run, so --cleanup-only can delete what a killed run left " + "(default: random). 1-8 lowercase letters or digits.", + ) + group.addoption( + "--keep-resources", + action="store_true", + help="Skip the end-of-run cleanup of leftovers; each test's own teardown still runs.", + ) + group.addoption( + "--cleanup-only", + action="store_true", + help="Run no tests; only delete what the run given by --run-id left behind.", + ) + + +def pytest_configure(config: pytest.Config) -> None: + run_id = config.getoption("run_id") + if run_id is not None and RUN_ID_PATTERN.fullmatch(run_id) is None: + raise pytest.UsageError("--run-id must be 1-8 lowercase letters or digits.") + if config.getoption("cleanup_only"): + if run_id is None: + raise pytest.UsageError("--cleanup-only needs the --run-id of the run to clean up.") + if config.getoption("keep_resources"): + raise pytest.UsageError("--cleanup-only and --keep-resources contradict each other.") + + +def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item]) -> None: + if config.getoption("cleanup_only"): + config.hook.pytest_deselected(items=list(items)) + items.clear() + + +# The run: one id and output directory, shared by the controller and every xdist worker + + +@pytest.hookimpl(optionalhook=True) +def pytest_configure_node(node: Any) -> None: + """xdist hook: runs on the controller for each worker before it starts.""" + run = _base_run_config(node.config) + node.workerinput[_HANDOFF] = {"run_id": run.run_id, "output": str(run.output)} + + +def _base_run_config(config: pytest.Config) -> RunConfig: + """The run's config without a built wheel; the same value in the controller and every worker.""" + if _RUN_KEY in config.stash: + return config.stash[_RUN_KEY] + handoff = getattr(config, "workerinput", {}).get(_HANDOFF) + if handoff: + run_id, output = handoff["run_id"], pathlib.Path(handoff["output"]) + else: + run_id = config.getoption("run_id") or uuid.uuid4().hex[:6] + output = run_output_dir(run_id) + option = config.getoption + run = RunConfig( + run_id=run_id, + output=output, + warehouse_id=option("warehouse_id"), + model=option("model"), + wheel=option("wheel"), + ) + config.stash[_RUN_KEY] = run + return run + + +def _build_target_workspace(config: pytest.Config) -> TargetWorkspace: + return make_target_workspace( + config.getoption("workspace"), + config.getoption("databricks_profile"), + _base_run_config(config), + ) + + +@pytest.fixture(scope="session") +def run_config(request: pytest.FixtureRequest) -> RunConfig: + run = _base_run_config(request.config) + run.output.mkdir(parents=True, exist_ok=True) + log(f"Agent Bricks E2E output: {run.output} (run {run.run_id})") + if run.wheel: + return run + # A directory per process: xdist workers build at the same time and must not share an output. + dist = run.output / "dist" / os.environ.get("PYTEST_XDIST_WORKER", "main") + return dataclasses.replace(run, wheel=_build_wheel(dist)) + + +def _build_wheel(dist: pathlib.Path) -> pathlib.Path: + """Build this checkout's wheel, which every generated project installs instead of PyPI's.""" + argv = ["uv", "build", "--wheel", "--out-dir", str(dist)] + log_command(argv, PACKAGE_DIR) + result = subprocess.run( + argv, cwd=PACKAGE_DIR, text=True, capture_output=True, timeout=600, check=False + ) + if result.returncode != 0: + raise E2EError(f"uv build failed ({result.returncode}):\n{result.stderr or result.stdout}") + wheels = sorted(dist.glob("*.whl")) + if not wheels: + raise E2EError("uv build produced no wheel") + return wheels[-1] + + +# The workspace + + +@pytest.fixture(scope="session") +def target_workspace(request: pytest.FixtureRequest) -> TargetWorkspace: + try: + return _build_target_workspace(request.config) + except TargetWorkspaceUnavailable as exc: + pytest.skip(str(exc)) + + +@pytest.fixture(scope="session") +def workspace_client(target_workspace: TargetWorkspace) -> Workspace: + return require_workspace(target_workspace, "Databricks SDK access") + + +@pytest.fixture(scope="session") +def catalog(request: pytest.FixtureRequest, target_workspace: TargetWorkspace) -> str: + workspace = require_workspace(target_workspace, "a UC catalog") + name = request.config.getoption("catalog") + workspace.require_catalog(name) + return name + + +# What the tools need, and the tools themselves + + +@pytest.fixture +def scratch_schema( + catalog: str, target_workspace: TargetWorkspace, run_config: RunConfig +) -> Iterator[str]: + """``catalog.schema`` of a schema private to this test, deleted with everything in it afterwards.""" + workspace = require_workspace(target_workspace, "a UC schema") + name = f"{scratch_schema_prefix(run_config.run_id)}{uuid.uuid4().hex[:6]}" + # RemoveAfter lets a workspace janitor delete schemas of runs that crashed before any cleanup. + full_name = workspace.create_schema(catalog, name, remove_after=now() + dt.timedelta(hours=2)) + try: + yield full_name + finally: + try: + workspace.delete_schema(full_name) + except Exception as exc: + log(f"cleanup warning | schema {full_name} | {exc}") + + +@pytest.fixture +def uc_function(scratch_schema: str, target_workspace: TargetWorkspace) -> UcFunction: + """The marker function and volume in ``scratch_schema``.""" + return create_uc_function(require_workspace(target_workspace, "a UC function"), scratch_schema) + + +@pytest.fixture(scope="session") +def genie_space(request: pytest.FixtureRequest, target_workspace: TargetWorkspace) -> str: + # A space needs a warehouse and seeded tables, so it is a long-lived one rather than per run. + space_id = request.config.getoption("genie_space_id") + if not space_id: + pytest.skip("no Genie space; pass --genie-space-id") + if re.fullmatch(r"[0-9a-f]{32}", space_id) is None: + raise E2EError(f"Genie space id {space_id!r} is not 32 lowercase hex characters.") + require_workspace(target_workspace, "a Genie space").require_genie_space(space_id) + return space_id + + +@pytest.fixture +def tools(request: pytest.FixtureRequest) -> tuple[Tool, ...]: + """The tools whose fixtures are all available; the rest are logged and left out.""" + available: list[Tool] = [] + for definition in TOOL_DEFINITIONS: + try: + resources = {name: request.getfixturevalue(name) for name in definition.requires} + except pytest.skip.Exception as skipped: + log(f"tool {definition.name} skipped | {skipped}") + continue + available.append(definition.build(**resources)) + if not available: + pytest.skip("no tool is available") + return tuple(available) + + +# The CLI under test, and deleting what each test created + + +def pytest_runtest_logreport(report: pytest.TestReport) -> None: + # Includes subtest reports, so the App's logs are captured when any subtest fails. + if report.failed: + _failed_tests.add(report.nodeid) + + +@pytest.fixture +def agentbricks_cli( + request: pytest.FixtureRequest, target_workspace: TargetWorkspace, run_config: RunConfig +) -> AgentbricksCli: + return AgentbricksCli(target_workspace, run_config, request.node.name) + + +@pytest.fixture +def cleanup(agentbricks_cli: AgentbricksCli) -> Iterator[Cleanup]: + """Register the delete commands for what the test creates; they run when it ends.""" + registered = Cleanup(agentbricks_cli) + yield registered + registered.run_all() + + +@pytest.fixture +def project( + request: pytest.FixtureRequest, + cleanup: Cleanup, + agentbricks_cli: AgentbricksCli, + target_workspace: TargetWorkspace, + run_config: RunConfig, +) -> Iterator[Project]: + """A run-scoped path and App name for the test to `agentbricks init`. + + Depends on ``cleanup`` so that its teardown runs first, while the test's App still exists. + """ + project = Project.allocate(run_config) + workspace = target_workspace.client + if workspace is not None: + # The test deletes whatever App has this name, so it must not belong to anyone else. + existing = agentbricks_cli.run( + f"--output json deployments get {project.app_name}", check=False + ) + if existing.returncode == 0: + raise E2EError(f"App {project.app_name} already exists; refusing to deploy over it.") + if not not_found(existing): + raise E2EError(f"Could not check App {project.app_name}: {existing.stderr}") + yield project + failed = request.node.nodeid in _failed_tests + if workspace is not None: + _record_app(agentbricks_cli, workspace, project.app_name, failed=failed) + # Otherwise a test that forgot the pin would pass against the release on PyPI. + if not failed and (project.path / "pyproject.toml").exists() and not project.is_pinned(): + raise E2EError( + f"{project.path} was never pinned to the runtime under test; call pin_runtime after init." + ) + + +def _record_app(cli: AgentbricksCli, workspace: Workspace, name: str, *, failed: bool) -> None: + """Keep what is lost once the test's cleanup deletes its App.""" + identity = app_identity(cli, workspace, name) + if identity.service_principal or identity.branch: + # Leftover cleanup deletes the App's Lakebase role from this record. + record_app_identity(cli.run_config.output, name, identity) + if failed and identity.service_principal: + workspace.app_logs(name, cli.logs_dir / f"deploy-runtime-{name}.log") + + +# Deleting what the whole run left behind + + +@pytest.hookimpl(trylast=True) +def pytest_sessionfinish(session: pytest.Session) -> None: + config = session.config + # Workers must not clean up: the controller outlives them and sees the last one finish. + if hasattr(config, "workerinput"): + return + cleanup_only = config.getoption("cleanup_only") + if cleanup_only: + run = _base_run_config(config) + else: + run = config.stash.get(_RUN_KEY, None) + # The output dir exists once any test has run, so its absence means nothing was created. + if run is None or not run.output.exists(): + return + if config.getoption("keep_resources"): + log("Resources retained (--keep-resources); skipping leftover cleanup.") + return + try: + target_workspace = _build_target_workspace(config) + workspace = target_workspace.client + except TargetWorkspaceUnavailable: + workspace = None + if workspace is None: + if cleanup_only: + raise pytest.UsageError("--cleanup-only needs a live workspace.") + return + if cleanup_only and not (run.output / "apps.jsonl").exists(): + log(f"no apps.jsonl in {run.output}; Lakebase leftovers of deleted Apps can't be found") + cli = AgentbricksCli(target_workspace, run, "cleanup") + failures = delete_run_leftovers( + cli, workspace, run.run_id, run.output, config.getoption("catalog") + ) + for failure in failures: + log(f"cleanup warning | {failure}") + if failures: + if session.exitstatus in (pytest.ExitCode.OK, pytest.ExitCode.NO_TESTS_COLLECTED): + session.exitstatus = pytest.ExitCode.TESTS_FAILED + elif cleanup_only: + # Every test was deselected, which pytest otherwise reports as "no tests collected". + session.exitstatus = pytest.ExitCode.OK diff --git a/integrations/agentbricks/tests/suites/test_deploy.py b/integrations/agentbricks/tests/suites/test_deploy.py new file mode 100644 index 000000000..e9fc72d7c --- /dev/null +++ b/integrations/agentbricks/tests/suites/test_deploy.py @@ -0,0 +1,43 @@ +"""`agentbricks deploy`: every tool works on the deployed App. + +`agentbricks deploy` returns once the App is running its new source, so the test invokes it at once. +""" + +from __future__ import annotations + +import json + +from harness.agentbricks_cli import AgentbricksCli +from harness.cleanup import Cleanup +from harness.project import Project +from tool_definitions import Tool + + +def test_deploy( + agentbricks_cli: AgentbricksCli, + project: Project, + cleanup: Cleanup, + tools: tuple[Tool, ...], + subtests, +) -> None: + agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) + project.pin_runtime() + app = project.app_name + # Before deploy: a deploy that fails partway can still leave the App behind. + cleanup(f"deployments delete {app} --yes") + for tool in tools: + if tool.add_command: + agentbricks_cli.run(f"tools add {tool.add_command} --source {project.path}") + project.write_files(tool.files) + + deployed = json.loads( + agentbricks_cli.run( + f"--output json deploy {app} --source {project.path}", timeout=2400 + ).stdout + ) + cleanup(f"memory stores delete memory-stores/{deployed['memory_store']} --yes") + cleanup(f"sessions stores delete {deployed['session_store']} --yes") + for tool in tools: + with subtests.test(tool=tool.name): + output = agentbricks_cli.invoke(app, tool.prompt) + tool.check(output) diff --git a/integrations/agentbricks/tests/suites/test_dev.py b/integrations/agentbricks/tests/suites/test_dev.py new file mode 100644 index 000000000..e507f625b --- /dev/null +++ b/integrations/agentbricks/tests/suites/test_dev.py @@ -0,0 +1,38 @@ +"""`agentbricks dev`: every tool added with `tools add` works on the local dev server. + +Each dev server gets a free port, passed as both --app-port and PORT: the template server binds +$PORT, so an inherited PORT would override --app-port. +""" + +from __future__ import annotations + +import pytest +from harness.agentbricks_cli import AgentbricksCli, free_port +from harness.project import Project +from tool_definitions import Tool + +# Local dev servers share run-local's proxy port, so all dev runs stay on one worker, in order. +pytestmark = pytest.mark.xdist_group("dev") + + +def test_dev( + agentbricks_cli: AgentbricksCli, project: Project, tools: tuple[Tool, ...], subtests +) -> None: + agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) + project.pin_runtime() + for tool in tools: + if tool.add_command: + agentbricks_cli.run(f"tools add {tool.add_command} --source {project.path}") + project.write_files(tool.files) + + port = free_port() + with agentbricks_cli.start( + f"dev --source {project.path} --app-port {port}", env={"PORT": str(port)} + ) as server: + server.wait_until_serving(port) + for tool in tools: + with subtests.test(tool=tool.name): + output = agentbricks_cli.invoke( + f"--url http://127.0.0.1:{port} --no-auth", tool.prompt + ) + tool.check(output) diff --git a/integrations/agentbricks/tests/suites/test_memory.py b/integrations/agentbricks/tests/suites/test_memory.py new file mode 100644 index 000000000..131668da0 --- /dev/null +++ b/integrations/agentbricks/tests/suites/test_memory.py @@ -0,0 +1,54 @@ +"""`agentbricks memory`: the deployed agent remembers a fact in one session and recalls it in another. + +Long-term memory is keyed by actor, so both invocations name the same one; the fact is random, so +only the memory store can supply it. +""" + +from __future__ import annotations + +import json +import uuid + +import pytest +from harness.agentbricks_cli import AgentbricksCli, excerpt, tool_results +from harness.cleanup import Cleanup +from harness.project import Project + + +# Deploy needs a live workspace. +@pytest.mark.usefixtures("workspace_client") +def test_memory(agentbricks_cli: AgentbricksCli, project: Project, cleanup: Cleanup) -> None: + agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) + project.pin_runtime() + app = project.app_name + # Before deploy: a deploy that fails partway can still leave the App behind. + cleanup(f"deployments delete {app} --yes") + deployed = json.loads( + agentbricks_cli.run( + f"--output json deploy {app} --source {project.path}", timeout=2400 + ).stdout + ) + cleanup(f"memory stores delete memory-stores/{deployed['memory_store']} --yes") + cleanup(f"sessions stores delete {deployed['session_store']} --yes") + + actor = f"e2e-{uuid.uuid4().hex[:8]}" + codename = f"AGENTBRICKS_MEMORY_{uuid.uuid4().hex[:12]}" + + # Each invocation is a new session; only the actor carries over. + remembered = agentbricks_cli.invoke( + app, + f"Remember this fact about me in long-term memory: my codename is {codename}.", + actor=actor, + ) + assert any( + r.get("name") == "remember" and r.get("status") == "success" + for r in tool_results(remembered) + ), f"The agent did not remember the fact: {excerpt(remembered)}" + + recalled = agentbricks_cli.invoke( + app, "What is my codename? Look it up in your long-term memory.", actor=actor + ) + assert any( + r.get("name") == "recall" and codename in str(r.get("content")) + for r in tool_results(recalled) + ), f"Long-term memory did not return the fact in a new session: {excerpt(recalled)}" diff --git a/integrations/agentbricks/tests/suites/test_sessions.py b/integrations/agentbricks/tests/suites/test_sessions.py new file mode 100644 index 000000000..aee7b208c --- /dev/null +++ b/integrations/agentbricks/tests/suites/test_sessions.py @@ -0,0 +1,46 @@ +"""`agentbricks sessions`: the deployed agent keeps a conversation's history, and only that one's. + +The word is random, so only the session store can supply it. +""" + +from __future__ import annotations + +import json +import uuid + +import pytest +from harness.agentbricks_cli import AgentbricksCli, excerpt, reply +from harness.cleanup import Cleanup +from harness.project import Project + +QUESTION = "What secret word did I tell you earlier in this conversation? Reply with just the word." + + +# Deploy needs a live workspace. +@pytest.mark.usefixtures("workspace_client") +def test_sessions(agentbricks_cli: AgentbricksCli, project: Project, cleanup: Cleanup) -> None: + agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) + project.pin_runtime() + app = project.app_name + # Before deploy: a deploy that fails partway can still leave the App behind. + cleanup(f"deployments delete {app} --yes") + deployed = json.loads( + agentbricks_cli.run( + f"--output json deploy {app} --source {project.path}", timeout=2400 + ).stdout + ) + cleanup(f"memory stores delete memory-stores/{deployed['memory_store']} --yes") + cleanup(f"sessions stores delete {deployed['session_store']} --yes") + + session = str(uuid.uuid4()) + word = f"AGENTBRICKS_SESSION_{uuid.uuid4().hex[:12]}" + + agentbricks_cli.invoke( + app, + f"The secret word for this conversation is {word}. Reply with just: noted.", + session_id=session, + ) + same = agentbricks_cli.invoke(app, QUESTION, session_id=session) + assert word in reply(same), f"The session lost its history: {excerpt(same)}" + other = agentbricks_cli.invoke(app, QUESTION) + assert word not in reply(other), f"Another session saw this one's history: {excerpt(other)}" diff --git a/integrations/agentbricks/tests/suites/test_tools.py b/integrations/agentbricks/tests/suites/test_tools.py new file mode 100644 index 000000000..12c9bc74f --- /dev/null +++ b/integrations/agentbricks/tests/suites/test_tools.py @@ -0,0 +1,38 @@ +"""`agentbricks tools`: after `tools remove`, the agent no longer has the tool. + +`tools add` is covered by `test_dev` and `test_deploy`, which invoke every added tool. +""" + +from __future__ import annotations + +import pytest +from harness.agentbricks_cli import AgentbricksCli, free_port +from harness.project import Project +from tool_definitions import Tool + +# Local dev servers share run-local's proxy port, so all dev runs stay on one worker, in order. +pytestmark = pytest.mark.xdist_group("dev") + + +def test_tools_remove( + agentbricks_cli: AgentbricksCli, project: Project, tools: tuple[Tool, ...], subtests +) -> None: + agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) + project.pin_runtime() + removed = [tool for tool in tools if tool.add_command and tool.tool_id] + for tool in removed: + agentbricks_cli.run(f"tools add {tool.add_command} --source {project.path}") + for tool in removed: + agentbricks_cli.run(f"tools remove {tool.tool_id} --source {project.path}") + + port = free_port() + with agentbricks_cli.start( + f"dev --source {project.path} --app-port {port}", env={"PORT": str(port)} + ) as server: + server.wait_until_serving(port) + for tool in removed: + with subtests.test(tool=tool.name): + output = agentbricks_cli.invoke( + f"--url http://127.0.0.1:{port} --no-auth", tool.prompt + ) + tool.check_not_called(output) diff --git a/integrations/agentbricks/tests/suites/test_tracing.py b/integrations/agentbricks/tests/suites/test_tracing.py new file mode 100644 index 000000000..ecf0314bf --- /dev/null +++ b/integrations/agentbricks/tests/suites/test_tracing.py @@ -0,0 +1,43 @@ +"""`agentbricks tracing`: an invocation of the deployed App leaves a trace in init's experiment. + +The experiment is the project's own, named after it, so any trace in it came from this test. +""" + +from __future__ import annotations + +import json +import time + +import pytest +from harness.agentbricks_cli import AgentbricksCli +from harness.cleanup import Cleanup +from harness.project import Project + + +# Deploy needs a live workspace. +@pytest.mark.usefixtures("workspace_client") +def test_tracing(agentbricks_cli: AgentbricksCli, project: Project, cleanup: Cleanup) -> None: + agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) + project.pin_runtime() + app = project.app_name + # Before deploy: a deploy that fails partway can still leave the App behind. + cleanup(f"deployments delete {app} --yes") + deployed = json.loads( + agentbricks_cli.run( + f"--output json deploy {app} --source {project.path}", timeout=2400 + ).stdout + ) + cleanup(f"memory stores delete memory-stores/{deployed['memory_store']} --yes") + cleanup(f"sessions stores delete {deployed['session_store']} --yes") + agentbricks_cli.invoke(app, "Reply with one word: hello.") + + # Traces are exported asynchronously, so a new one can take a little while to show up. + deadline = time.monotonic() + 180 + while True: + traces = json.loads( + agentbricks_cli.run(f"--output json tracing list --source {project.path}").stdout + ) + if traces or time.monotonic() > deadline: + break + time.sleep(15) + assert traces, "No trace in the project's experiment after invoking the App" diff --git a/integrations/agentbricks/tests/suites/tool_definitions.py b/integrations/agentbricks/tests/suites/tool_definitions.py new file mode 100644 index 000000000..7c4e20ee1 --- /dev/null +++ b/integrations/agentbricks/tests/suites/tool_definitions.py @@ -0,0 +1,169 @@ +"""The tools under test: how each is added, prompted and verified.""" + +from __future__ import annotations + +import dataclasses +import json +from collections.abc import Callable, Mapping +from typing import Any + +from harness.agentbricks_cli import excerpt, tool_results +from uc_objects import UcFunction + +PYTHON_MARKER_FILE = "agent/tools/matrix_marker.py" +PYTHON_MARKER_SOURCE = ( + "from langchain_core.tools import tool\n\n\n" + "@tool\n" + "def matrix_marker(value: str) -> str:\n" + ' """Return the deterministic AgentBricks E2E marker."""\n' + " return 'AGENTBRICKS_PYTHON_OK'\n" +) +UC_MARKER = "AGENTBRICKS_UC_OK:matrix" +PYTHON_MARKER = "AGENTBRICKS_PYTHON_OK" + + +# Whether one tool step of an `endpoint invoke` response is this tool returning its result. +Called = Callable[[dict[str, Any]], bool] + + +@dataclasses.dataclass(frozen=True) +class Tool: + name: str + prompt: str + called: Called + # What follows `agentbricks tools add`; None when the tool is a user-owned file instead. + add_command: str | None + # User-owned files the tool needs. + files: Mapping[str, str] + # The id `tools add` binds it under, which `tools remove` takes. + tool_id: str | None + + def check(self, serialized: str) -> None: + """Assert the `endpoint invoke --output json` output shows the tool ran and returned its result.""" + assert self._calls(serialized), f"No successful {self.name} result: {excerpt(serialized)}" + + def check_not_called(self, serialized: str) -> None: + """Assert the agent answered without the tool returning its result.""" + assert not self._calls(serialized), f"{self.name} still ran: {excerpt(serialized)}" + + def _calls(self, serialized: str) -> list[dict[str, Any]]: + return [ + r for r in tool_results(serialized) if r.get("status") == "success" and self.called(r) + ] + + +@dataclasses.dataclass(frozen=True) +class ToolDefinition: + """A tool and the fixtures it needs; ``build`` takes them as keyword arguments, by fixture name.""" + + name: str + requires: tuple[str, ...] + build: Callable[..., Tool] + + +def returned(marker: str) -> Called: + """The step returned ``marker``, which only the tool can produce.""" + return lambda r: marker in str(r.get("content")) + + +def web_search_called(r: dict[str, Any]) -> bool: + return r.get("name") == "web_search" and "https://" in str(r.get("content")) + + +def genie_called(r: dict[str, Any]) -> bool: + # Genie reports a failed or timed-out answer inside a successful tool step, as its status. + if r.get("name") not in ("genie_ask", "genie_poll"): + return False + content = r.get("content") + try: + message = json.loads(content) if isinstance(content, str) else content + except ValueError: + return False + return ( + isinstance(message, dict) + and message.get("status") == "COMPLETED" + and bool(message.get("attachments")) + ) + + +def sandbox_tool(uc_function: UcFunction) -> Tool: + return Tool( + name="sandbox", + prompt=( + "You must call the sandbox tool. In the sandbox, use Python to read the " + f"entire text file at {uc_function.volume_file_path}. Return only its exact contents; " + "do not fabricate them." + ), + called=returned(uc_function.volume_marker), + add_command=f"sandbox --scope volume:{uc_function.volume}", + files={}, + tool_id="sandbox", + ) + + +def web_search_tool(**_resources: object) -> Tool: + return Tool( + name="mcp", + prompt=( + "You must use a tool from the configured system.ai.web_search MCP server. " + "Search official Databricks documentation for Model Context Protocol, then " + "return the title and https URL of one result. Do not answer from memory." + ), + called=web_search_called, + add_command="mcp system.ai.web_search --auth app", + files={}, + tool_id="web_search", + ) + + +def python_tool() -> Tool: + return Tool( + name="python", + prompt=( + "You must call the matrix_marker Python tool with value 'matrix'. " + "Return its exact result." + ), + called=returned(PYTHON_MARKER), + add_command=None, + files={PYTHON_MARKER_FILE: PYTHON_MARKER_SOURCE}, + tool_id=None, + ) + + +def uc_function_tool(uc_function: UcFunction) -> Tool: + return Tool( + name="uc_function", + prompt=( + "You must call the Unity Catalog function tool that computes the Agent Bricks marker, " + "with value 'matrix'. Return the called tool's exact result." + ), + called=returned(UC_MARKER), + add_command=f"uc-function {uc_function.function} --name agentbricks_uc_marker", + files={}, + tool_id="agentbricks_uc_marker", + ) + + +def genie_tool(genie_space: str) -> Tool: + return Tool( + name="genie", + prompt=( + "You must call the genie_ask tool and ask what data is available in the " + "configured Genie space. Return a one-sentence summary based only on the tool " + "response." + ), + called=genie_called, + add_command=f"genie-agent {genie_space} --name genie --auth app", + files={}, + tool_id="genie", + ) + + +# MCP needs a workspace that serves system.ai.web_search; Python needs nothing. +TOOL_DEFINITIONS = ( + ToolDefinition("sandbox", ("uc_function",), sandbox_tool), + ToolDefinition("mcp", ("workspace_client",), web_search_tool), + ToolDefinition("python", (), python_tool), + ToolDefinition("uc_function", ("uc_function",), uc_function_tool), + ToolDefinition("genie", ("genie_space",), genie_tool), +) diff --git a/integrations/agentbricks/tests/suites/uc_objects.py b/integrations/agentbricks/tests/suites/uc_objects.py new file mode 100644 index 000000000..21c40deb8 --- /dev/null +++ b/integrations/agentbricks/tests/suites/uc_objects.py @@ -0,0 +1,51 @@ +"""The UC objects the tools suite uses, created fresh in each test's temporary schema.""" + +from __future__ import annotations + +import dataclasses +import uuid + +from harness.run import E2EError +from harness.workspace import Workspace + +FUNCTION = "uc_marker" +VOLUME = "ab_volume" +# The MCP server exposes a function as catalog__schema__function, capped at 64 characters. +MAX_TOOL_NAME_LENGTH = 64 + + +@dataclasses.dataclass(frozen=True) +class UcFunction: + """The scratch UC objects the tools use.""" + + function: str + volume: str + volume_file_path: str + volume_marker: str + + +def create_uc_function(workspace: Workspace, scratch_schema: str) -> UcFunction: + """Create the marker function and volume in ``scratch_schema``, whose deletion removes them.""" + catalog, _, schema = scratch_schema.partition(".") + exposed_tool_name = f"{scratch_schema}.{FUNCTION}".replace(".", "__") + if len(exposed_tool_name) > MAX_TOOL_NAME_LENGTH: + raise E2EError( + f"The UC function's MCP tool name would exceed {MAX_TOOL_NAME_LENGTH} characters: " + f"{exposed_tool_name!r}. Use a shorter catalog." + ) + workspace.sql( + f"CREATE FUNCTION `{catalog}`.`{schema}`.`{FUNCTION}`" + "(value STRING) RETURNS STRING " + "COMMENT 'Deterministic Agent Bricks E2E marker tool' " + "RETURN concat('AGENTBRICKS_UC_OK:', value)" + ) + workspace.create_volume(catalog, schema, VOLUME) + marker = f"AGENTBRICKS_VOLUME_{uuid.uuid4().hex}" + file_path = f"/Volumes/{catalog}/{schema}/{VOLUME}/marker.txt" + workspace.upload(file_path, marker.encode()) + return UcFunction( + function=f"{scratch_schema}.{FUNCTION}", + volume=f"{scratch_schema}.{VOLUME}", + volume_file_path=file_path, + volume_marker=marker, + )