From fbc714f84ddbd57b70d4409ce2772da2d57dd616 Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Wed, 7 Oct 2026 23:46:20 +0000 Subject: [PATCH 01/45] [agentbricks] add pytest matrix e2e suite + suite runner Co-authored-by: Isaac --- .../agentbricks/tests/agentbricks_cli.py | 632 ++++++++++ integrations/agentbricks/tests/common.py | 332 ++++++ integrations/agentbricks/tests/conftest.py | 53 + .../memory_store/bundle/databricks.yml | 26 + .../tests/matrices/memory_store/suite.yaml | 18 + .../memory_store/test_memory_store.py | 40 + .../session_store/bundle/databricks.yml | 26 + .../tests/matrices/session_store/suite.yaml | 18 + .../session_store/test_session_store.py | 39 + .../matrices/tools/bundle/databricks.yml | 26 + .../tests/matrices/tools/conftest.py | 11 + .../tests/matrices/tools/suite.yaml | 25 + .../tests/matrices/tools/test_deploy.py | 64 + .../tests/matrices/tools/test_dev.py | 46 + .../tests/matrices/tools/test_grants.py | 54 + .../tests/matrices/tools/test_selection.py | 59 + .../agentbricks/tests/matrices/tools/tools.py | 241 ++++ .../matrices/tracing/bundle/databricks.yml | 26 + .../tests/matrices/tracing/suite.yaml | 17 + .../tests/matrices/tracing/test_tracing.py | 47 + integrations/agentbricks/tests/run_suites.py | 1061 +++++++++++++++++ .../agentbricks/tests/workspace_client.py | 514 ++++++++ 22 files changed, 3375 insertions(+) create mode 100644 integrations/agentbricks/tests/agentbricks_cli.py create mode 100644 integrations/agentbricks/tests/common.py create mode 100644 integrations/agentbricks/tests/conftest.py create mode 100644 integrations/agentbricks/tests/matrices/memory_store/bundle/databricks.yml create mode 100644 integrations/agentbricks/tests/matrices/memory_store/suite.yaml create mode 100644 integrations/agentbricks/tests/matrices/memory_store/test_memory_store.py create mode 100644 integrations/agentbricks/tests/matrices/session_store/bundle/databricks.yml create mode 100644 integrations/agentbricks/tests/matrices/session_store/suite.yaml create mode 100644 integrations/agentbricks/tests/matrices/session_store/test_session_store.py create mode 100644 integrations/agentbricks/tests/matrices/tools/bundle/databricks.yml create mode 100644 integrations/agentbricks/tests/matrices/tools/conftest.py create mode 100644 integrations/agentbricks/tests/matrices/tools/suite.yaml create mode 100644 integrations/agentbricks/tests/matrices/tools/test_deploy.py create mode 100644 integrations/agentbricks/tests/matrices/tools/test_dev.py create mode 100644 integrations/agentbricks/tests/matrices/tools/test_grants.py create mode 100644 integrations/agentbricks/tests/matrices/tools/test_selection.py create mode 100644 integrations/agentbricks/tests/matrices/tools/tools.py create mode 100644 integrations/agentbricks/tests/matrices/tracing/bundle/databricks.yml create mode 100644 integrations/agentbricks/tests/matrices/tracing/suite.yaml create mode 100644 integrations/agentbricks/tests/matrices/tracing/test_tracing.py create mode 100644 integrations/agentbricks/tests/run_suites.py create mode 100644 integrations/agentbricks/tests/workspace_client.py diff --git a/integrations/agentbricks/tests/agentbricks_cli.py b/integrations/agentbricks/tests/agentbricks_cli.py new file mode 100644 index 000000000..83b727bd9 --- /dev/null +++ b/integrations/agentbricks/tests/agentbricks_cli.py @@ -0,0 +1,632 @@ +"""The ``agentbricks`` CLI and the agents it produces: the system under test. + +Every call here is a real ``agentbricks`` subprocess, so a regression in the CLI fails the test +that makes it. The only non-subprocess part is ``Agent.invoke``, which POSTs to a running agent. +Nothing in this module knows about any particular matrix; tests say what to bind. +""" + +from __future__ import annotations + +import concurrent.futures +import contextlib +import dataclasses +import json +import os +import pathlib +import re +import shlex +import shutil +import signal +import socket +import subprocess +import time +import urllib.error +import urllib.request +import uuid +from collections.abc import Callable, Iterator, Mapping, Sequence +from typing import Any, Protocol + +import tomli +import tomlkit +from common import ( + E2E_MODEL, + FRAMEWORKS, + Evidence, + Inputs, + MatrixError, + Transcript, + last_lines, + last_nonempty_line, + now, +) +from workspace_client import Workspace, cleanup_app + +# `agentbricks init` binds a session store, memory store, and tracing experiment by default. +_UNBIND_COMMANDS = { + "session_store": ("sessions", "unbind"), + "memory_store": ("memory", "unbind"), + "tracing": ("tracing", "unbind"), +} +_STORE_KINDS = {"memory": ("memory", "stores"), "sessions": ("sessions", "stores")} +_DIRECT_HEADER = ( + 'schema_version = 1\n\n[agent]\nframework = "{framework}"\nserver = "agentbricks"\n' +) + + +class ToolSpec(Protocol): + """What ``write_manifest`` needs from a tool: its [[tools]] TOML and any user-owned files.""" + + toml: str + files: Mapping[str, str] + + +@dataclasses.dataclass +class Project: + authoring: str + path: pathlib.Path + app_name: str + features: tuple[str, ...] = () + memory_store_config: str | None = None + session_store_config: str | None = None + memory_resource: dict[str, Any] | None = None + session_resource: dict[str, Any] | None = None + memory_store_name: str | None = None + session_store_name: str | None = None + experiment_name: str | None = None + # Set only once this project registered the App, so cleanup never deletes one it did not create. + app_registered: bool = False + + +class Agent: + """A running agent: a dev server or a deployed App.""" + + def __init__( + self, + base_url: str, + headers: Callable[[], dict[str, str]], + log_path: pathlib.Path, + transcript: Transcript, + *, + app_name: str | None = None, + app: dict[str, Any] | None = None, + ): + self.base_url = base_url + self.log_path = log_path + self.app_name = app_name + self.app = app + self._headers = headers + self._transcript = transcript + + @property + def runtime(self) -> str: + return "deploy" if self.app_name else "dev" + + def invoke(self, prompt: str, label: str = "invoke") -> str: + """POST the prompt to /api/invocations and return the serialized response.""" + url = f"{self.base_url}/api/invocations" + last: Exception | None = None + invocation_id = str(uuid.uuid4()) + body = { + "id": invocation_id, + # session_id is a top-level invocation field; the adapter rejects it nested in input. + "session_id": invocation_id, + "input": { + "model": E2E_MODEL, + "messages": [{"role": "user", "content": prompt}], + }, + } + for attempt in range(1, 4): + try: + response = _monitored( + label, + lambda: _http_json(url, body, self._headers()), + self._transcript, + timeout=360, + ) + return json.dumps(response, sort_keys=True, default=str) + except Exception as exc: + last = exc + self._transcript.write(f"attempt {attempt}/3 | {label} | {exc}") + if attempt < 3: + time.sleep(15) + raise MatrixError(f"{label} failed after 3 attempts: {last}") + + def curl(self, prompt: str) -> str: + """A redacted curl equivalent of ``invoke``, recorded as the evidence row's command.""" + auth = " -H 'Authorization: Bearer '" if self.app_name else "" + body = json.dumps( + { + "id": "", + "session_id": "", + "input": { + "model": E2E_MODEL, + "messages": [{"role": "user", "content": prompt}], + }, + } + ) + return ( + f"curl -sS -X POST {shlex.quote(f'{self.base_url}/api/invocations')}" + f" -H 'Content-Type: application/json'{auth} --data {shlex.quote(body)}" + ) + + +class AgentbricksCli: + """Authors, runs and deploys projects with the CLI, and deletes what it created on ``close``.""" + + def __init__( + self, + inputs: Inputs, + workspace: Workspace, + transcript: Transcript, + evidence: Evidence, + ): + self.inputs = inputs + self.workspace = workspace + self.transcript = transcript + self.evidence = evidence + self.logs_dir = inputs.output / "logs" + self._projects: list[Project] = [] + self._sequence = 0 + + # Process plumbing + + def argv(self, *args: str) -> list[str]: + profile = ( + ["--profile", self.inputs.databricks_profile] if self.inputs.databricks_profile else [] + ) + return [str(self.inputs.agentbricks_bin), *profile, *args] + + def cli( + self, + *args: str, + cwd: pathlib.Path | None = None, + timeout: float = 300, + check: bool = True, + ) -> subprocess.CompletedProcess[str]: + argv = self.argv(*args) + self.transcript.command(argv, cwd) + result = subprocess.run( + argv, cwd=cwd, text=True, capture_output=True, timeout=timeout, check=False + ) + if result.stdout.strip(): + self.transcript.write(result.stdout) + if result.stderr.strip(): + self.transcript.write(result.stderr) + if check and result.returncode != 0: + raise MatrixError( + f"Command failed ({result.returncode}): {shlex.join(argv)}\n" + f"{result.stderr or result.stdout}" + ) + return result + + def _run_long(self, label: str, argv: Sequence[str], *, timeout: float) -> None: + self.transcript.command(argv) + log_path = self.logs_dir / f"{label}.log" + log_path.parent.mkdir(parents=True, exist_ok=True) + started = time.monotonic() + with log_path.open("w", encoding="utf-8") as log_file: + process = subprocess.Popen( + list(argv), + text=True, + stdout=log_file, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + next_tick = 60.0 + while process.poll() is None: + elapsed = time.monotonic() - started + if elapsed >= timeout: + terminate_process_group(process, grace=0) + raise MatrixError(f"{label} timed out after {timeout:.0f}s; log: {log_path}") + if elapsed >= next_tick: + self.transcript.write( + f"tick {now():%H:%M} | {label} | running | {last_nonempty_line(log_path)}" + ) + next_tick += 60.0 + time.sleep(2) + self.transcript.write(log_path.read_text(encoding="utf-8", errors="replace")) + if process.returncode != 0: + raise MatrixError(f"{label} failed ({process.returncode}); log: {log_path}") + self.transcript.write(f"tick {now():%H:%M} | {label} | success") + + def _next(self) -> int: + self._sequence += 1 + return self._sequence + + # Authoring + + def new_project(self, authoring: str, keep: Sequence[str] = ()) -> Project: + """`agentbricks init` a uniquely named project, pin the runtime under test, and leave only + the ``keep`` features (session_store, memory_store, tracing) bound. + + CLI authoring unbinds every other default feature. Direct authoring replaces the whole + manifest later in ``write_manifest``, which binds ``keep`` itself. + """ + unknown = set(keep) - set(_UNBIND_COMMANDS) + if unknown: + raise MatrixError(f"Unknown features to keep: {sorted(unknown)}") + path = self.inputs.output / "projects" / f"t-{uuid.uuid4().hex[:8]}" + path.parent.mkdir(parents=True, exist_ok=True) + project = Project(authoring, path, f"agent-bricks-{path.name}", features=tuple(keep)) + self._projects.append(project) + self._run_long( + f"init-{path.name}", + self.argv("init", "--framework", FRAMEWORKS[0], str(path)), + timeout=600, + ) + if self.inputs.bridge_sha: + self._pin_bridge_sources(path) + else: + self._pin_project_wheel(path) + # Registered before any store exists so cleanup knows the app name if creation fails. + self.evidence.register_project(project.app_name, authoring=authoring) + if authoring == "cli": + for feature, command in _UNBIND_COMMANDS.items(): + if feature not in keep: + self.cli(*command, "--source", str(path)) + self._sync_manifest(project) + return project + + def tools_add( + self, project: Project, *args: str, check: bool = True, json_output: bool = False + ) -> subprocess.CompletedProcess[str]: + return self._tools("add", project, args, check, json_output) + + def tools_remove( + self, project: Project, *args: str, check: bool = True, json_output: bool = False + ) -> subprocess.CompletedProcess[str]: + return self._tools("remove", project, args, check, json_output) + + def _tools( + self, + verb: str, + project: Project, + args: Sequence[str], + check: bool, + json_output: bool, + ) -> subprocess.CompletedProcess[str]: + prefix = ("--output", "json") if json_output else () + return self.cli(*prefix, "tools", verb, *args, "--source", str(project.path), check=check) + + def manifest(self, project: Project) -> dict[str, Any]: + return tomli.loads((project.path / "agent.toml").read_text(encoding="utf-8")) + + def write_file(self, project: Project, relative_path: str, content: str) -> None: + """Write a user-owned file in the project, as a developer would.""" + target = project.path / relative_path + target.parent.mkdir(parents=True, exist_ok=True) + self.transcript.file_step(target, "user-owned file") + target.write_text(content, encoding="utf-8") + + def write_manifest(self, project: Project, tool_specs: Sequence[ToolSpec]) -> None: + """Direct authoring: write agent.toml with exactly these tools and the project's features.""" + sections = [_DIRECT_HEADER.format(framework=FRAMEWORKS[0])] + sections.extend(spec.toml for spec in tool_specs if spec.toml) + if "session_store" in project.features: + sections.append(f'[session_store]\nname = "{project.path.name}-sessions"\n') + if "memory_store" in project.features: + sections.append(f'[memory_store]\nname = "{project.path.name}-memory"\n') + if "tracing" in project.features: + # Follows default_experiment_name's shape so direct authoring carries the same + # tracing binding `agentbricks init` gives CLI authoring. + slug = re.sub(r"[^a-z0-9-]+", "-", project.path.name.lower()).strip("-") or "agent" + experiment = f"/Shared/agentbricks_traces/{slug}-{self.inputs.run_suffix}" + sections.append(f'[tracing]\nexperiment_name = "{experiment}"\n') + target = project.path / "agent.toml" + self.transcript.file_step(target, "direct authoring; no agentbricks tools command") + target.write_text("\n".join(sections), encoding="utf-8") + for spec in tool_specs: + for relative_path, content in spec.files.items(): + self.write_file(project, relative_path, content) + self._sync_manifest(project) + + def _sync_manifest(self, project: Project) -> None: + """Read the authored agent.toml and create the stores it declares.""" + manifest = self.manifest(project) + project.memory_store_config = manifest.get("memory_store", {}).get("name") + project.session_store_config = manifest.get("session_store", {}).get("name") + project.experiment_name = manifest.get("tracing", {}).get("experiment_name") + if project.memory_store_config and not project.memory_store_name: + result = self.cli( + "--output", + "json", + "memory", + "stores", + "create", + "--name", + project.memory_store_config, + ) + project.memory_resource = json.loads(result.stdout) + name = project.memory_resource.get("name") + if isinstance(name, str) and name: + project.memory_store_name = name + if project.session_store_config and not project.session_store_name: + result = self.cli( + "--output", + "json", + "sessions", + "stores", + "create", + "--name", + project.session_store_config, + ) + project.session_resource = json.loads(result.stdout) + project.session_store_name = project.session_store_config + self.evidence.register_project( + project.app_name, + authoring=project.authoring, + memory_store_name=project.memory_store_name, + session_store_name=project.session_store_name, + ) + + def _pin_project_wheel(self, project: pathlib.Path) -> None: + wheel = self.inputs.wheel + 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()} + self.transcript.file_step( + pyproject, + f"pin databricks-agentbricks runtime to {vendored_wheel.relative_to(project)}", + ) + pyproject.write_text(tomlkit.dumps(document), encoding="utf-8") + + def _pin_bridge_sources(self, project: pathlib.Path) -> None: + bridge_sha = self.inputs.bridge_sha + pyproject = project / "pyproject.toml" + document = tomlkit.parse(pyproject.read_text(encoding="utf-8")) + dependencies = document["project"]["dependencies"] + if not any( + str(dependency).startswith("databricks-langchain") for dependency in dependencies + ): + dependencies.append("databricks-langchain>=0.17.0") + if "tool" not in document: + document["tool"] = tomlkit.table() + if "uv" not in document["tool"]: + document["tool"]["uv"] = tomlkit.table() + if "sources" not in document["tool"]["uv"]: + document["tool"]["uv"]["sources"] = tomlkit.table() + for package, subdirectory in ( + ("databricks-agentbricks", "integrations/agentbricks"), + ("databricks-langchain", "integrations/langchain"), + ): + source = tomlkit.inline_table() + source.update( + { + "git": "https://github.com/databricks/databricks-ai-bridge.git", + "rev": bridge_sha, + "subdirectory": subdirectory, + } + ) + document["tool"]["uv"]["sources"][package] = source + self.transcript.file_step(pyproject, f"pin Agent Bricks and LangChain to {bridge_sha}") + pyproject.write_text(tomlkit.dumps(document), encoding="utf-8") + + # Running + + @contextlib.contextmanager + def dev(self, project: Project) -> Iterator[Agent]: + """`agentbricks dev` the project locally for the duration of the block.""" + label = f"dev-{project.path.name}-{self._next()}" + log_path = self.logs_dir / f"{label}.log" + port = _free_port() + argv = self.argv( + "dev", "--source", str(project.path), "--app-port", str(port), "--prepare-environment" + ) + log_path.parent.mkdir(parents=True, exist_ok=True) + self.transcript.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, + ) + try: + self._wait_for_local(process, port, label, log_path) + yield Agent(f"http://127.0.0.1:{port}", lambda: {}, log_path, self.transcript) + finally: + if process.poll() is None: + terminate_process_group(process) + self.transcript.write(f"{label} | stopped") + + def deploy(self, project: Project, app: str | None = None) -> Agent: + """`agentbricks deploy` the project; deploying again updates the same App.""" + name = app or project.app_name + if not project.app_registered: + self.workspace.assert_app_absent(name) + # Deploy can create the App and then fail while waiting for it, so register first. + self.evidence.register_app(name) + project.app_registered = True + label = f"deploy-{name}-{self._next()}" + self._run_long( + label, + self.argv("deploy", name, "--source", str(project.path)), + timeout=2400, + ) + deployed = self.workspace.wait_for_app(name) + url = str(deployed.get("url") or "").rstrip("/") + if not url: + raise MatrixError(f"App {name} has no URL: {deployed}") + return Agent( + url, + lambda: self.workspace.app_headers, + self.logs_dir / f"{label}.log", + self.transcript, + app_name=name, + app=deployed, + ) + + def _wait_for_local( + self, process: subprocess.Popen[str], port: int, label: str, log_path: pathlib.Path + ) -> None: + started = time.monotonic() + next_tick = 60.0 + while True: + if process.poll() is not None: + raise MatrixError( + f"{label} exited {process.returncode}: {last_lines(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 > 1200: + raise MatrixError(f"{label} did not become reachable: {last_lines(log_path, 30)}") + if elapsed >= next_tick: + self.transcript.write( + f"tick {now():%H:%M} | {label} | starting | {last_nonempty_line(log_path)}" + ) + next_tick += 60 + time.sleep(5) + + # Cleanup + + def delete_store(self, kind: str, name: str) -> subprocess.CompletedProcess[str]: + return delete_store( + self.inputs.agentbricks_bin, self.inputs.databricks_profile, self.transcript, kind, name + ) + + def close(self) -> None: + """Delete every App, store and role this CLI created; failures are left for the sweep.""" + for project in self._projects: + if project.app_registered: + self._attach_app_logs(project) + info = { + "memory_store_name": project.memory_store_name, + "session_store_name": project.session_store_name, + } + try: + results = cleanup_app( + self.workspace, + self.delete_store, + project.app_name, + info, + has_app=project.app_registered, + ) + except Exception as exc: + self.transcript.write(f"cleanup warning | {project.app_name} | {exc}") + continue + if any(result.get("status") == "failed" for result in results): + self.transcript.write( + f"cleanup incomplete | {project.app_name} | deferred to the runner sweep" + ) + continue + self.evidence.mark_cleaned(project.app_name, results) + + def _attach_app_logs(self, project: Project) -> None: + # The App is deleted next, so a failed cell must capture its logs while it exists. + failed = any( + row["status"] == "fail" and row.get("app_name") == project.app_name + for row in self.evidence.read_rows() + ) + if not failed: + return + log_path = self.workspace.app_logs( + project.app_name, self.logs_dir / f"deploy-runtime-{project.app_name}.log" + ) + if log_path is not None: + self.evidence.attach_artifact_to_failed_rows(project.app_name, log_path) + + +def delete_store( + agentbricks: pathlib.Path | str, + profile: str | None, + transcript: Transcript, + kind: str, + name: str, +) -> subprocess.CompletedProcess[str]: + """`agentbricks stores delete`; also used by the runner's final sweep.""" + command = _STORE_KINDS[kind] + argv = [ + str(agentbricks), + *(["--profile", profile] if profile else []), + *command, + "delete", + name, + "--yes", + ] + transcript.command(argv) + result = subprocess.run(argv, text=True, capture_output=True, timeout=600, check=False) + for text in (result.stdout, result.stderr): + if text.strip(): + transcript.write(text) + return result + + +def terminate_process_group(process: subprocess.Popen[str], *, grace: float = 20) -> None: + try: + os.killpg(process.pid, signal.SIGTERM) + except ProcessLookupError: + return + if grace <= 0: + return + try: + process.wait(timeout=grace) + except subprocess.TimeoutExpired: + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass + + +def _free_port() -> int: + with socket.socket() as probe: + probe.bind(("127.0.0.1", 0)) + return probe.getsockname()[1] + + +def _monitored( + label: str, + operation: Callable[[], dict[str, Any]], + transcript: Transcript, + *, + timeout: float, +) -> dict[str, Any]: + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(operation) + started = time.monotonic() + while True: + try: + return future.result( + timeout=min(60, max(1, timeout - (time.monotonic() - started))) + ) + except concurrent.futures.TimeoutError: + elapsed = time.monotonic() - started + transcript.write(f"tick {now():%H:%M} | {label} | running | {elapsed:.0f}s") + if elapsed >= timeout: + raise MatrixError(f"{label} timed out after {timeout:.0f}s") from None + + +def _http_json(url: str, body: dict[str, Any], headers: dict[str, str]) -> dict[str, Any]: + request = urllib.request.Request( + url, + data=json.dumps(body).encode(), + headers={"Content-Type": "application/json", **headers}, + method="POST", + ) + try: + with urllib.request.urlopen(request, timeout=340) as response: + payload = response.read().decode() + except urllib.error.HTTPError as exc: + detail = exc.read().decode(errors="replace") + raise MatrixError(f"HTTP {exc.code} from {url}: {detail}") from exc + try: + value = json.loads(payload) + except json.JSONDecodeError as exc: + raise MatrixError(f"Invalid JSON from {url}: {payload[:2000]}") from exc + if not isinstance(value, dict): + raise MatrixError(f"Expected object response from {url}, got {type(value).__name__}") + return value diff --git a/integrations/agentbricks/tests/common.py b/integrations/agentbricks/tests/common.py new file mode 100644 index 000000000..20f6dffaa --- /dev/null +++ b/integrations/agentbricks/tests/common.py @@ -0,0 +1,332 @@ +"""Matrix-agnostic pieces shared by the test modules and ``run_suites.py``. + +``run_suites.py`` provisions the workspace dependencies and passes their resolved values in a JSON +file named by ``AGENTBRICKS_TEST_INPUTS``. This module loads that file, owns the matrix constants +(so the runner and the tests cannot drift), and records the evidence rows the runner turns into +``evidence.json``. +""" + +from __future__ import annotations + +import contextlib +import dataclasses +import datetime as dt +import json +import os +import pathlib +import shlex +import sys +import threading +import time +from collections.abc import Iterator, Sequence +from typing import Any + +INPUTS_ENV = "AGENTBRICKS_TEST_INPUTS" + +FRAMEWORKS = ("langgraph",) +AUTHORING_PATHS = ("cli", "direct") +RUNTIMES = ("dev", "deploy") +TOOL_KINDS = ("sandbox", "mcp", "python", "uc_function", "genie") +E2E_MODEL = "system.ai.gpt-5-2" +TOOL_RESOURCE_PREFIX = "agentbricks-tool-" + +# The (authoring, runtime, kind) evidence cells each matrix must produce a passing row for. The +# runner's verify step reads this, keyed by suite name. +MATRIX_CELLS: dict[str, tuple[tuple[str, str, str], ...]] = { + "tools": tuple( + (authoring, runtime, kind) + for authoring in AUTHORING_PATHS + for runtime in RUNTIMES + for kind in TOOL_KINDS + ) + + tuple((authoring, "deploy", "selection") for authoring in AUTHORING_PATHS), + "session_store": (("cli", "deploy", "session_store"),), + "memory_store": (("cli", "deploy", "memory_store"),), + "tracing": (("cli", "deploy", "tracing"),), +} + +EXPECTED = { + "sandbox": "the exact hidden marker read from the temporary UC volume file", + "python": "AGENTBRICKS_PYTHON_OK", + "uc_function": "AGENTBRICKS_UC_OK:matrix", + "mcp": "a web-search tool call and a non-empty https result", + "genie": "a genie_ask tool call and a non-empty Genie response", + "selection": "the UC function tool's exact result, and no call to the Python marker tool", + "session_store": "the declared session store exists", + "memory_store": "the declared memory store exists", + "tracing": "the tracing experiment is bound", +} + + +class MatrixError(RuntimeError): + """A reproducible setup or execution failure.""" + + +def now() -> dt.datetime: + return dt.datetime.now(dt.timezone.utc) + + +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" + + +class Transcript: + def __init__(self, path: pathlib.Path): + self.path = path + self.path.parent.mkdir(parents=True, exist_ok=True) + self._lock = threading.Lock() + + def write(self, text: str) -> None: + line = text.rstrip() + "\n" + with self._lock: + with self.path.open("a", encoding="utf-8") as output: + output.write(line) + sys.stdout.write(line) + sys.stdout.flush() + + def command(self, argv: Sequence[str], cwd: pathlib.Path | None = None) -> None: + prefix = f"cd {shlex.quote(str(cwd))} && " if cwd else "" + self.write(f"$ {prefix}{shlex.join(list(argv))}") + + def file_step(self, path: pathlib.Path, description: str) -> None: + self.write(f"# write {path}: {description}") + + +@dataclasses.dataclass(frozen=True) +class Inputs: + """The workspace dependencies ``run_suites.py`` resolved for this suite run.""" + + suite: str + databricks_profile: str | None + app_auth_profile: str | None + preprovisioned_app_catalog_access: bool + bridge_sha: str | None + wheel: pathlib.Path + agentbricks_bin: pathlib.Path + output: pathlib.Path + run_suffix: str + catalog: str + scratch_schema: str + uc_function: str + transitive_uc_function: str + uc_volume: str + volume_marker: str + volume_file_path: str + genie_space_id: str + warehouse_id: str + + @classmethod + def from_env(cls) -> Inputs: + path = os.environ.get(INPUTS_ENV) + if not path: + raise MatrixError(f"{INPUTS_ENV} is not set; run these tests through run_suites.py.") + data = json.loads(pathlib.Path(path).read_text(encoding="utf-8")) + return cls( + suite=data["suite"], + databricks_profile=data.get("databricks_profile"), + app_auth_profile=data.get("app_auth_profile") or data.get("databricks_profile"), + preprovisioned_app_catalog_access=bool(data.get("preprovisioned_app_catalog_access")), + bridge_sha=data.get("bridge_sha"), + wheel=pathlib.Path(data["wheel"]), + agentbricks_bin=pathlib.Path(data["agentbricks_bin"]), + output=pathlib.Path(data["output_dir"]), + run_suffix=data["run_suffix"], + catalog=data["catalog"], + scratch_schema=data["schema"], + uc_function=data["uc_function"], + transitive_uc_function=data["transitive_uc_function"], + uc_volume=data["uc_volume"], + volume_marker=data["volume_marker"], + volume_file_path=data["volume_file_path"], + genie_space_id=data["genie_space_id"], + warehouse_id=data["warehouse_id"], + ) + + +@dataclasses.dataclass +class EvidenceRow: + framework: str + authoring: str + runtime: str + tool_kind: str + status: str + command: str + expected: str + actual: str + duration_seconds: float + artifact_paths: list[str] + app_name: str | None = None + app_url: str | None = None + error: str | None = None + + +class Evidence: + """Appends the files under the output directory that the runner turns into ``evidence.json``.""" + + def __init__(self, output: pathlib.Path): + self.rows_path = output / "rows.jsonl" + self.registry_path = output / "registry.jsonl" + self.cleanup_path = output / "cleanup.jsonl" + self.grant_dir = output / "grant-checks" + self._lock = threading.Lock() + + def register_app(self, name: str) -> None: + self._append(self.registry_path, {"kind": "app", "app_name": name}) + + def register_project(self, app_name: str, **fields: Any) -> None: + self._append(self.registry_path, {"kind": "project", "app_name": app_name, **fields}) + + def mark_cleaned(self, app_name: str, results: Sequence[dict[str, Any]]) -> None: + for result in results: + self._append(self.cleanup_path, result) + self._append(self.registry_path, {"kind": "cleaned", "app_name": app_name}) + + def record_row(self, row: EvidenceRow) -> None: + self._append(self.rows_path, dataclasses.asdict(row)) + + def read_rows(self) -> list[dict[str, Any]]: + if not self.rows_path.exists(): + return [] + with self._lock: + lines = self.rows_path.read_text(encoding="utf-8").splitlines() + return [json.loads(line) for line in lines if line.strip()] + + def record_setup_failure( + self, + authoring: str, + runtime: str, + exc: Exception, + log_path: pathlib.Path | None, + kinds: Sequence[str], + app_name: str | None = None, + ) -> None: + """Fail the cells whose agent could not be set up, so the evidence names them.""" + existing = { + row["tool_kind"] + for row in self.read_rows() + if row["authoring"] == authoring and row["runtime"] == runtime + } + for kind in kinds: + if kind in existing: + continue + self.record_row( + EvidenceRow( + framework=FRAMEWORKS[0], + authoring=authoring, + runtime=runtime, + tool_kind=kind, + status="fail", + command="runtime setup", + expected=EXPECTED[kind], + actual="", + duration_seconds=0.0, + artifact_paths=[str(log_path)] if log_path else [], + app_name=app_name, + error=str(exc), + ) + ) + + def attach_artifact_to_failed_rows(self, app_name: str, artifact: pathlib.Path) -> bool: + rows = self.read_rows() + failed = [ + row for row in rows if row["status"] == "fail" and row.get("app_name") == app_name + ] + for row in failed: + row["artifact_paths"].append(str(artifact)) + if failed: + with self._lock: + temporary = self.rows_path.with_suffix(".jsonl.tmp") + temporary.write_text( + "".join(json.dumps(row) + "\n" for row in rows), encoding="utf-8" + ) + os.replace(temporary, self.rows_path) + return bool(failed) + + def record_grant_check(self, app_name: str, **fields: Any) -> None: + self.grant_dir.mkdir(parents=True, exist_ok=True) + path = self.grant_dir / f"{app_name}.json" + current = json.loads(path.read_text(encoding="utf-8")) if path.exists() else {} + current.update(fields) + current["app_name"] = app_name + path.write_text(json.dumps(current, indent=2), encoding="utf-8") + + def _append(self, path: pathlib.Path, record: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with self._lock: + with path.open("a", encoding="utf-8") as output: + output.write(json.dumps(record) + "\n") + + +@dataclasses.dataclass +class Outcome: + """What a cell observed; the row's ``actual`` and links are derived from it.""" + + serialized: str = "" + command: str = "" + log_path: pathlib.Path | None = None + app_name: str | None = None + app_url: str | None = None + + +@contextlib.contextmanager +def recorded_row( + evidence: Evidence, + authoring: str, + runtime: str, + kind: str, + *, + command: str, + expected: str, + marker: str | None = None, +) -> Iterator[Outcome]: + """Run one matrix cell and append its pass/fail evidence row whatever happens.""" + outcome = Outcome(command=command) + started = time.monotonic() + error: str | None = None + try: + yield outcome + except Exception as exc: + error = str(exc) or type(exc).__name__ + raise + finally: + evidence.record_row( + EvidenceRow( + framework=FRAMEWORKS[0], + authoring=authoring, + runtime=runtime, + tool_kind=kind, + status="pass" if error is None else "fail", + command=outcome.command, + expected=expected, + actual=evidence_excerpt(outcome.serialized, marker), + duration_seconds=round(time.monotonic() - started, 3), + artifact_paths=[str(outcome.log_path)] if outcome.log_path else [], + app_name=outcome.app_name, + app_url=outcome.app_url, + error=error, + ) + ) + + +def evidence_excerpt(serialized: str, required_marker: str | None) -> str: + limit = 6000 + if len(serialized) <= limit: + return serialized + head_length = limit // 2 + head = serialized[:head_length] + if required_marker is None or required_marker in head: + tail = serialized[-head_length:] + else: + marker_offset = serialized.find(required_marker) + start = max(head_length, marker_offset - head_length // 2) + tail = serialized[start : start + head_length] + return f"{head}\n... response truncated ...\n{tail}" diff --git a/integrations/agentbricks/tests/conftest.py b/integrations/agentbricks/tests/conftest.py new file mode 100644 index 000000000..fa0d08981 --- /dev/null +++ b/integrations/agentbricks/tests/conftest.py @@ -0,0 +1,53 @@ +"""Fixtures shared by every matrix under ``matrices/``. + +The CLI wrapper is function-scoped: each test authors, runs and deploys its own projects, and the +wrapper deletes every App and store it created when the test ends. +""" + +from __future__ import annotations + +from collections.abc import Iterator + +import pytest +from agentbricks_cli import AgentbricksCli +from common import Evidence, Inputs, Transcript +from workspace_client import Workspace + + +@pytest.fixture(scope="session") +def inputs() -> Inputs: + return Inputs.from_env() + + +@pytest.fixture(scope="session") +def transcript(inputs: Inputs) -> Transcript: + return Transcript(inputs.output / "commands.log") + + +@pytest.fixture(scope="session") +def evidence(inputs: Inputs) -> Evidence: + return Evidence(inputs.output) + + +@pytest.fixture(scope="session") +def workspace_client(inputs: Inputs, transcript: Transcript) -> Workspace: + workspace = Workspace( + inputs.databricks_profile, + transcript, + app_auth_profile=inputs.app_auth_profile, + warehouse_id=inputs.warehouse_id, + preprovisioned_app_catalog_access=inputs.preprovisioned_app_catalog_access, + ) + workspace.check_app_auth() + return workspace + + +@pytest.fixture +def agentbricks_cli( + inputs: Inputs, workspace_client: Workspace, transcript: Transcript, evidence: Evidence +) -> Iterator[AgentbricksCli]: + cli = AgentbricksCli(inputs, workspace_client, transcript, evidence) + try: + yield cli + finally: + cli.close() diff --git a/integrations/agentbricks/tests/matrices/memory_store/bundle/databricks.yml b/integrations/agentbricks/tests/matrices/memory_store/bundle/databricks.yml new file mode 100644 index 000000000..a322534b7 --- /dev/null +++ b/integrations/agentbricks/tests/matrices/memory_store/bundle/databricks.yml @@ -0,0 +1,26 @@ +# Per-run scratch schema for this matrix; the UC function and volume live inside it. +bundle: + name: agentbricks-memory-store-scratch + +variables: + catalog_name: {} + run_id: {} + +targets: + nightly: + default: true + variables: + catalog_name: main + dev: + variables: + catalog_name: main + +workspace: + # State is per run so concurrent or repeated runs never rename each other's schema. + root_path: ~/.bundle/${bundle.name}/${bundle.target}/${var.run_id} + +resources: + schemas: + scratch_schema: + catalog_name: ${var.catalog_name} + name: agentbricks_agent_tools_e2e_${var.run_id} diff --git a/integrations/agentbricks/tests/matrices/memory_store/suite.yaml b/integrations/agentbricks/tests/matrices/memory_store/suite.yaml new file mode 100644 index 000000000..973a30a98 --- /dev/null +++ b/integrations/agentbricks/tests/matrices/memory_store/suite.yaml @@ -0,0 +1,18 @@ +# STUB: copy-me manifest for the memory_store matrix. The memory store is created per project by +# the test fixture and deleted by the runner's teardown registry, so there is no workspace +# dependency to declare for it. Add an `existing` requirement here if a real case needs a pre- +# existing resource. +name: memory_store +tests: [matrices/memory_store/] +runs: + nightly: + timeout_minutes: 90 +requires: + - name: catalog + kind: existing + value: ${AGENTBRICKS_E2E_CATALOG:-main} + check: [catalogs, get] + - name: scratch_schema + kind: dab + bundle: bundle + catalog: catalog diff --git a/integrations/agentbricks/tests/matrices/memory_store/test_memory_store.py b/integrations/agentbricks/tests/matrices/memory_store/test_memory_store.py new file mode 100644 index 000000000..bd379002e --- /dev/null +++ b/integrations/agentbricks/tests/matrices/memory_store/test_memory_store.py @@ -0,0 +1,40 @@ +"""STUB: a single-feature agent with a memory store; the store the CLI declared was created. + +Real case still to write: the agent writes a memory in one invocation and recalls it in another. +""" + +from __future__ import annotations + +import json + +import pytest +from agentbricks_cli import AgentbricksCli +from common import EXPECTED, Evidence, recorded_row + + +def test_memory_store_created(agentbricks_cli: AgentbricksCli, evidence: Evidence) -> None: + project = agentbricks_cli.new_project("cli", keep=("memory_store",)) + command = f"agentbricks memory stores create --name {project.memory_store_config}" + with recorded_row( + evidence, + "cli", + "deploy", + "memory_store", + command=command, + expected=EXPECTED["memory_store"], + ) as row: + row.app_name = project.app_name + row.log_path = agentbricks_cli.deploy(project).log_path + resource = project.memory_resource or {} + row.serialized = json.dumps(resource, sort_keys=True) + + assert resource.get("display_name") == project.memory_store_config, resource + assert str(resource.get("name")).startswith("memory-stores/"), resource + assert project.memory_store_name == resource["name"] + + +@pytest.mark.xfail(reason="STUB: no memory write/recall case written yet", run=False) +def test_memory_recalled_across_invocations( + agentbricks_cli: AgentbricksCli, evidence: Evidence +) -> None: + raise NotImplementedError diff --git a/integrations/agentbricks/tests/matrices/session_store/bundle/databricks.yml b/integrations/agentbricks/tests/matrices/session_store/bundle/databricks.yml new file mode 100644 index 000000000..9bc6c7d19 --- /dev/null +++ b/integrations/agentbricks/tests/matrices/session_store/bundle/databricks.yml @@ -0,0 +1,26 @@ +# Per-run scratch schema for this matrix; the UC function and volume live inside it. +bundle: + name: agentbricks-session-store-scratch + +variables: + catalog_name: {} + run_id: {} + +targets: + nightly: + default: true + variables: + catalog_name: main + dev: + variables: + catalog_name: main + +workspace: + # State is per run so concurrent or repeated runs never rename each other's schema. + root_path: ~/.bundle/${bundle.name}/${bundle.target}/${var.run_id} + +resources: + schemas: + scratch_schema: + catalog_name: ${var.catalog_name} + name: agentbricks_agent_tools_e2e_${var.run_id} diff --git a/integrations/agentbricks/tests/matrices/session_store/suite.yaml b/integrations/agentbricks/tests/matrices/session_store/suite.yaml new file mode 100644 index 000000000..c5797d0ad --- /dev/null +++ b/integrations/agentbricks/tests/matrices/session_store/suite.yaml @@ -0,0 +1,18 @@ +# STUB: copy-me manifest for the session_store matrix. The session store is created per project by +# the test fixture and deleted by the runner's teardown registry, so there is no workspace +# dependency to declare for it. Add an `existing` requirement here if a real case needs a pre- +# existing resource. +name: session_store +tests: [matrices/session_store/] +runs: + nightly: + timeout_minutes: 90 +requires: + - name: catalog + kind: existing + value: ${AGENTBRICKS_E2E_CATALOG:-main} + check: [catalogs, get] + - name: scratch_schema + kind: dab + bundle: bundle + catalog: catalog diff --git a/integrations/agentbricks/tests/matrices/session_store/test_session_store.py b/integrations/agentbricks/tests/matrices/session_store/test_session_store.py new file mode 100644 index 000000000..fc35e1b4d --- /dev/null +++ b/integrations/agentbricks/tests/matrices/session_store/test_session_store.py @@ -0,0 +1,39 @@ +"""STUB: a single-feature agent with a session store; the store the CLI declared was created. + +Real case still to write: conversation state persists across invocations of one session_id and is +isolated between sessions. +""" + +from __future__ import annotations + +import json + +import pytest +from agentbricks_cli import AgentbricksCli +from common import EXPECTED, Evidence, recorded_row + + +def test_session_store_created(agentbricks_cli: AgentbricksCli, evidence: Evidence) -> None: + project = agentbricks_cli.new_project("cli", keep=("session_store",)) + command = f"agentbricks sessions stores create --name {project.session_store_config}" + with recorded_row( + evidence, + "cli", + "deploy", + "session_store", + command=command, + expected=EXPECTED["session_store"], + ) as row: + row.app_name = project.app_name + row.log_path = agentbricks_cli.deploy(project).log_path + resource = project.session_resource or {} + row.serialized = json.dumps(resource, sort_keys=True) + + assert resource.get("session_store_name") == project.session_store_config, resource + + +@pytest.mark.xfail(reason="STUB: no session-persistence case written yet", run=False) +def test_session_persists_across_invocations( + agentbricks_cli: AgentbricksCli, evidence: Evidence +) -> None: + raise NotImplementedError diff --git a/integrations/agentbricks/tests/matrices/tools/bundle/databricks.yml b/integrations/agentbricks/tests/matrices/tools/bundle/databricks.yml new file mode 100644 index 000000000..34e7bb212 --- /dev/null +++ b/integrations/agentbricks/tests/matrices/tools/bundle/databricks.yml @@ -0,0 +1,26 @@ +# Per-run scratch schema for this matrix; the UC function and volume live inside it. +bundle: + name: agentbricks-tools-scratch + +variables: + catalog_name: {} + run_id: {} + +targets: + nightly: + default: true + variables: + catalog_name: main + dev: + variables: + catalog_name: main + +workspace: + # State is per run so concurrent or repeated runs never rename each other's schema. + root_path: ~/.bundle/${bundle.name}/${bundle.target}/${var.run_id} + +resources: + schemas: + scratch_schema: + catalog_name: ${var.catalog_name} + name: agentbricks_agent_tools_e2e_${var.run_id} diff --git a/integrations/agentbricks/tests/matrices/tools/conftest.py b/integrations/agentbricks/tests/matrices/tools/conftest.py new file mode 100644 index 000000000..a4de93e51 --- /dev/null +++ b/integrations/agentbricks/tests/matrices/tools/conftest.py @@ -0,0 +1,11 @@ +"""Tools matrix: every test runs once per authoring path.""" + +from __future__ import annotations + +import pytest +from common import AUTHORING_PATHS + + +@pytest.fixture(params=AUTHORING_PATHS) +def authoring(request: pytest.FixtureRequest) -> str: + return request.param diff --git a/integrations/agentbricks/tests/matrices/tools/suite.yaml b/integrations/agentbricks/tests/matrices/tools/suite.yaml new file mode 100644 index 000000000..7449bdc83 --- /dev/null +++ b/integrations/agentbricks/tests/matrices/tools/suite.yaml @@ -0,0 +1,25 @@ +# Dependency manifest for the LangGraph x CLI/direct x dev/deploy tool matrix. run_suites.py +# provisions these in order; `existing` is only validated, never created or deleted. +name: tools +tests: [matrices/tools/] +runs: + nightly: + timeout_minutes: 240 +requires: + - name: catalog + kind: existing + value: ${AGENTBRICKS_E2E_CATALOG:-main} + check: [catalogs, get] + - name: scratch_schema + kind: dab + bundle: bundle + catalog: catalog + # Nested + marker functions and a volume, created imperatively in the schema and swept by its destroy. + - name: uc_function + kind: uc_function + schema: scratch_schema + - name: genie_space + kind: existing + value: ${AGENTBRICKS_E2E_GENIE_SPACE_ID} + check: [genie, get-space] + pattern: "^[0-9a-f]{32}$" diff --git a/integrations/agentbricks/tests/matrices/tools/test_deploy.py b/integrations/agentbricks/tests/matrices/tools/test_deploy.py new file mode 100644 index 000000000..ca9ca6a77 --- /dev/null +++ b/integrations/agentbricks/tests/matrices/tools/test_deploy.py @@ -0,0 +1,64 @@ +"""Each tool works on a deployed App, and its UC grant is applied, per authoring path. + +CLI authoring is incremental: add one tool, redeploy the same App, verify it, then add the next. +After every redeploy the grants of all tools added so far must still hold, which is the idempotency +check. Direct authoring writes every tool into agent.toml and deploys once. A subtest per tool +reports each tool independently and the sequence continues past a failure. +""" + +from __future__ import annotations + +from agentbricks_cli import AgentbricksCli +from common import Evidence, Inputs +from tools import bind, invoke_and_check, setup_failures, tool_row, tools_for +from workspace_client import Workspace + + +def test_deploy( + agentbricks_cli: AgentbricksCli, + workspace_client: Workspace, + inputs: Inputs, + evidence: Evidence, + authoring: str, + subtests, +) -> None: + tools = tools_for(inputs) + project = agentbricks_cli.new_project(authoring) + app = project.app_name + + def assert_granted(tool) -> None: + if tool.grant is not None: + assert workspace_client.granted(app, tool.grant), ( + f"{tool.name}: {tool.grant} not granted to {app}: " + f"{sorted(workspace_client.granted_tuples(app))}" + ) + + def grant_transitive_if_needed(tool) -> None: + # The nested function is deliberately outside Agent Bricks' grants; the tool only works + # once this manual grant is in place. + if tool.needs_transitive_grant: + workspace_client.grant_transitive(app, inputs.transitive_uc_function) + + if authoring == "cli": + for index, tool in enumerate(tools): + with ( + subtests.test(tool=tool.name), + tool_row(evidence, authoring, "deploy", tool) as row, + ): + row.app_name = app + bind(agentbricks_cli, project, tool) + agent = agentbricks_cli.deploy(project, app) + grant_transitive_if_needed(tool) + invoke_and_check(row, agent, tool) + for added in tools[: index + 1]: + assert_granted(added) + return + + with setup_failures(evidence, authoring, "deploy", tools, app): + agentbricks_cli.write_manifest(project, tools) + agent = agentbricks_cli.deploy(project, app) + for tool in tools: + with subtests.test(tool=tool.name), tool_row(evidence, authoring, "deploy", tool) as row: + grant_transitive_if_needed(tool) + invoke_and_check(row, agent, tool) + assert_granted(tool) diff --git a/integrations/agentbricks/tests/matrices/tools/test_dev.py b/integrations/agentbricks/tests/matrices/tools/test_dev.py new file mode 100644 index 000000000..23b76049e --- /dev/null +++ b/integrations/agentbricks/tests/matrices/tools/test_dev.py @@ -0,0 +1,46 @@ +"""Each tool works in `agentbricks dev`, per authoring path. + +CLI authoring is incremental: add one tool, restart dev, verify it, then add the next on the same +project. Direct authoring writes every tool into agent.toml and runs dev once. A subtest per tool +reports each tool independently and the sequence continues past a failure. +""" + +from __future__ import annotations + +import contextlib + +from agentbricks_cli import AgentbricksCli +from common import Evidence, Inputs +from tools import ( + bind, + invoke_and_check, + reject_unavailable_mcp, + setup_failures, + tool_row, + tools_for, +) + + +def test_dev( + agentbricks_cli: AgentbricksCli, inputs: Inputs, evidence: Evidence, authoring: str, subtests +) -> None: + tools = tools_for(inputs) + project = agentbricks_cli.new_project(authoring) + + if authoring == "cli": + with subtests.test(step="rejects-unavailable-mcp"): + reject_unavailable_mcp(agentbricks_cli, project) + for tool in tools: + with subtests.test(tool=tool.name), tool_row(evidence, authoring, "dev", tool) as row: + bind(agentbricks_cli, project, tool) + with agentbricks_cli.dev(project) as agent: + invoke_and_check(row, agent, tool) + return + + with contextlib.ExitStack() as stack: + with setup_failures(evidence, authoring, "dev", tools): + agentbricks_cli.write_manifest(project, tools) + agent = stack.enter_context(agentbricks_cli.dev(project)) + for tool in tools: + with subtests.test(tool=tool.name), tool_row(evidence, authoring, "dev", tool) as row: + invoke_and_check(row, agent, tool) diff --git a/integrations/agentbricks/tests/matrices/tools/test_grants.py b/integrations/agentbricks/tests/matrices/tools/test_grants.py new file mode 100644 index 000000000..5fce3cd95 --- /dev/null +++ b/integrations/agentbricks/tests/matrices/tools/test_grants.py @@ -0,0 +1,54 @@ +"""Deploy grants that belong to no single tool (deploy-only): nothing transitive, nothing lost.""" + +from __future__ import annotations + +from agentbricks_cli import AgentbricksCli +from common import TOOL_RESOURCE_PREFIX, Evidence, Inputs +from tools import bind, tools_for +from workspace_client import Workspace + + +def test_grants( + agentbricks_cli: AgentbricksCli, + workspace_client: Workspace, + inputs: Inputs, + evidence: Evidence, + authoring: str, + subtests, +) -> None: + uc_function = next(tool for tool in tools_for(inputs) if tool.name == "uc_function") + # Tracing stays bound so the App carries a resource Agent Bricks does not own. + project = agentbricks_cli.new_project(authoring, keep=("tracing",)) + app = project.app_name + if authoring == "cli": + bind(agentbricks_cli, project, uc_function) + else: + agentbricks_cli.write_manifest(project, [uc_function]) + agentbricks_cli.deploy(project, app) + state: dict = {} + + with subtests.test(check="transitive-function-not-auto-granted"): + # Raises if the hidden nested function is already granted or listed as an App resource. + state.update(workspace_client.transitive_state(app, inputs.transitive_uc_function)) + + with subtests.test(check="manual-grant-is-direct-and-effective"): + granted = workspace_client.grant_transitive(app, inputs.transitive_uc_function) + if state: + evidence.record_grant_check( + app, + authoring=authoring, + service_principal_client_id=state["service_principal_client_id"], + initial={ + "transitive_direct_privileges": state["direct_privileges"], + "transitive_effective_privileges": state["effective_privileges"], + }, + post_manual_transitive_grant=granted, + manual_transitive_grant_applied=True, + ) + assert "EXECUTE" in granted["direct_privileges"] + assert "EXECUTE" in granted["effective_privileges"] + + with subtests.test(check="non-tool-resource-preserved"): + unrelated = workspace_client.non_tool_resources(app) + assert unrelated, "No non-tool App resource; the tracing experiment should be attached" + assert not any(r["name"].startswith(TOOL_RESOURCE_PREFIX) for r in unrelated) diff --git a/integrations/agentbricks/tests/matrices/tools/test_selection.py b/integrations/agentbricks/tests/matrices/tools/test_selection.py new file mode 100644 index 000000000..16dbd9d3d --- /dev/null +++ b/integrations/agentbricks/tests/matrices/tools/test_selection.py @@ -0,0 +1,59 @@ +"""MULTI-TOOL SELECTION check (deploy-only): with every tool bound, the agent picks the one asked for. + +The incremental tests cannot catch a regression where binding several tools makes the model, or +the tool surface, route to the wrong one, so this is the one test that binds them all at once. +""" + +from __future__ import annotations + +from agentbricks_cli import AgentbricksCli +from common import EXPECTED, Evidence, Inputs, recorded_row +from tools import PYTHON_MARKER, UC_MARKER, bind, tools_for +from workspace_client import Workspace + +# Names the Unity Catalog function without its tool id, and the Python helper only to exclude it. +PROMPT = ( + "Use the Unity Catalog function tool to compute the Agent Bricks marker for the value " + "'matrix', not the local Python helper tool. Return the called tool's exact result." +) + + +def test_selects_uc_function_not_python_marker( + agentbricks_cli: AgentbricksCli, + workspace_client: Workspace, + inputs: Inputs, + evidence: Evidence, + authoring: str, +) -> None: + tools = tools_for(inputs) + project = agentbricks_cli.new_project(authoring) + app = project.app_name + with recorded_row( + evidence, + authoring, + "deploy", + "selection", + command="agentbricks deploy", + expected=EXPECTED["selection"], + marker=UC_MARKER, + ) as row: + row.app_name = app + if authoring == "cli": + for tool in tools: + bind(agentbricks_cli, project, tool) + else: + agentbricks_cli.write_manifest(project, tools) + agent = agentbricks_cli.deploy(project, app) + # The nested function is deliberately outside Agent Bricks' grants. + workspace_client.grant_transitive(app, inputs.transitive_uc_function) + row.command = agent.curl(PROMPT) + row.log_path = agent.log_path + row.app_url = agent.base_url + row.serialized = agent.invoke(PROMPT, "deploy-selection") + + assert UC_MARKER in row.serialized, ( + f"The UC function tool was not called: {row.serialized[:2000]}" + ) + assert PYTHON_MARKER not in row.serialized, ( + f"The Python marker tool was called too: {row.serialized[:2000]}" + ) diff --git a/integrations/agentbricks/tests/matrices/tools/tools.py b/integrations/agentbricks/tests/matrices/tools/tools.py new file mode 100644 index 000000000..8bcd49e70 --- /dev/null +++ b/integrations/agentbricks/tests/matrices/tools/tools.py @@ -0,0 +1,241 @@ +"""The tools under test: how each is added, prompted, verified, and what a deploy grants for it.""" + +from __future__ import annotations + +import contextlib +import dataclasses +import json +from collections.abc import Iterator, Mapping, Sequence + +from agentbricks_cli import Agent, AgentbricksCli, Project +from common import EXPECTED, Evidence, Inputs, MatrixError, Outcome, recorded_row +from workspace_client import GrantTuple + +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" + + +@dataclasses.dataclass(frozen=True) +class Tool: + name: str + prompt: str + # The exact text the response must contain; None for tools checked by their tool-call evidence. + marker: str | None + # `agentbricks tools add` arguments; None when the tool is a user-owned file instead. + add_args: tuple[str, ...] | None + # The [[tools]] entry for direct authoring, and any user-owned files the tool needs. + toml: str + files: Mapping[str, str] + # The id agent.toml declares it under, and whether it must be App-authed. + tool_id: str | None + app_auth: bool + # The App resource a deploy grants for it; None when it needs no UC grant. + grant: GrantTuple | None = None + # The UC function wraps a nested one that Agent Bricks deliberately does not grant. + needs_transitive_grant: bool = False + + def check(self, serialized: str) -> None: + lowered = serialized.lower() + if self.marker is not None: + if self.marker not in serialized: + raise AssertionError(f"Missing exact marker {self.marker!r}: {serialized[:2000]}") + elif self.name == "mcp": + tool_evidence = any(v in lowered for v in ("web_search", "web search", "search")) + if not tool_evidence or "https" not in lowered or len(serialized) < 80: + raise AssertionError( + f"Missing web-search execution/result evidence: {serialized[:2000]}" + ) + elif self.name == "genie": + if ( + not any(v in lowered for v in ("genie_ask", "genie", "conversation_id")) + or len(serialized) < 80 + ): + raise AssertionError( + f"Missing Genie execution/result evidence: {serialized[:2000]}" + ) + else: + raise MatrixError(f"No semantic assertion is defined for tool {self.name!r}.") + + +def tools_for(inputs: Inputs) -> tuple[Tool, ...]: + """The TOOLS table, resolved against this run's scratch volume, function and Genie space.""" + uc_tool_name = inputs.uc_function.replace(".", "__") + return ( + Tool( + name="sandbox", + prompt=( + "You must call the sandbox tool. In the sandbox, use Python to read the " + f"entire text file at {inputs.volume_file_path}. Return only its exact contents; " + "do not fabricate them." + ), + marker=inputs.volume_marker, + add_args=("sandbox", "--scope", f"volume:{inputs.uc_volume}", "--auth", "app"), + toml=( + '[[tools]]\nid = "sandbox"\n' + 'source = { kind = "sandbox", service = "system.ai.sandbox" }\n' + "policy = { downscope = [\n" + f' {{ resource = "volume:{inputs.uc_volume}", permission = "read_only" }},\n' + "] }\n" + ), + files={}, + tool_id="sandbox", + app_auth=True, + grant=("uc_securable", inputs.uc_volume, "VOLUME", "READ_VOLUME"), + ), + 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." + ), + marker=None, + add_args=("mcp", "system.ai.web_search", "--auth", "app"), + toml=( + '[[tools]]\nid = "web_search"\n' + 'source = { kind = "mcp", service = "system.ai.web_search" }\n' + ), + files={}, + tool_id="web_search", + app_auth=True, + ), + Tool( + name="python", + prompt=( + "You must call the matrix_marker Python tool with value 'matrix'. " + "Return its exact result." + ), + marker=PYTHON_MARKER, + add_args=None, + toml="", + files={PYTHON_MARKER_FILE: PYTHON_MARKER_SOURCE}, + tool_id=None, + app_auth=False, + ), + Tool( + name="uc_function", + prompt=( + f"You must call the tool named {uc_tool_name} with value 'matrix'. " + "Return the called tool's exact result." + ), + marker=UC_MARKER, + add_args=("uc-function", inputs.uc_function, "--name", "agentbricks_uc_marker"), + toml=( + '[[tools]]\nid = "agentbricks_uc_marker"\n' + f'source = {{ kind = "uc_function", function = "{inputs.uc_function}" }}\n' + ), + files={}, + tool_id="agentbricks_uc_marker", + app_auth=False, + grant=("uc_securable", inputs.uc_function, "FUNCTION", "EXECUTE"), + needs_transitive_grant=True, + ), + 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." + ), + marker=None, + add_args=("genie-agent", inputs.genie_space_id, "--name", "genie", "--auth", "app"), + toml=( + '[[tools]]\nid = "genie"\nauth = "app"\n' + f'source = {{ kind = "genie_agent", space_id = "{inputs.genie_space_id}" }}\n' + ), + files={}, + tool_id="genie", + app_auth=True, + grant=("genie_space", inputs.genie_space_id, "GENIE_SPACE", "CAN_RUN"), + ), + ) + + +def bind(cli: AgentbricksCli, project: Project, tool: Tool) -> None: + """CLI authoring: add one tool and confirm agent.toml declares it as intended.""" + if tool.add_args is not None: + cli.tools_add(project, *tool.add_args) + for relative_path, content in tool.files.items(): + cli.write_file(project, relative_path, content) + if tool.tool_id is None: + return + declared = {t["id"]: t for t in cli.manifest(project).get("tools", [])} + if tool.tool_id not in declared: + raise MatrixError(f"tools add did not declare {tool.tool_id!r}: {sorted(declared)}") + if tool.app_auth and declared[tool.tool_id].get("auth") != "app": + raise MatrixError(f"CLI-authored {tool.tool_id} binding is not App-auth.") + + +def reject_unavailable_mcp(cli: AgentbricksCli, project: Project) -> None: + """`tools add` must refuse an MCP service that does not exist and leave agent.toml untouched.""" + manifest = project.path / "agent.toml" + before = manifest.read_bytes() + rejected = cli.tools_add( + project, + "mcp", + "system.ai.missing_service", + "--name", + "broken_mcp", + check=False, + json_output=True, + ) + if rejected.returncode == 0 or manifest.read_bytes() != before: + raise MatrixError( + "agentbricks tools add accepted an unavailable MCP service or changed agent.toml" + ) + if json.loads(rejected.stderr).get("error", {}).get("code") not in { + "NOT_FOUND", + "RESOURCE_DOES_NOT_EXIST", + }: + raise MatrixError(f"Unexpected MCP validation error: {rejected.stderr}") + cli.tools_remove(project, "mcp", "system.ai.missing_service") + if any(tool["id"] == "broken_mcp" for tool in cli.manifest(project).get("tools", [])): + raise MatrixError("agentbricks tools remove left the broken MCP binding in agent.toml") + + +def tool_row(evidence: Evidence, authoring: str, runtime: str, tool: Tool): + """One evidence row for one tool in one authoring x runtime cell.""" + return recorded_row( + evidence, + authoring, + runtime, + tool.name, + command=f"agentbricks {'dev' if runtime == 'dev' else 'deploy'}", + expected=tool.marker or EXPECTED[tool.name], + marker=tool.marker, + ) + + +def invoke_and_check(outcome: Outcome, agent: Agent, tool: Tool) -> None: + outcome.command = agent.curl(tool.prompt) + outcome.log_path = agent.log_path + outcome.app_name = agent.app_name + outcome.app_url = agent.base_url if agent.app_name else None + outcome.serialized = agent.invoke(tool.prompt, f"{agent.runtime}-{tool.name}") + tool.check(outcome.serialized) + + +@contextlib.contextmanager +def setup_failures( + evidence: Evidence, + authoring: str, + runtime: str, + tools: Sequence[Tool], + app_name: str | None = None, +) -> Iterator[None]: + """Fail every tool's cell when the shared agent they all need cannot be set up.""" + try: + yield + except Exception as exc: + evidence.record_setup_failure( + authoring, runtime, exc, None, [tool.name for tool in tools], app_name + ) + raise diff --git a/integrations/agentbricks/tests/matrices/tracing/bundle/databricks.yml b/integrations/agentbricks/tests/matrices/tracing/bundle/databricks.yml new file mode 100644 index 000000000..83558e3f0 --- /dev/null +++ b/integrations/agentbricks/tests/matrices/tracing/bundle/databricks.yml @@ -0,0 +1,26 @@ +# Per-run scratch schema for this matrix; the UC function and volume live inside it. +bundle: + name: agentbricks-tracing-scratch + +variables: + catalog_name: {} + run_id: {} + +targets: + nightly: + default: true + variables: + catalog_name: main + dev: + variables: + catalog_name: main + +workspace: + # State is per run so concurrent or repeated runs never rename each other's schema. + root_path: ~/.bundle/${bundle.name}/${bundle.target}/${var.run_id} + +resources: + schemas: + scratch_schema: + catalog_name: ${var.catalog_name} + name: agentbricks_agent_tools_e2e_${var.run_id} diff --git a/integrations/agentbricks/tests/matrices/tracing/suite.yaml b/integrations/agentbricks/tests/matrices/tracing/suite.yaml new file mode 100644 index 000000000..d86ed9f2e --- /dev/null +++ b/integrations/agentbricks/tests/matrices/tracing/suite.yaml @@ -0,0 +1,17 @@ +# STUB: copy-me manifest for the tracing matrix. The tracing experiment is get-or-created by +# deploy under /Shared/agentbricks_traces, so there is no workspace dependency to declare for it. +# Add an `existing` requirement here if a real case needs a pre-existing resource. +name: tracing +tests: [matrices/tracing/] +runs: + nightly: + timeout_minutes: 90 +requires: + - name: catalog + kind: existing + value: ${AGENTBRICKS_E2E_CATALOG:-main} + check: [catalogs, get] + - name: scratch_schema + kind: dab + bundle: bundle + catalog: catalog diff --git a/integrations/agentbricks/tests/matrices/tracing/test_tracing.py b/integrations/agentbricks/tests/matrices/tracing/test_tracing.py new file mode 100644 index 000000000..2754087b2 --- /dev/null +++ b/integrations/agentbricks/tests/matrices/tracing/test_tracing.py @@ -0,0 +1,47 @@ +"""STUB: a single-feature agent with tracing; the experiment is bound and attached to the App. + +Real case still to write: an invocation exports a trace to the experiment. +""" + +from __future__ import annotations + +import json + +import pytest +from agentbricks_cli import AgentbricksCli +from common import EXPECTED, TOOL_RESOURCE_PREFIX, Evidence, recorded_row +from workspace_client import Workspace + + +def test_tracing_experiment_bound( + agentbricks_cli: AgentbricksCli, workspace_client: Workspace, evidence: Evidence +) -> None: + project = agentbricks_cli.new_project("cli", keep=("tracing",)) + with recorded_row( + evidence, + "cli", + "deploy", + "tracing", + command="agentbricks deploy (agent.toml [tracing] experiment_name)", + expected=EXPECTED["tracing"], + ) as row: + row.app_name = project.app_name + row.log_path = agentbricks_cli.deploy(project).log_path + assert project.experiment_name, "agent.toml binds no tracing experiment" + row.serialized = json.dumps({"experiment_name": project.experiment_name}) + + # Deploy attaches the experiment as an App resource that Agent Bricks does not own. + experiments = [ + r for r in workspace_client.non_tool_resources(project.app_name) if "experiment" in r + ] + assert experiments, "No experiment App resource among the non-tool resources" + assert not any(r["name"].startswith(TOOL_RESOURCE_PREFIX) for r in experiments) + expected_id = workspace_client.experiment_id(project.experiment_name) + assert any(str(r["experiment"].get("experiment_id")) == expected_id for r in experiments), ( + f"{project.experiment_name} ({expected_id}) not in {experiments}" + ) + + +@pytest.mark.xfail(reason="STUB: no trace-export case written yet", run=False) +def test_invocation_exports_trace(agentbricks_cli: AgentbricksCli, evidence: Evidence) -> None: + raise NotImplementedError diff --git a/integrations/agentbricks/tests/run_suites.py b/integrations/agentbricks/tests/run_suites.py new file mode 100644 index 000000000..944b754a4 --- /dev/null +++ b/integrations/agentbricks/tests/run_suites.py @@ -0,0 +1,1061 @@ +#!/usr/bin/env python3 +"""Provision a suite's workspace dependencies, run its pytest tests, and tear everything down. + +One entry point for local and CI runs. Auth comes from the environment (ambient Databricks +credentials, or DATABRICKS_CONFIG_PROFILE locally): + + uv run python tests/run_suites.py --profile nightly + +Each ``matrices//suite.yaml`` manifest declares the dependencies (``requires``) its pytest directory +needs. Three kinds exist: ``existing`` (validated, never created or deleted), ``dab`` (a bundle +deployed before and destroyed after the tests), and ``uc_function`` (imperative SQL inside a +``dab`` schema, swept by the schema's destroy). The runner hands the resolved values to pytest +through the JSON file named by ``AGENTBRICKS_TEST_INPUTS`` and owns the evidence file. +""" + +from __future__ import annotations + +import argparse +import dataclasses +import functools +import hashlib +import json +import os +import pathlib +import re +import shlex +import signal +import subprocess +import sys +import tempfile +import time +import uuid +import zipfile +from collections.abc import Sequence +from typing import Any + +import yaml +from agentbricks_cli import delete_store +from common import ( + AUTHORING_PATHS, + FRAMEWORKS, + INPUTS_ENV, + MatrixError, + Transcript, + last_nonempty_line, + now, +) +from databricks.sdk import WorkspaceClient +from workspace_client import Workspace, cleanup_app + +TESTS_DIR = pathlib.Path(__file__).resolve().parent +MATRICES_DIR = TESTS_DIR / "matrices" +PACKAGE_DIR = TESTS_DIR.parent +REPOSITORY_ROOT = TESTS_DIR.parents[2] +PROFILES = ("ci", "nightly") +REQUIREMENT_KINDS = ("existing", "dab", "uc_function") +# Time held back from the suite timeout so teardown can finish after pytest is stopped. +TEARDOWN_RESERVE_SECONDS = 180 + +_WHEEL_SOURCE_FILES = ( + "databricks_agentkit/_api_client.py", + "databricks_agentbricks/tool_access.py", + "databricks_agentbricks/app_resources.py", + "databricks_agentbricks/cli/deploy.py", +) +_ENV_REFERENCE = re.compile(r"\$\{([A-Za-z_][A-Za-z0-9_]*)(?::-([^}]*))?\}") + + +@dataclasses.dataclass +class Suite: + name: str + path: pathlib.Path + tests: list[str] + runs: dict[str, dict[str, Any]] + requires: list[dict[str, Any]] + + +@dataclasses.dataclass +class Options: + run_profile: str + output: pathlib.Path + wheel: pathlib.Path | None + commit_sha: str | None + source_root: pathlib.Path + genie_space_id: str | None + bridge_sha: str | None + warehouse_id: str | None + databricks_profile: str | None + app_auth_profile: str | None + preprovisioned_app_catalog_access: bool + keep_resources: bool + + +class Shell: + """Subprocess and Databricks CLI wrappers that log every command.""" + + def __init__(self, profile: str | None, output: pathlib.Path): + self.profile = profile + self.output = output + self.transcript = Transcript(output / "commands.log") + self.warehouse_id: str | None = None + + def profile_args(self) -> list[str]: + return ["--profile", self.profile] if self.profile else [] + + def auth_label(self) -> str: + return f"profile {self.profile!r}" if self.profile else "ambient environment credentials" + + def run( + self, + argv: Sequence[str], + *, + cwd: pathlib.Path | None = None, + env: dict[str, str] | None = None, + timeout: float = 300, + log: bool = True, + check: bool = True, + ) -> subprocess.CompletedProcess[str]: + if log: + self.transcript.command(argv, cwd) + result = subprocess.run( + list(argv), + cwd=cwd, + env=env, + text=True, + capture_output=True, + timeout=timeout, + check=False, + ) + if log and result.stdout.strip(): + self.transcript.write(result.stdout) + if log and result.stderr.strip(): + self.transcript.write(result.stderr) + if check and result.returncode != 0: + raise MatrixError( + f"Command failed ({result.returncode}): {shlex.join(list(argv))}\n" + f"{result.stderr or result.stdout}" + ) + return result + + def run_long( + self, + label: str, + argv: Sequence[str], + *, + cwd: pathlib.Path | None = None, + env: dict[str, str] | None = None, + timeout: float = 1800, + ) -> str: + self.transcript.command(argv, cwd) + log_path = self.output / "logs" / f"{label}.log" + log_path.parent.mkdir(parents=True, exist_ok=True) + started = time.monotonic() + with log_path.open("w", encoding="utf-8") as log_file: + process = subprocess.Popen( + list(argv), + cwd=cwd, + env=env, + text=True, + stdout=log_file, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + next_tick = 60.0 + while process.poll() is None: + elapsed = time.monotonic() - started + if elapsed >= timeout: + os.killpg(process.pid, signal.SIGTERM) + raise MatrixError(f"{label} timed out after {timeout:.0f}s; log: {log_path}") + if elapsed >= next_tick: + self.transcript.write( + f"tick {now():%H:%M} | {label} | running | {last_nonempty_line(log_path)}" + ) + next_tick += 60.0 + time.sleep(2) + output = log_path.read_text(encoding="utf-8", errors="replace") + self.transcript.write(output) + if process.returncode != 0: + raise MatrixError(f"{label} failed ({process.returncode}); log: {log_path}") + self.transcript.write(f"tick {now():%H:%M} | {label} | success") + return output + + def databricks(self, args: Sequence[str], *, timeout: float = 300) -> Any: + result = self.run( + ["databricks", *args, *self.profile_args(), "--output", "json"], + timeout=timeout, + ) + try: + return json.loads(result.stdout or "{}") + except json.JSONDecodeError as exc: + raise MatrixError(f"Databricks CLI returned invalid JSON: {result.stdout}") from exc + + +class Toolchain(Shell): + """The wheel under test, its venv, and the workspace identity shared by every suite.""" + + def __init__(self, options: Options): + super().__init__(options.databricks_profile, options.output) + self.options = options + self.app_auth_profile = options.app_auth_profile or options.databricks_profile + self.wheel = options.wheel.resolve() if options.wheel else None + self.commit_sha = options.commit_sha.lower() if options.commit_sha else None + self.source_provenance: dict[str, Any] = {} + self.versions: dict[str, str] = {} + self.host: str | None = None + self.runner_venv = options.output / "runner-venv" + self.agentbricks = self.runner_venv / "bin" / "agentbricks" + + def _app_auth_label(self) -> str: + return ( + f"App auth profile {self.app_auth_profile!r}" + if self.app_auth_profile + else "ambient environment credentials" + ) + + def bootstrap(self) -> None: + self.output.mkdir(parents=True, exist_ok=True) + source_root = self.options.source_root.resolve() + if self.commit_sha is None: + self.commit_sha = ( + self.run(["git", "rev-parse", "HEAD"], cwd=source_root).stdout.strip().lower() + ) + if re.fullmatch(r"[0-9a-f]{40}", self.commit_sha) is None: + raise MatrixError("--commit-sha must be a full 40-character hexadecimal Git SHA") + if self.options.bridge_sha is not None and ( + re.fullmatch(r"[0-9a-f]{40}", self.options.bridge_sha) is None + ): + raise MatrixError("--bridge-sha must be a 40-character lowercase Git commit SHA.") + if self.wheel is None: + dist = self.output / "dist" + self.run( + ["uv", "build", "--wheel", "--out-dir", str(dist)], cwd=PACKAGE_DIR, timeout=600 + ) + wheels = sorted(dist.glob("*.whl")) + if not wheels: + raise MatrixError("uv build produced no wheel") + self.wheel = wheels[-1].resolve() + self.source_provenance = _source_provenance(source_root, self.commit_sha, self.wheel) + self.run(["uv", "venv", str(self.runner_venv)], timeout=300) + self.run( + [ + "uv", + "pip", + "install", + "--python", + str(self.runner_venv / "bin" / "python"), + str(self.wheel), + ], + timeout=600, + ) + self.run([str(self.agentbricks), "tools", "--help"]) + version_commands = { + "agentbricks": [str(self.agentbricks), "--version"], + "databricks": ["databricks", "version"], + "uv": ["uv", "--version"], + "python": [str(self.runner_venv / "bin" / "python"), "--version"], + } + for name, command in version_commands.items(): + result = self.run(command, log=False, check=False) + version = (result.stdout or result.stderr).strip() + if result.returncode != 0 or not version: + raise MatrixError(f"Could not record {name} version: {version or 'no output'}") + self.versions[name] = version + workspace_client = WorkspaceClient(profile=self.profile) + app_auth_client = WorkspaceClient(profile=self.app_auth_profile) + if not workspace_client.config.host: + raise MatrixError(f"Could not resolve a host from {self.auth_label()}.") + if not app_auth_client.config.host: + raise MatrixError(f"Could not resolve a host from {self._app_auth_label()}.") + self.host = workspace_client.config.host.rstrip("/") + app_auth_host = app_auth_client.config.host.rstrip("/") + if app_auth_host != self.host: + raise MatrixError(f"{self._app_auth_label()} targets {app_auth_host}, not {self.host}.") + if app_auth_client.config.auth_type == "pat": + raise MatrixError( + f"{self._app_auth_label()} uses a PAT. " + "Databricks Apps /api routes require OAuth; run `databricks auth login` " + "for a profile on the same workspace." + ) + if not app_auth_client.config.authenticate().get("Authorization"): + raise MatrixError(f"Could not resolve credentials from {self._app_auth_label()}.") + + +class SuiteRunner(Shell): + def __init__(self, suite: Suite, toolchain: Toolchain, options: Options): + super().__init__(toolchain.profile, options.output / suite.name) + self.suite = suite + self.toolchain = toolchain + self.options = options + self.run_id = uuid.uuid4().hex[:8] + self.run_suffix = self.run_id[:6] + self.cleanup_required = not options.keep_resources + self.cleanup_complete = False + self.cleanup_results: list[dict[str, Any]] = [] + self.started_at = now().isoformat() + self.ended_at: str | None = None + self.resolved: dict[str, str] = {} + self.catalog: str | None = None + self.schema: str | None = None + self.uc_function: str | None = None + self.transitive_uc_function: str | None = None + self.uc_volume: str | None = None + self.volume_marker: str | None = None + self.volume_file_path: str | None = None + self.genie_space_id: str | None = None + self.bundle_dir: pathlib.Path | None = None + self.bundle_env: dict[str, str] = {} + self.bundle_target = "nightly" if options.run_profile == "nightly" else "dev" + self.inputs_path = self.output / "inputs.json" + self.rows_path = self.output / "rows.jsonl" + self.registry_path = self.output / "registry.jsonl" + self.cleanup_path = self.output / "cleanup.jsonl" + self.grant_dir = self.output / "grant-checks" + self._workspace: Workspace | None = None + + @property + def workspace(self) -> Workspace: + if self._workspace is None: + self._workspace = Workspace(self.profile, self.transcript) + return self._workspace + + # Provisioning + + def provision(self) -> None: + for requirement in self.suite.requires: + kind = requirement.get("kind") + if kind == "existing": + self._provision_existing(requirement) + elif kind == "dab": + self._provision_dab(requirement) + elif kind == "uc_function": + self._provision_uc_function(requirement) + else: + raise MatrixError(f"Unknown requirement kind {kind!r} in {self.suite.name}.") + self.genie_space_id = self.resolved.get("genie_space") + + def _provision_existing(self, requirement: dict[str, Any]) -> None: + name = requirement["name"] + value = _expand(str(requirement["value"]), self._overrides()) + pattern = requirement.get("pattern") + if pattern and re.fullmatch(pattern, value) is None: + raise MatrixError(f"Requirement {name!r} value {value!r} does not match {pattern!r}.") + check = requirement.get("check") + if check: + self.databricks([*check, value]) + self.resolved[name] = value + if name == "catalog": + self.catalog = value + + def _provision_dab(self, requirement: dict[str, Any]) -> None: + name = requirement["name"] + catalog = self.resolved.get(requirement["catalog"]) + if not catalog: + raise MatrixError(f"Requirement {name!r} needs {requirement['catalog']!r} first.") + self.bundle_dir = (self.suite.path.parent / requirement["bundle"]).resolve() + self.bundle_env = { + **os.environ, + "BUNDLE_VAR_catalog_name": catalog, + "BUNDLE_VAR_run_id": self.run_id, + } + bundle_args = ["-t", self.bundle_target, *self.profile_args()] + self.run_long( + "bundle-deploy", + ["databricks", "bundle", "deploy", *bundle_args], + cwd=self.bundle_dir, + env=self.bundle_env, + timeout=900, + ) + summary = json.loads( + self.run( + ["databricks", "bundle", "summary", *bundle_args, "--output", "json"], + cwd=self.bundle_dir, + env=self.bundle_env, + timeout=120, + ).stdout + ) + resource = summary.get("resources", {}).get("schemas", {}).get(name) + if not isinstance(resource, dict): + raise MatrixError(f"bundle summary has no schemas.{name}: {summary}") + resolved_id = str(resource.get("id") or "") + if resolved_id.count(".") == 1: + schema = resolved_id + else: + schema = f"{resource.get('catalog_name') or catalog}.{resource.get('name')}" + if schema.count(".") != 1 or "None" in schema: + raise MatrixError(f"Could not resolve the deployed schema from {resource}.") + self.schema = schema + self.resolved[name] = schema + + def _provision_uc_function(self, requirement: dict[str, Any]) -> None: + schema = self.resolved.get(requirement["schema"]) + if not schema: + raise MatrixError(f"Requirement {requirement['name']!r} needs a deployed schema first.") + self.warehouse_id = self.workspace.start_warehouse(self.options.warehouse_id) + catalog, _, schema_name = schema.partition(".") + suffix = uuid.uuid4().hex[:8] + nested_function_name = f"agentbricks_nested_{suffix}" + self.transitive_uc_function = f"{catalog}.{schema_name}.{nested_function_name}" + self.workspace.sql( + f"CREATE OR REPLACE FUNCTION `{catalog}`.`{schema_name}`.`{nested_function_name}`" + "(value STRING) RETURNS STRING " + "COMMENT 'Transitive Agent Bricks E2E marker; never declared in agent.toml' " + "RETURN concat('AGENTBRICKS_UC_OK:', value)" + ) + # Leave room for catalog and schema in the 64-character MCP tool name. + function_name = f"ab_uc_{suffix}" + self.uc_function = f"{catalog}.{schema_name}.{function_name}" + volume_name = f"agentbricks_volume_{suffix}" + self.uc_volume = f"{catalog}.{schema_name}.{volume_name}" + self.workspace.sql(f"CREATE VOLUME `{catalog}`.`{schema_name}`.`{volume_name}`") + self.volume_marker = f"AGENTBRICKS_VOLUME_{uuid.uuid4().hex}" + self.volume_file_path = f"/Volumes/{catalog}/{schema_name}/{volume_name}/marker.txt" + self.workspace.upload(self.volume_file_path, self.volume_marker.encode()) + exposed_tool_name = self.uc_function.replace(".", "__") + if len(exposed_tool_name) > 64: + raise MatrixError( + "The UC function's MCP tool name would exceed 64 characters: " + f"{exposed_tool_name!r}. Use a shorter catalog or run id." + ) + self.workspace.sql( + f"CREATE OR REPLACE FUNCTION `{catalog}`.`{schema_name}`.`{function_name}`" + "(value STRING) RETURNS STRING " + "COMMENT 'Deterministic Agent Bricks E2E marker tool' " + f"RETURN `{catalog}`.`{schema_name}`.`{nested_function_name}`(value)" + ) + self.resolved[requirement["name"]] = self.uc_function + + def _overrides(self) -> dict[str, str]: + return ( + {"AGENTBRICKS_E2E_GENIE_SPACE_ID": self.options.genie_space_id} + if self.options.genie_space_id + else {} + ) + + def write_inputs(self) -> None: + toolchain = self.toolchain + inputs = { + "suite": self.suite.name, + "run_profile": self.options.run_profile, + "databricks_profile": self.profile, + "app_auth_profile": toolchain.app_auth_profile, + "preprovisioned_app_catalog_access": self.options.preprovisioned_app_catalog_access, + "bridge_sha": self.options.bridge_sha, + "wheel": str(toolchain.wheel), + "agentbricks_bin": str(toolchain.agentbricks), + "output_dir": str(self.output), + "run_id": self.run_id, + "run_suffix": self.run_suffix, + "host": toolchain.host, + "catalog": self.catalog, + "schema": self.schema, + "uc_function": self.uc_function, + "transitive_uc_function": self.transitive_uc_function, + "uc_volume": self.uc_volume, + "volume_marker": self.volume_marker, + "volume_file_path": self.volume_file_path, + "genie_space_id": self.genie_space_id, + "warehouse_id": self.warehouse_id, + "requirements": self.resolved, + } + self.inputs_path.parent.mkdir(parents=True, exist_ok=True) + self.inputs_path.write_text(json.dumps(inputs, indent=2), encoding="utf-8") + self.inputs_path.chmod(0o600) + + # Test execution + + def run_pytest(self, budget_seconds: float) -> int: + tests = [str(TESTS_DIR / relative) for relative in self.suite.tests] + argv = [ + sys.executable, + "-m", + "pytest", + *tests, + "-v", + "-ra", + "--capture=tee-sys", + "-p", + "no:cacheprovider", + "--junitxml", + str(self.output / "junit.xml"), + ] + env = {**os.environ, INPUTS_ENV: str(self.inputs_path)} + self.transcript.command(argv) + process = subprocess.Popen(argv, cwd=TESTS_DIR, env=env, start_new_session=True) + try: + return process.wait(timeout=budget_seconds) + except subprocess.TimeoutExpired: + self.transcript.write(f"pytest exceeded its {budget_seconds:.0f}s budget; stopping it.") + # SIGINT lets pytest unwind its fixtures (the dev servers) before teardown starts. + _signal_group(process, signal.SIGINT) + try: + process.wait(timeout=120) + except subprocess.TimeoutExpired: + _signal_group(process, signal.SIGKILL) + process.wait() + return 124 + + # Teardown + + def _registry(self) -> tuple[list[str], dict[str, dict[str, Any]], set[str]]: + apps: list[str] = [] + projects: dict[str, dict[str, Any]] = {} + cleaned: set[str] = set() + if self.registry_path.exists(): + for line in self.registry_path.read_text(encoding="utf-8").splitlines(): + if not line.strip(): + continue + event = json.loads(line) + name = event["app_name"] + if event["kind"] == "app" and name not in apps: + apps.append(name) + elif event["kind"] == "project": + projects.setdefault(name, {}).update(event) + elif event["kind"] == "cleaned": + cleaned.add(name) + return apps, projects, cleaned + + def teardown(self) -> None: + # Each test deletes its own App and stores; this sweeps whatever a test could not, such as + # after a crash, and folds the tests' own cleanup results into the evidence. + apps, projects, cleaned = self._registry() + self.cleanup_results.extend(self._read_jsonl(self.cleanup_path)) + for name in dict.fromkeys([*apps, *projects]): + if name not in cleaned: + self.cleanup_results.extend( + cleanup_app( + self.workspace, + functools.partial( + delete_store, + self.toolchain.agentbricks, + self.profile, + self.transcript, + ), + name, + projects.get(name, {}), + has_app=name in apps, + ) + ) + for label, function in ( + ("UC function", self.uc_function), + ("transitive UC function", self.transitive_uc_function), + ): + if not function: + continue + catalog, schema, function_name = function.split(".") + try: + self.workspace.sql( + f"DROP FUNCTION IF EXISTS `{catalog}`.`{schema}`.`{function_name}`" + ) + self.cleanup_results.append( + {"resource": f"function:{function}", "status": "deleted"} + ) + except Exception as exc: + self.transcript.write(f"cleanup warning | {label} | {exc}") + self.cleanup_results.append( + {"resource": f"function:{function}", "status": "failed", "detail": str(exc)} + ) + if self.uc_volume: + catalog, schema, volume_name = self.uc_volume.split(".") + if self.volume_file_path: + try: + self.workspace.delete_file(self.volume_file_path) + self.cleanup_results.append( + {"resource": f"file:{self.volume_file_path}", "status": "deleted"} + ) + except Exception as exc: + self.transcript.write(f"cleanup warning | UC volume file | {exc}") + self.cleanup_results.append( + { + "resource": f"file:{self.volume_file_path}", + "status": "failed", + "detail": str(exc), + } + ) + try: + self.workspace.sql(f"DROP VOLUME IF EXISTS `{catalog}`.`{schema}`.`{volume_name}`") + self.cleanup_results.append( + {"resource": f"volume:{self.uc_volume}", "status": "deleted"} + ) + except Exception as exc: + self.transcript.write(f"cleanup warning | UC volume | {exc}") + self.cleanup_results.append( + { + "resource": f"volume:{self.uc_volume}", + "status": "failed", + "detail": str(exc), + } + ) + if self.bundle_dir is not None: + resource_name = f"bundle:{self.schema or self.suite.name}" + try: + self.run_long( + "bundle-destroy", + [ + "databricks", + "bundle", + "destroy", + "--auto-approve", + "-t", + self.bundle_target, + *self.profile_args(), + ], + cwd=self.bundle_dir, + env=self.bundle_env, + timeout=900, + ) + self.cleanup_results.append({"resource": resource_name, "status": "deleted"}) + except Exception as exc: + self.transcript.write(f"cleanup warning | bundle destroy | {exc}") + self.cleanup_results.append( + {"resource": resource_name, "status": "failed", "detail": str(exc)} + ) + self.cleanup_complete = not any( + result.get("status") == "failed" for result in self.cleanup_results + ) + + # Evidence + + def _read_jsonl(self, path: pathlib.Path) -> list[dict[str, Any]]: + if not path.exists(): + return [] + return [ + json.loads(line) + for line in path.read_text(encoding="utf-8").splitlines() + if line.strip() + ] + + def _grant_checks(self) -> list[dict[str, Any]]: + checks = [] + for path in sorted(self.grant_dir.glob("*.json")) if self.grant_dir.exists() else (): + recorded = json.loads(path.read_text(encoding="utf-8")) + checks.append( + { + "app_name": recorded.get("app_name"), + "authoring": recorded.get("authoring"), + "service_principal_client_id": recorded.get("service_principal_client_id"), + "initial": recorded.get("initial", {}), + "post_manual_transitive_grant": recorded.get( + "post_manual_transitive_grant", {} + ), + "manual_transitive_grant_applied": recorded.get( + "manual_transitive_grant_applied", False + ), + } + ) + return checks + + def write_evidence(self) -> pathlib.Path: + toolchain = self.toolchain + wheel = toolchain.wheel + payload = { + "schema_version": 1, + "suite": self.suite.name, + "run_profile": self.options.run_profile, + "run_id": self.run_id, + "commit_sha": toolchain.commit_sha, + "source_provenance": toolchain.source_provenance, + "template_repo": "wheel://databricks-agentbricks", + "template_ref": _sha256(wheel) if wheel else None, + "workspace_host": toolchain.host, + "started_at": self.started_at, + "ended_at": self.ended_at, + "versions": toolchain.versions, + "profile": self.profile, + "app_auth_profile": toolchain.app_auth_profile, + "bridge_sha": self.options.bridge_sha, + "preprovisioned_app_catalog_access": self.options.preprovisioned_app_catalog_access, + "generated_at": now().isoformat(), + "wheel": str(wheel) if wheel else None, + "wheel_sha256": _sha256(wheel) if wheel else None, + "schema": self.schema, + "uc_function": self.uc_function, + "transitive_uc_function": self.transitive_uc_function, + "uc_volume": self.uc_volume, + "volume_marker": self.volume_marker, + "volume_file_path": self.volume_file_path, + "genie_space_id": self.genie_space_id, + "warehouse_id": self.warehouse_id, + "grant_checks": self._grant_checks(), + "cleanup_required": self.cleanup_required, + "cleanup_complete": self.cleanup_complete, + "cleanup": self.cleanup_results, + "rows": self._read_jsonl(self.rows_path), + } + target = self.output / "evidence.json" + temporary = target.with_suffix(".json.tmp") + temporary.write_text(json.dumps(payload, indent=2), encoding="utf-8") + os.replace(temporary, target) + return target + + # Orchestration + + def execute(self) -> int: + timeout_minutes = self.suite.runs[self.options.run_profile]["timeout_minutes"] + started = time.monotonic() + pytest_returncode: int | None = None + try: + self.provision() + self.write_inputs() + remaining = ( + timeout_minutes * 60 - TEARDOWN_RESERVE_SECONDS - (time.monotonic() - started) + ) + pytest_returncode = self.run_pytest(max(60.0, remaining)) + except Exception as exc: + self.transcript.write(f"suite {self.suite.name} failed before tests finished: {exc}") + finally: + if self.cleanup_required: + self.teardown() + else: + self.transcript.write("Resources retained (--keep-resources); skipping teardown.") + self.ended_at = now().isoformat() + evidence = self.write_evidence() + verified = verify_evidence(evidence) + return 0 if pytest_returncode == 0 and verified == 0 else 1 + + +def _signal_group(process: subprocess.Popen[Any], sig: signal.Signals) -> None: + try: + os.killpg(process.pid, sig) + except ProcessLookupError: + pass + + +def _expand(value: str, overrides: dict[str, str]) -> str: + def replace(match: re.Match[str]) -> str: + name, default = match.group(1), match.group(2) + resolved = overrides.get(name) or os.environ.get(name) or default + if not resolved: + raise MatrixError(f"Required environment variable {name} is not set.") + return resolved + + return _ENV_REFERENCE.sub(replace, value) + + +def _sha256(path: pathlib.Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as input_file: + for chunk in iter(lambda: input_file.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _source_provenance( + source_root: pathlib.Path, commit_sha: str, wheel: pathlib.Path +) -> dict[str, Any]: + def git(*args: str, text: bool = True) -> subprocess.CompletedProcess: + result = subprocess.run( + ["git", *args], + cwd=source_root, + capture_output=True, + text=text, + check=False, + ) + if result.returncode != 0: + detail = result.stderr if text else result.stderr.decode(errors="replace") + raise MatrixError(f"Could not inspect source provenance: {detail.strip()}") + return result + + head = git("rev-parse", "HEAD").stdout.strip().lower() + if head != commit_sha.lower(): + raise MatrixError( + f"Claimed commit {commit_sha} does not match source checkout HEAD {head}." + ) + status = git("status", "--porcelain", "--untracked-files=no").stdout.strip() + source_hashes: dict[str, str] = {} + try: + with zipfile.ZipFile(wheel) as archive: + for member in _WHEEL_SOURCE_FILES: + source_path = source_root / "integrations" / "agentbricks" / "src" / member + source_bytes = source_path.read_bytes() + try: + wheel_bytes = archive.read(member) + except KeyError as exc: + raise MatrixError(f"Built wheel is missing source module {member}.") from exc + if wheel_bytes != source_bytes: + raise MatrixError( + f"Built wheel module {member} does not match source checkout." + ) + source_hashes[member] = hashlib.sha256(wheel_bytes).hexdigest() + except zipfile.BadZipFile as exc: + raise MatrixError(f"Built wheel is not a readable zip archive: {wheel}") from exc + diff = git("diff", "--binary", "HEAD", text=False).stdout + return { + "source_head_sha": head, + "source_dirty": bool(status), + "source_diff_sha256": hashlib.sha256(diff).hexdigest(), + "wheel_source_matches": True, + "wheel_source_sha256": source_hashes, + } + + +def _verify_tools_evidence(document: dict[str, Any], rows: list[dict[str, Any]]) -> int: + """Sandbox-marker and grant proofs, which only the tools matrix produces.""" + grant_checks = document.get("grant_checks", []) + sandbox_markers = { + "sandbox": document.get("volume_marker"), + } + volume_file_path = document.get("volume_file_path") + if ( + any(not isinstance(marker, str) or not marker for marker in sandbox_markers.values()) + or not isinstance(volume_file_path, str) + or not volume_file_path.startswith("/Volumes/") + ): + sys.stdout.write("sandbox evidence: hidden marker provenance is missing or invalid\n") + return 1 + for tool_kind, marker in sandbox_markers.items(): + matching_rows = [row for row in rows if row.get("tool_kind") == tool_kind] + if any( + row.get("expected") != marker or marker not in str(row.get("actual", "")) + for row in matching_rows + ): + sys.stdout.write( + f"sandbox evidence: exact hidden marker proof failed for {tool_kind}\n" + ) + return 1 + # One deploy per authoring path proves the transitive exclusion and the manual grant. Repeat-deploy + # idempotency is proven by the incremental CLI deploy rows, which re-check earlier tools' grants. + transitive = [check for check in grant_checks if check.get("manual_transitive_grant_applied")] + expected_transitive = len(FRAMEWORKS) * len(AUTHORING_PATHS) + if len(transitive) != expected_transitive or any( + "EXECUTE" in check.get("initial", {}).get("transitive_direct_privileges", []) + or "EXECUTE" in check.get("initial", {}).get("transitive_effective_privileges", []) + or "EXECUTE" + not in check.get("post_manual_transitive_grant", {}).get("direct_privileges", []) + or "EXECUTE" + not in check.get("post_manual_transitive_grant", {}).get("effective_privileges", []) + for check in transitive + ): + sys.stdout.write( + f"grant evidence: expected {expected_transitive} transitive exclusion/manual-grant " + f"proofs, got {len(transitive)} (or one failed)\n" + ) + return 1 + return 0 + + +def verify_evidence(path: pathlib.Path, *, require_cleanup: bool = True) -> int: + document = json.loads(path.read_text(encoding="utf-8")) + provenance_fields = ("commit_sha", "workspace_host", "started_at", "ended_at") + missing_provenance = [field for field in provenance_fields if not document.get(field)] + template_repo = document.get("template_repo") + template_ref = document.get("template_ref") + if ( + not isinstance(template_repo, str) + or not template_repo + or not isinstance(template_ref, str) + or not template_ref + ): + missing_provenance.append("template_repo/template_ref") + elif template_repo.startswith("wheel://") and template_ref != document.get("wheel_sha256"): + missing_provenance.append("wheel_template_ref") + source_provenance = document.get("source_provenance") + if ( + not isinstance(source_provenance, dict) + or source_provenance.get("source_head_sha") != document.get("commit_sha") + or source_provenance.get("wheel_source_matches") is not True + or not isinstance(source_provenance.get("wheel_source_sha256"), dict) + or not source_provenance.get("wheel_source_sha256") + ): + missing_provenance.append("source_provenance") + versions = document.get("versions") + if not isinstance(versions, dict) or any( + not versions.get(name) for name in ("agentbricks", "databricks", "uv", "python") + ): + missing_provenance.append("versions") + if missing_provenance: + sys.stdout.write(f"evidence provenance missing: {sorted(missing_provenance)}\n") + return 1 + suite = document.get("suite") + kinds = MATRIX_KINDS.get(suite) + if kinds is None: + sys.stdout.write( + f"evidence suite {suite!r} is not a known matrix: {sorted(MATRIX_KINDS)}\n" + ) + return 1 + rows = document.get("rows", []) + expected = { + (framework, authoring, runtime, tool) + for framework in FRAMEWORKS + for authoring, runtime, tool in cells + } + actual = { + (row["framework"], row["authoring"], row["runtime"], row["tool_kind"]) for row in rows + } + duplicates = len(rows) - len(actual) + passed = sum(row.get("status") == "pass" for row in rows) + failed = sum(row.get("status") == "fail" for row in rows) + skipped = len(expected - actual) + sys.stdout.write(f"{passed} passed, {failed} failed, {skipped} skipped\n") + if actual != expected or duplicates or passed != len(expected): + if expected - actual: + sys.stdout.write(f"missing cells: {sorted(expected - actual)}\n") + if duplicates: + sys.stdout.write(f"duplicate rows: {duplicates}\n") + return 1 + if suite == "tools" and _verify_tools_evidence(document, rows): + return 1 + cleanup_required = document.get("cleanup_required") + if not isinstance(cleanup_required, bool): + sys.stdout.write("cleanup evidence: cleanup_required was not recorded\n") + return 1 + cleanup = document.get("cleanup", []) + if ( + require_cleanup + and cleanup_required + and ( + document.get("cleanup_complete") is not True + or not isinstance(cleanup, list) + or any(result.get("status") == "failed" for result in cleanup) + or any( + result.get("resource", "").startswith("app:") + and result.get("status") == "deleted" + and not result.get("confirmed_absent_at") + for result in cleanup + ) + ) + ): + sys.stdout.write("cleanup evidence: required cleanup is incomplete or failed\n") + if isinstance(cleanup, list): + for result in cleanup: + if result.get("status") == "failed": + sys.stdout.write( + f"cleanup failed | {result.get('resource')} | {result.get('detail')}\n" + ) + # teardown() always stamps confirmed_absent_at on deleted apps, so + # this only fires for evidence from an older runner or a hand-edited file. + elif ( + result.get("resource", "").startswith("app:") + and result.get("status") == "deleted" + and not result.get("confirmed_absent_at") + ): + sys.stdout.write( + f"cleanup unconfirmed | {result.get('resource')} |" + " missing confirmed_absent_at\n" + ) + return 1 + if suite == "tools": + grant_checks = document.get("grant_checks", []) + transitive = [c for c in grant_checks if c.get("manual_transitive_grant_applied")] + sys.stdout.write(f"{len(transitive)} transitive grant proofs passed\n") + return 0 + + +def load_suites() -> list[Suite]: + suites = [] + for path in sorted(MATRICES_DIR.glob("*/suite.yaml")): + document = yaml.safe_load(path.read_text(encoding="utf-8")) + for requirement in document.get("requires", []): + if requirement.get("kind") not in REQUIREMENT_KINDS: + raise MatrixError(f"{path.parent.name}: unknown requirement kind in {requirement}.") + suites.append( + Suite( + name=document["name"], + path=path, + tests=list(document["tests"]), + runs=document.get("runs", {}), + requires=list(document.get("requires", [])), + ) + ) + return suites + + +def select_suites(profile: str, names: Sequence[str]) -> list[Suite]: + available = load_suites() + if names: + by_name = {suite.name: suite for suite in available} + unknown = sorted(set(names) - set(by_name)) + if unknown: + raise MatrixError(f"Unknown suite(s) {unknown}; available: {sorted(by_name)}") + selected = [by_name[name] for name in dict.fromkeys(names)] + missing = [suite.name for suite in selected if profile not in suite.runs] + if missing: + raise MatrixError(f"Suite(s) {missing} do not run under the {profile!r} profile.") + return selected + return [suite for suite in available if profile in suite.runs] + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--profile", choices=PROFILES, help="Run profile selecting suites.") + parser.add_argument("--suite", action="append", default=[], help="Suite name; repeatable.") + parser.add_argument( + "--output", type=pathlib.Path, help="Output directory (default: a tempdir)." + ) + parser.add_argument("--wheel", type=pathlib.Path, help="Prebuilt wheel (default: build one).") + parser.add_argument("--commit-sha", help="Source commit (default: HEAD of --source-root).") + parser.add_argument("--source-root", type=pathlib.Path, default=REPOSITORY_ROOT) + parser.add_argument( + "--genie-space-id", + help="Overrides AGENTBRICKS_E2E_GENIE_SPACE_ID for the genie_space requirement.", + ) + parser.add_argument( + "--bridge-sha", help="Immutable bridge commit for generated App dependencies." + ) + parser.add_argument("--warehouse-id") + parser.add_argument( + "--databricks-profile", + help="Databricks CLI profile; default is ambient credentials (DATABRICKS_CONFIG_PROFILE).", + ) + parser.add_argument( + "--app-auth-profile", + help="OAuth profile for deployed App /api calls; defaults to --databricks-profile.", + ) + parser.add_argument( + "--preprovisioned-app-catalog-access", + action="store_true", + help="Skip per-App USE CATALOG grants because catalog access is pre-provisioned.", + ) + parser.add_argument("--keep-resources", action="store_true") + parser.add_argument("--verify-evidence", type=pathlib.Path) + args = parser.parse_args() + if args.verify_evidence is None and args.profile is None: + parser.error("--profile is required unless --verify-evidence is used") + return args + + +def main() -> int: + args = parse_args() + if args.verify_evidence: + return verify_evidence(args.verify_evidence) + try: + suites = select_suites(args.profile, args.suite) + except MatrixError as exc: + sys.stderr.write(f"{exc}\n") + return 2 + if not suites: + sys.stderr.write(f"No suites run under the {args.profile!r} profile.\n") + return 2 + output = args.output or pathlib.Path(tempfile.mkdtemp(prefix="agentbricks-suites-")) + options = Options( + run_profile=args.profile, + output=output.expanduser().resolve(), + wheel=args.wheel, + commit_sha=args.commit_sha, + source_root=args.source_root, + genie_space_id=args.genie_space_id or os.environ.get("AGENTBRICKS_E2E_GENIE_SPACE_ID"), + bridge_sha=args.bridge_sha, + warehouse_id=args.warehouse_id, + databricks_profile=args.databricks_profile, + app_auth_profile=args.app_auth_profile, + preprovisioned_app_catalog_access=args.preprovisioned_app_catalog_access, + keep_resources=args.keep_resources, + ) + toolchain = Toolchain(options) + try: + toolchain.bootstrap() + except MatrixError as exc: + toolchain.transcript.write(f"bootstrap failed: {exc}") + return 1 + failed = [ + suite.name for suite in suites if SuiteRunner(suite, toolchain, options).execute() != 0 + ] + if failed: + sys.stdout.write(f"suites failed: {failed}; output: {options.output}\n") + return 1 + sys.stdout.write(f"suites passed: {[suite.name for suite in suites]}\n") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/integrations/agentbricks/tests/workspace_client.py b/integrations/agentbricks/tests/workspace_client.py new file mode 100644 index 000000000..1158140e9 --- /dev/null +++ b/integrations/agentbricks/tests/workspace_client.py @@ -0,0 +1,514 @@ +"""Databricks SDK helpers for the matrix tests and runner: the workspace, which is not under test. + +Everything here goes through the SDK except ``app_logs``, since the SDK has no Apps log API, and +bundles, which the runner drives with the Databricks CLI. +""" + +from __future__ import annotations + +import datetime as dt +import io +import pathlib +import re +import subprocess +import time +from collections.abc import Callable, Sequence +from typing import Any + +from common import TOOL_RESOURCE_PREFIX, MatrixError, Transcript, now +from databricks.sdk import WorkspaceClient +from databricks.sdk.errors import DatabricksError, NotFound +from databricks.sdk.service.catalog import SecurableType +from databricks.sdk.service.sql import ExecuteStatementRequestOnWaitTimeout, State, StatementState + +GrantTuple = tuple[str, str, str, str] + + +class Workspace: + def __init__( + self, + profile: str | None, + transcript: Transcript, + *, + app_auth_profile: str | None = None, + warehouse_id: str | None = None, + preprovisioned_app_catalog_access: bool = False, + ): + self.profile = profile + self.app_auth_profile = app_auth_profile or profile + self.transcript = transcript + self.warehouse_id = warehouse_id + self.preprovisioned_app_catalog_access = preprovisioned_app_catalog_access + self.client = WorkspaceClient(profile=profile) + self._app_auth_client: WorkspaceClient | None = None + + # App auth: Databricks Apps /api routes need OAuth, so a PAT is rejected. + + def _app_auth_label(self) -> str: + return ( + f"App auth profile {self.app_auth_profile!r}" + if self.app_auth_profile + else "ambient environment credentials" + ) + + def _app_client(self) -> WorkspaceClient: + if self._app_auth_client is None: + self._app_auth_client = WorkspaceClient(profile=self.app_auth_profile) + return self._app_auth_client + + def check_app_auth(self) -> None: + if self._app_client().config.auth_type == "pat": + raise MatrixError( + f"{self._app_auth_label()} uses a PAT. " + "Databricks Apps /api routes require OAuth; run `databricks auth login` " + "for a profile on the same workspace." + ) + + @property + def app_headers(self) -> dict[str, str]: + # Resolved per call so OAuth tokens refresh during a long run. + authorization = self._app_client().config.authenticate().get("Authorization") + if not authorization: + raise MatrixError(f"Could not resolve credentials from {self._app_auth_label()}.") + return {"Authorization": authorization} + + # 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 MatrixError("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) + self.transcript.write(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 sql(self, statement: str, *, timeout: float = 600) -> None: + if self.warehouse_id is None: + raise MatrixError("SQL warehouse was not selected.") + self.transcript.write(f"$ sql: {statement}") + response = self.client.statement_execution.execute_statement( + statement=statement, + warehouse_id=self.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 MatrixError(f"SQL response has no statement_id: {response}") + if timeout <= 0: + raise MatrixError(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 MatrixError(f"SQL failed: {response.as_dict()}") + + def upload(self, path: str, data: bytes) -> None: + self.client.files.upload(path, io.BytesIO(data), overwrite=True) + self.transcript.write(f"# uploaded {path}") + + def delete_file(self, path: str) -> None: + try: + self.client.files.delete(path) + except NotFound: + pass + + # Apps + + def app(self, name: str) -> dict[str, Any]: + return self.client.apps.get(name).as_dict() + + def assert_app_absent(self, name: str) -> None: + try: + self.client.apps.get(name) + except NotFound: + return + except DatabricksError as exc: + raise MatrixError(f"Could not verify App {name} is absent: {exc}") from exc + raise MatrixError(f"App {name} already exists; refusing to deploy over it.") + + def wait_for_app(self, name: str) -> dict[str, Any]: + """The App once it is ACTIVE with a URL.""" + started = time.monotonic() + next_tick = 0.0 + while time.monotonic() - started < 1200: + app = self.app(name) + compute = app.get("compute_status", {}) + state = compute.get("state") if isinstance(compute, dict) else None + if state == "ACTIVE" and app.get("url"): + return app + elapsed = time.monotonic() - started + if elapsed >= next_tick: + self.transcript.write(f"tick {now():%H:%M} | app-{name} | {state or 'UNKNOWN'}") + next_tick += 60 + time.sleep(15) + raise MatrixError(f"App {name} did not become ACTIVE.") + + 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: + self.transcript.write(f"App runtime log capture warning for {name}: {exc}") + return None + self.transcript.write(f"App runtime logs captured: {log_path}") + return log_path + + # Grants + + def granted(self, app_name: str, grant: GrantTuple) -> bool: + """Whether Agent Bricks granted the App ``grant``, as (kind, name, type, permission).""" + return grant in app_resource_tuples(tool_resources(self.app(app_name))) + + def granted_tuples(self, app_name: str) -> set[GrantTuple]: + return app_resource_tuples(tool_resources(self.app(app_name))) + + def experiment_id(self, name: str) -> str: + experiment = self.client.experiments.get_by_name(name).experiment + if experiment is None or not experiment.experiment_id: + raise MatrixError(f"Experiment {name!r} does not exist.") + return str(experiment.experiment_id) + + def non_tool_resources(self, app_name: str) -> list[dict[str, Any]]: + return non_tool_resources(self.app(app_name)) + + def transitive_state(self, app_name: str, function: str) -> dict[str, Any]: + """The App principal's privileges on a nested function before any manual grant.""" + app = self.app(app_name) + principal = self._principal(app) + direct = direct_privileges(self.client, SecurableType.FUNCTION, function, principal) + if "EXECUTE" in direct: + raise MatrixError( + f"Agent Bricks granted the transitive function directly: {function} -> {direct}" + ) + effective = effective_privileges(self.client, SecurableType.FUNCTION, function, principal) + if "EXECUTE" in effective: + raise MatrixError( + "The transitive function was already effective before the manual grant: " + f"{function} -> {effective}" + ) + if any( + resource.get("uc_securable", {}).get("securable_full_name") == function + for resource in tool_resources(app) + ): + raise MatrixError( + "The transitive function appeared in Agent Bricks-owned Apps resources." + ) + return { + "service_principal_client_id": principal, + "direct_privileges": direct, + "effective_privileges": effective, + } + + def grant_transitive(self, app_name: str, function: str) -> dict[str, Any]: + """Grant the App principal EXECUTE on a nested function and confirm it is direct and effective.""" + principal = self._principal(self.app(app_name)) + catalog, schema, function_name = function.split(".") + quoted_principal = f"`{principal.replace('`', '``')}`" + statements = [] + if not self.preprovisioned_app_catalog_access: + statements.append(f"GRANT USE CATALOG ON CATALOG `{catalog}` TO {quoted_principal}") + statements.extend( + ( + f"GRANT USE SCHEMA ON SCHEMA `{catalog}`.`{schema}` TO {quoted_principal}", + f"GRANT EXECUTE ON FUNCTION `{catalog}`.`{schema}`.`{function_name}` " + f"TO {quoted_principal}", + ) + ) + for statement in statements: + self.sql(statement) + direct = direct_privileges(self.client, SecurableType.FUNCTION, function, principal) + effective = effective_privileges(self.client, SecurableType.FUNCTION, function, principal) + if "EXECUTE" not in effective: + raise MatrixError( + "Manual EXECUTE grant did not become effective for the transitive function: " + f"{function} -> {effective}" + ) + if "EXECUTE" not in direct: + raise MatrixError( + "Manual EXECUTE grant was not persisted directly for the transitive function: " + f"{function} -> {direct}" + ) + return { + "granted_at": now().isoformat(), + "direct_privileges": direct, + "effective_privileges": effective, + } + + @staticmethod + def _principal(app: dict[str, Any]) -> str: + principal = app.get("service_principal_client_id") + if not principal: + raise MatrixError(f"App response has no service_principal_client_id: {app}") + return str(principal) + + # Cleanup + + def delete_app(self, name: str) -> None: + self.client.apps.delete(name) + self._wait_for_app_deleted(name) + + def _wait_for_app_deleted(self, name: str, timeout: float = 1200) -> None: + started = time.monotonic() + next_tick = 0.0 + while True: + try: + self.client.apps.get(name) + except NotFound: + self.transcript.write(f"tick {now():%H:%M} | delete-{name} | absent") + return + except DatabricksError as exc: + raise MatrixError(f"Could not verify deletion of App {name!r}: {exc}") from exc + elapsed = time.monotonic() - started + if elapsed >= timeout: + raise MatrixError( + f"App {name!r} still existed {timeout:.0f}s after delete returned." + ) + if elapsed >= next_tick: + self.transcript.write(f"tick {now():%H:%M} | delete-{name} | deleting") + next_tick += 60 + time.sleep(15) + + def runtime_store(self, app_name: str) -> dict[str, Any]: + return self.client.api_client.do("GET", f"/api/2.0/agents/runtime-stores/{app_name}") + + def delete_runtime_store(self, app_name: str) -> None: + """Delete the deploy-created Runtime Store; a missing store counts as deleted.""" + try: + self.client.api_client.do("DELETE", f"/api/2.0/agents/runtime-stores/{app_name}") + except NotFound: + pass + + def app_role_target(self, app_name: str) -> str | None: + """The Lakebase role of the App's principal, only if ownership is verified end to end.""" + try: + principal = self.app(app_name).get("service_principal_client_id") + store = self.runtime_store(app_name) + owner = store.get("owner", {}).get("app", {}) + branch = store.get("storage_backend", {}).get("lakebase", {}).get("branch") + if ( + not principal + or store.get("name") != f"runtime-stores/{app_name}" + or owner.get("name") != app_name + or owner.get("service_principal_id") != principal + or not isinstance(branch, str) + or not re.fullmatch(r"projects/[^/]+/branches/[^/]+", branch) + ): + self.transcript.write( + f"cleanup warning | Lakebase role for {app_name} | ownership not verified" + ) + return None + for role in self.client.postgres.list_roles(parent=branch): + data = role.as_dict() + status = data.get("status", {}) + name = data.get("name") + if ( + status.get("postgres_role") == principal + and status.get("identity_type") == "SERVICE_PRINCIPAL" + and isinstance(name, str) + and name.startswith(f"{branch}/roles/") + ): + return name + except Exception as exc: + self.transcript.write(f"cleanup warning | Lakebase role lookup for {app_name} | {exc}") + return None + + def delete_role(self, role_name: str) -> None: + operation = self.client.postgres.delete_role(name=role_name) + wait = getattr(operation, "wait", None) + if callable(wait): + wait() + + +def cleanup_app( + workspace: Workspace, + delete_store: Callable[[str, str], subprocess.CompletedProcess[str]], + name: str, + info: dict[str, Any], + *, + has_app: bool, +) -> list[dict[str, Any]]: + """Delete one project's stores, and its App, runtime store and Lakebase role if it has an App. + + ``delete_store(kind, store)`` is the agentbricks CLI's store delete, which has no SDK surface. + """ + transcript = workspace.transcript + results: list[dict[str, Any]] = [] + has_stores = bool(info.get("memory_store_name") or info.get("session_store_name")) + # Looked up before anything is deleted, since it needs the App and its runtime store. + role_target = workspace.app_role_target(name) if has_app and has_stores else None + stores_ok = True + for label, kind, store in ( + ("memory store", "memory", info.get("memory_store_name")), + ("session store", "sessions", info.get("session_store_name")), + ): + if not store: + continue + try: + result = delete_store(kind, store) + if result.returncode != 0: + stores_ok = False + results.append( + { + "resource": f"{label}:{store}", + "status": "failed", + "detail": (result.stderr or result.stdout).strip(), + } + ) + else: + results.append({"resource": f"{label}:{store}", "status": "deleted"}) + except Exception as exc: + stores_ok = False + transcript.write(f"cleanup warning | {label} {store} | {exc}") + results.append({"resource": f"{label}:{store}", "status": "failed", "detail": str(exc)}) + if not has_app: + return results + + runtime_store_deleted = False + if has_stores: + # The Runtime Store owns a dedicated Lakebase database; it must be deleted before the + # role, or the role delete fails on database ownership. + try: + workspace.delete_runtime_store(name) + runtime_store_deleted = True + results.append({"resource": f"runtime-store:{name}", "status": "deleted"}) + except Exception as exc: + results.append( + {"resource": f"runtime-store:{name}", "status": "failed", "detail": str(exc)} + ) + try: + workspace.delete_app(name) + except Exception as exc: + results.append({"resource": f"app:{name}", "status": "failed", "detail": str(exc)}) + return results + results.append( + {"resource": f"app:{name}", "status": "deleted", "confirmed_absent_at": now().isoformat()} + ) + if role_target and stores_ok and runtime_store_deleted: + try: + workspace.delete_role(role_target) + results.append({"resource": f"lakebase-role:{role_target}", "status": "deleted"}) + except Exception as exc: + transcript.write(f"cleanup warning | Lakebase role {role_target} | {exc}") + results.append( + { + "resource": f"lakebase-role:{role_target}", + "status": "failed", + "detail": str(exc), + } + ) + return results + + +# Grant helpers + + +def _split_resources(app: dict[str, Any]) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: + resources = app.get("resources") or [] + if not isinstance(resources, list): + raise MatrixError(f"App resources are not a list: {resources}") + + def name(resource: dict[str, Any]) -> str: + return str(resource.get("name", "")) + + owned = [r for r in resources if name(r).startswith(TOOL_RESOURCE_PREFIX)] + other = [r for r in resources if not name(r).startswith(TOOL_RESOURCE_PREFIX)] + return sorted(owned, key=name), sorted(other, key=name) + + +def tool_resources(app: dict[str, Any]) -> list[dict[str, Any]]: + """The App resources Agent Bricks owns, i.e. those it names with ``TOOL_RESOURCE_PREFIX``.""" + return _split_resources(app)[0] + + +def non_tool_resources(app: dict[str, Any]) -> list[dict[str, Any]]: + """The App resources Agent Bricks does not own, e.g. the tracing experiment.""" + return _split_resources(app)[1] + + +def app_resource_tuples(resources: Sequence[dict[str, Any]]) -> set[GrantTuple]: + actual = set() + for resource in resources: + if "uc_securable" in resource: + value = resource["uc_securable"] + actual.add( + ( + "uc_securable", + value.get("securable_full_name"), + value.get("securable_type"), + value.get("permission"), + ) + ) + elif "genie_space" in resource: + value = resource["genie_space"] + actual.add( + ("genie_space", value.get("space_id"), "GENIE_SPACE", value.get("permission")) + ) + return actual + + +def effective_privileges( + client: WorkspaceClient, securable_type: SecurableType, full_name: str, principal: str +) -> list[str]: + privileges: set[str] = set() + page_token: str | None = None + while True: + response = client.grants.get_effective( + securable_type.value, + full_name, + max_results=0, + principal=principal, + **({"page_token": page_token} if page_token else {}), + ) + for assignment in response.privilege_assignments or (): + if assignment.principal != principal: + continue + privileges.update( + privilege.privilege.value + for privilege in assignment.privileges or () + if privilege.privilege is not None + ) + page_token = response.next_page_token + if not page_token: + return sorted(privileges) + + +def direct_privileges( + client: WorkspaceClient, securable_type: SecurableType, full_name: str, principal: str +) -> list[str]: + privileges: set[str] = set() + page_token: str | None = None + while True: + response = client.grants.get( + securable_type.value, + full_name, + max_results=0, + principal=principal, + **({"page_token": page_token} if page_token else {}), + ) + for assignment in response.privilege_assignments or (): + if assignment.principal == principal: + privileges.update(privilege.value for privilege in assignment.privileges or ()) + page_token = response.next_page_token + if not page_token: + return sorted(privileges) From 4bb52a53d7f426cfe0f99021e94236efe719f469 Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Wed, 7 Oct 2026 23:49:04 +0000 Subject: [PATCH 02/45] [agentbricks] drop unbind from e2e single-feature authoring Co-authored-by: Isaac --- .../agentbricks/tests/agentbricks_cli.py | 41 +++++-------------- .../memory_store/test_memory_store.py | 4 +- .../session_store/test_session_store.py | 4 +- .../tests/matrices/tools/test_grants.py | 4 +- .../tests/matrices/tracing/test_tracing.py | 4 +- 5 files changed, 19 insertions(+), 38 deletions(-) diff --git a/integrations/agentbricks/tests/agentbricks_cli.py b/integrations/agentbricks/tests/agentbricks_cli.py index 83b727bd9..be5c31dce 100644 --- a/integrations/agentbricks/tests/agentbricks_cli.py +++ b/integrations/agentbricks/tests/agentbricks_cli.py @@ -41,12 +41,6 @@ ) from workspace_client import Workspace, cleanup_app -# `agentbricks init` binds a session store, memory store, and tracing experiment by default. -_UNBIND_COMMANDS = { - "session_store": ("sessions", "unbind"), - "memory_store": ("memory", "unbind"), - "tracing": ("tracing", "unbind"), -} _STORE_KINDS = {"memory": ("memory", "stores"), "sessions": ("sessions", "stores")} _DIRECT_HEADER = ( 'schema_version = 1\n\n[agent]\nframework = "{framework}"\nserver = "agentbricks"\n' @@ -65,7 +59,6 @@ class Project: authoring: str path: pathlib.Path app_name: str - features: tuple[str, ...] = () memory_store_config: str | None = None session_store_config: str | None = None memory_resource: dict[str, Any] | None = None @@ -235,19 +228,15 @@ def _next(self) -> int: # Authoring - def new_project(self, authoring: str, keep: Sequence[str] = ()) -> Project: - """`agentbricks init` a uniquely named project, pin the runtime under test, and leave only - the ``keep`` features (session_store, memory_store, tracing) bound. + def new_project(self, authoring: str) -> Project: + """`agentbricks init` a uniquely named project and pin the runtime under test. - CLI authoring unbinds every other default feature. Direct authoring replaces the whole - manifest later in ``write_manifest``, which binds ``keep`` itself. + `init` binds a session store, a memory store and a tracing experiment by default, and CLI + authoring keeps them. Direct authoring replaces the whole manifest in ``write_manifest``. """ - unknown = set(keep) - set(_UNBIND_COMMANDS) - if unknown: - raise MatrixError(f"Unknown features to keep: {sorted(unknown)}") path = self.inputs.output / "projects" / f"t-{uuid.uuid4().hex[:8]}" path.parent.mkdir(parents=True, exist_ok=True) - project = Project(authoring, path, f"agent-bricks-{path.name}", features=tuple(keep)) + project = Project(authoring, path, f"agent-bricks-{path.name}") self._projects.append(project) self._run_long( f"init-{path.name}", @@ -261,9 +250,6 @@ def new_project(self, authoring: str, keep: Sequence[str] = ()) -> Project: # Registered before any store exists so cleanup knows the app name if creation fails. self.evidence.register_project(project.app_name, authoring=authoring) if authoring == "cli": - for feature, command in _UNBIND_COMMANDS.items(): - if feature not in keep: - self.cli(*command, "--source", str(path)) self._sync_manifest(project) return project @@ -299,19 +285,14 @@ def write_file(self, project: Project, relative_path: str, content: str) -> None target.write_text(content, encoding="utf-8") def write_manifest(self, project: Project, tool_specs: Sequence[ToolSpec]) -> None: - """Direct authoring: write agent.toml with exactly these tools and the project's features.""" + """Direct authoring: write agent.toml with exactly these tools plus a tracing binding.""" sections = [_DIRECT_HEADER.format(framework=FRAMEWORKS[0])] sections.extend(spec.toml for spec in tool_specs if spec.toml) - if "session_store" in project.features: - sections.append(f'[session_store]\nname = "{project.path.name}-sessions"\n') - if "memory_store" in project.features: - sections.append(f'[memory_store]\nname = "{project.path.name}-memory"\n') - if "tracing" in project.features: - # Follows default_experiment_name's shape so direct authoring carries the same - # tracing binding `agentbricks init` gives CLI authoring. - slug = re.sub(r"[^a-z0-9-]+", "-", project.path.name.lower()).strip("-") or "agent" - experiment = f"/Shared/agentbricks_traces/{slug}-{self.inputs.run_suffix}" - sections.append(f'[tracing]\nexperiment_name = "{experiment}"\n') + # Follows default_experiment_name's shape so direct authoring carries the same tracing + # binding `agentbricks init` gives CLI authoring. + slug = re.sub(r"[^a-z0-9-]+", "-", project.path.name.lower()).strip("-") or "agent" + experiment = f"/Shared/agentbricks_traces/{slug}-{self.inputs.run_suffix}" + sections.append(f'[tracing]\nexperiment_name = "{experiment}"\n') target = project.path / "agent.toml" self.transcript.file_step(target, "direct authoring; no agentbricks tools command") target.write_text("\n".join(sections), encoding="utf-8") diff --git a/integrations/agentbricks/tests/matrices/memory_store/test_memory_store.py b/integrations/agentbricks/tests/matrices/memory_store/test_memory_store.py index bd379002e..ac59fbe5f 100644 --- a/integrations/agentbricks/tests/matrices/memory_store/test_memory_store.py +++ b/integrations/agentbricks/tests/matrices/memory_store/test_memory_store.py @@ -1,4 +1,4 @@ -"""STUB: a single-feature agent with a memory store; the store the CLI declared was created. +"""STUB: an `init` agent (default bindings); the memory store the CLI declared was created. Real case still to write: the agent writes a memory in one invocation and recalls it in another. """ @@ -13,7 +13,7 @@ def test_memory_store_created(agentbricks_cli: AgentbricksCli, evidence: Evidence) -> None: - project = agentbricks_cli.new_project("cli", keep=("memory_store",)) + project = agentbricks_cli.new_project("cli") command = f"agentbricks memory stores create --name {project.memory_store_config}" with recorded_row( evidence, diff --git a/integrations/agentbricks/tests/matrices/session_store/test_session_store.py b/integrations/agentbricks/tests/matrices/session_store/test_session_store.py index fc35e1b4d..7ce71b7b2 100644 --- a/integrations/agentbricks/tests/matrices/session_store/test_session_store.py +++ b/integrations/agentbricks/tests/matrices/session_store/test_session_store.py @@ -1,4 +1,4 @@ -"""STUB: a single-feature agent with a session store; the store the CLI declared was created. +"""STUB: an `init` agent (default bindings); the session store the CLI declared was created. Real case still to write: conversation state persists across invocations of one session_id and is isolated between sessions. @@ -14,7 +14,7 @@ def test_session_store_created(agentbricks_cli: AgentbricksCli, evidence: Evidence) -> None: - project = agentbricks_cli.new_project("cli", keep=("session_store",)) + project = agentbricks_cli.new_project("cli") command = f"agentbricks sessions stores create --name {project.session_store_config}" with recorded_row( evidence, diff --git a/integrations/agentbricks/tests/matrices/tools/test_grants.py b/integrations/agentbricks/tests/matrices/tools/test_grants.py index 5fce3cd95..9cc463b5b 100644 --- a/integrations/agentbricks/tests/matrices/tools/test_grants.py +++ b/integrations/agentbricks/tests/matrices/tools/test_grants.py @@ -17,8 +17,8 @@ def test_grants( subtests, ) -> None: uc_function = next(tool for tool in tools_for(inputs) if tool.name == "uc_function") - # Tracing stays bound so the App carries a resource Agent Bricks does not own. - project = agentbricks_cli.new_project(authoring, keep=("tracing",)) + # `init` and direct authoring both bind tracing, a resource Agent Bricks does not own. + project = agentbricks_cli.new_project(authoring) app = project.app_name if authoring == "cli": bind(agentbricks_cli, project, uc_function) diff --git a/integrations/agentbricks/tests/matrices/tracing/test_tracing.py b/integrations/agentbricks/tests/matrices/tracing/test_tracing.py index 2754087b2..5dbdd25c2 100644 --- a/integrations/agentbricks/tests/matrices/tracing/test_tracing.py +++ b/integrations/agentbricks/tests/matrices/tracing/test_tracing.py @@ -1,4 +1,4 @@ -"""STUB: a single-feature agent with tracing; the experiment is bound and attached to the App. +"""STUB: an `init` agent (default bindings); the tracing experiment is bound and attached to the App. Real case still to write: an invocation exports a trace to the experiment. """ @@ -16,7 +16,7 @@ def test_tracing_experiment_bound( agentbricks_cli: AgentbricksCli, workspace_client: Workspace, evidence: Evidence ) -> None: - project = agentbricks_cli.new_project("cli", keep=("tracing",)) + project = agentbricks_cli.new_project("cli") with recorded_row( evidence, "cli", From 8c8e213f8d01d2be1d3528bb5637114c07260188 Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Thu, 8 Oct 2026 20:25:38 +0000 Subject: [PATCH 03/45] [agentbricks] run e2e matrices in parallel and fix live-run issues Co-authored-by: Isaac --- integrations/agentbricks/pyproject.toml | 1 + .../agentbricks/tests/agentbricks_cli.py | 31 +++++++-- integrations/agentbricks/tests/common.py | 9 ++- .../memory_store/bundle/databricks.yml | 2 +- .../session_store/bundle/databricks.yml | 2 +- .../matrices/tools/bundle/databricks.yml | 2 +- .../tests/matrices/tools/suite.yaml | 4 +- .../tests/matrices/tools/test_dev.py | 4 ++ .../matrices/tracing/bundle/databricks.yml | 2 +- integrations/agentbricks/tests/run_suites.py | 66 ++++++++++++------- .../agentbricks/tests/workspace_client.py | 7 +- 11 files changed, 90 insertions(+), 40 deletions(-) diff --git a/integrations/agentbricks/pyproject.toml b/integrations/agentbricks/pyproject.toml index 58f25c38b..208c9084e 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", ] diff --git a/integrations/agentbricks/tests/agentbricks_cli.py b/integrations/agentbricks/tests/agentbricks_cli.py index be5c31dce..e7b38c1b5 100644 --- a/integrations/agentbricks/tests/agentbricks_cli.py +++ b/integrations/agentbricks/tests/agentbricks_cli.py @@ -68,6 +68,7 @@ class Project: experiment_name: str | None = None # Set only once this project registered the App, so cleanup never deletes one it did not create. app_registered: bool = False + environment_prepared: bool = False class Agent: @@ -124,6 +125,17 @@ def invoke(self, prompt: str, label: str = "invoke") -> str: time.sleep(15) raise MatrixError(f"{label} failed after 3 attempts: {last}") + def warm_up(self, deadline_seconds: float = 180.0) -> None: + """Invoke until the App answers once. A new App's session store can drop the first calls.""" + started = time.monotonic() + while True: + try: + self.invoke("Reply with the single word OK.", label="warm-up") + return + except MatrixError: + if time.monotonic() - started > deadline_seconds: + raise + def curl(self, prompt: str) -> str: """A redacted curl equivalent of ``invoke``, recorded as the evidence row's command.""" auth = " -H 'Authorization: Bearer '" if self.app_name else "" @@ -286,7 +298,7 @@ def write_file(self, project: Project, relative_path: str, content: str) -> None def write_manifest(self, project: Project, tool_specs: Sequence[ToolSpec]) -> None: """Direct authoring: write agent.toml with exactly these tools plus a tracing binding.""" - sections = [_DIRECT_HEADER.format(framework=FRAMEWORKS[0])] + sections: list[str] = [_DIRECT_HEADER.format(framework=FRAMEWORKS[0])] sections.extend(spec.toml for spec in tool_specs if spec.toml) # Follows default_experiment_name's shape so direct authoring carries the same tracing # binding `agentbricks init` gives CLI authoring. @@ -397,9 +409,9 @@ def dev(self, project: Project) -> Iterator[Agent]: label = f"dev-{project.path.name}-{self._next()}" log_path = self.logs_dir / f"{label}.log" port = _free_port() - argv = self.argv( - "dev", "--source", str(project.path), "--app-port", str(port), "--prepare-environment" - ) + # Only the first start prepares the venv; run-local refuses to recreate an existing one. + prepare = [] if project.environment_prepared else ["--prepare-environment"] + argv = self.argv("dev", "--source", str(project.path), "--app-port", str(port), *prepare) log_path.parent.mkdir(parents=True, exist_ok=True) self.transcript.command(argv) with log_path.open("w", encoding="utf-8") as log_file: @@ -409,9 +421,12 @@ def dev(self, project: Project) -> Iterator[Agent]: stdout=log_file, stderr=subprocess.STDOUT, start_new_session=True, + # The template server binds $PORT, so an inherited PORT would override --app-port. + env={**os.environ, "PORT": str(port)}, ) try: self._wait_for_local(process, port, label, log_path) + project.environment_prepared = True yield Agent(f"http://127.0.0.1:{port}", lambda: {}, log_path, self.transcript) finally: if process.poll() is None: @@ -421,7 +436,8 @@ def dev(self, project: Project) -> Iterator[Agent]: def deploy(self, project: Project, app: str | None = None) -> Agent: """`agentbricks deploy` the project; deploying again updates the same App.""" name = app or project.app_name - if not project.app_registered: + creating = not project.app_registered + if creating: self.workspace.assert_app_absent(name) # Deploy can create the App and then fail while waiting for it, so register first. self.evidence.register_app(name) @@ -436,7 +452,7 @@ def deploy(self, project: Project, app: str | None = None) -> Agent: url = str(deployed.get("url") or "").rstrip("/") if not url: raise MatrixError(f"App {name} has no URL: {deployed}") - return Agent( + agent = Agent( url, lambda: self.workspace.app_headers, self.logs_dir / f"{label}.log", @@ -444,6 +460,9 @@ def deploy(self, project: Project, app: str | None = None) -> Agent: app_name=name, app=deployed, ) + if creating: + agent.warm_up() + return agent def _wait_for_local( self, process: subprocess.Popen[str], port: int, label: str, log_path: pathlib.Path diff --git a/integrations/agentbricks/tests/common.py b/integrations/agentbricks/tests/common.py index 20f6dffaa..bc23e24df 100644 --- a/integrations/agentbricks/tests/common.py +++ b/integrations/agentbricks/tests/common.py @@ -173,9 +173,12 @@ class Evidence: """Appends the files under the output directory that the runner turns into ``evidence.json``.""" def __init__(self, output: pathlib.Path): - self.rows_path = output / "rows.jsonl" - self.registry_path = output / "registry.jsonl" - self.cleanup_path = output / "cleanup.jsonl" + # One set of files per pytest-xdist worker; the runner merges them, so no file is shared + # between processes (attach_artifact_to_failed_rows rewrites its rows file in place). + worker = os.environ.get("PYTEST_XDIST_WORKER", "main") + self.rows_path = output / f"rows.{worker}.jsonl" + self.registry_path = output / f"registry.{worker}.jsonl" + self.cleanup_path = output / f"cleanup.{worker}.jsonl" self.grant_dir = output / "grant-checks" self._lock = threading.Lock() diff --git a/integrations/agentbricks/tests/matrices/memory_store/bundle/databricks.yml b/integrations/agentbricks/tests/matrices/memory_store/bundle/databricks.yml index a322534b7..ae61d6a26 100644 --- a/integrations/agentbricks/tests/matrices/memory_store/bundle/databricks.yml +++ b/integrations/agentbricks/tests/matrices/memory_store/bundle/databricks.yml @@ -23,4 +23,4 @@ resources: schemas: scratch_schema: catalog_name: ${var.catalog_name} - name: agentbricks_agent_tools_e2e_${var.run_id} + name: ab_e2e_${var.run_id} diff --git a/integrations/agentbricks/tests/matrices/session_store/bundle/databricks.yml b/integrations/agentbricks/tests/matrices/session_store/bundle/databricks.yml index 9bc6c7d19..ce606e4aa 100644 --- a/integrations/agentbricks/tests/matrices/session_store/bundle/databricks.yml +++ b/integrations/agentbricks/tests/matrices/session_store/bundle/databricks.yml @@ -23,4 +23,4 @@ resources: schemas: scratch_schema: catalog_name: ${var.catalog_name} - name: agentbricks_agent_tools_e2e_${var.run_id} + name: ab_e2e_${var.run_id} diff --git a/integrations/agentbricks/tests/matrices/tools/bundle/databricks.yml b/integrations/agentbricks/tests/matrices/tools/bundle/databricks.yml index 34e7bb212..ca46ec79e 100644 --- a/integrations/agentbricks/tests/matrices/tools/bundle/databricks.yml +++ b/integrations/agentbricks/tests/matrices/tools/bundle/databricks.yml @@ -23,4 +23,4 @@ resources: schemas: scratch_schema: catalog_name: ${var.catalog_name} - name: agentbricks_agent_tools_e2e_${var.run_id} + name: ab_e2e_${var.run_id} diff --git a/integrations/agentbricks/tests/matrices/tools/suite.yaml b/integrations/agentbricks/tests/matrices/tools/suite.yaml index 7449bdc83..ad984fb99 100644 --- a/integrations/agentbricks/tests/matrices/tools/suite.yaml +++ b/integrations/agentbricks/tests/matrices/tools/suite.yaml @@ -4,7 +4,9 @@ name: tools tests: [matrices/tools/] runs: nightly: - timeout_minutes: 240 + timeout_minutes: 60 + # Deploy tests run in parallel; the dev tests share one worker (xdist_group "dev"). + workers: 6 requires: - name: catalog kind: existing diff --git a/integrations/agentbricks/tests/matrices/tools/test_dev.py b/integrations/agentbricks/tests/matrices/tools/test_dev.py index 23b76049e..32ba2d781 100644 --- a/integrations/agentbricks/tests/matrices/tools/test_dev.py +++ b/integrations/agentbricks/tests/matrices/tools/test_dev.py @@ -9,6 +9,7 @@ import contextlib +import pytest from agentbricks_cli import AgentbricksCli from common import Evidence, Inputs from tools import ( @@ -20,6 +21,9 @@ tools_for, ) +# 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, inputs: Inputs, evidence: Evidence, authoring: str, subtests diff --git a/integrations/agentbricks/tests/matrices/tracing/bundle/databricks.yml b/integrations/agentbricks/tests/matrices/tracing/bundle/databricks.yml index 83558e3f0..f4c86aeac 100644 --- a/integrations/agentbricks/tests/matrices/tracing/bundle/databricks.yml +++ b/integrations/agentbricks/tests/matrices/tracing/bundle/databricks.yml @@ -23,4 +23,4 @@ resources: schemas: scratch_schema: catalog_name: ${var.catalog_name} - name: agentbricks_agent_tools_e2e_${var.run_id} + name: ab_e2e_${var.run_id} diff --git a/integrations/agentbricks/tests/run_suites.py b/integrations/agentbricks/tests/run_suites.py index 944b754a4..73591a0b8 100644 --- a/integrations/agentbricks/tests/run_suites.py +++ b/integrations/agentbricks/tests/run_suites.py @@ -40,6 +40,7 @@ AUTHORING_PATHS, FRAMEWORKS, INPUTS_ENV, + MATRIX_CELLS, MatrixError, Transcript, last_nonempty_line, @@ -89,6 +90,7 @@ class Options: app_auth_profile: str | None preprovisioned_app_catalog_access: bool keep_resources: bool + pytest_args: list[str] = dataclasses.field(default_factory=list) class Shell: @@ -307,9 +309,10 @@ def __init__(self, suite: Suite, toolchain: Toolchain, options: Options): self.bundle_env: dict[str, str] = {} self.bundle_target = "nightly" if options.run_profile == "nightly" else "dev" self.inputs_path = self.output / "inputs.json" - self.rows_path = self.output / "rows.jsonl" - self.registry_path = self.output / "registry.jsonl" - self.cleanup_path = self.output / "cleanup.jsonl" + # Per-worker files written by common.Evidence; glob patterns merge them. + self.rows_glob = "rows.*.jsonl" + self.registry_glob = "registry.*.jsonl" + self.cleanup_glob = "cleanup.*.jsonl" self.grant_dir = self.output / "grant-checks" self._workspace: Workspace | None = None @@ -464,6 +467,11 @@ def write_inputs(self) -> None: # Test execution + def _parallel_args(self) -> list[str]: + # loadgroup keeps tests sharing an xdist_group (the local dev runs) on one worker. + workers = int(self.suite.runs[self.options.run_profile].get("workers", 1)) + return ["-n", str(workers), "--dist", "loadgroup"] if workers > 1 else [] + def run_pytest(self, budget_seconds: float) -> int: tests = [str(TESTS_DIR / relative) for relative in self.suite.tests] argv = [ @@ -478,8 +486,13 @@ def run_pytest(self, budget_seconds: float) -> int: "no:cacheprovider", "--junitxml", str(self.output / "junit.xml"), + *self._parallel_args(), + *self.options.pytest_args, ] env = {**os.environ, INPUTS_ENV: str(self.inputs_path)} + if self.options.databricks_profile: + # The CLI auth token lookup is by host; pin the profile when several share one host. + env["DATABRICKS_CONFIG_PROFILE"] = self.options.databricks_profile self.transcript.command(argv) process = subprocess.Popen(argv, cwd=TESTS_DIR, env=env, start_new_session=True) try: @@ -501,25 +514,21 @@ def _registry(self) -> tuple[list[str], dict[str, dict[str, Any]], set[str]]: apps: list[str] = [] projects: dict[str, dict[str, Any]] = {} cleaned: set[str] = set() - if self.registry_path.exists(): - for line in self.registry_path.read_text(encoding="utf-8").splitlines(): - if not line.strip(): - continue - event = json.loads(line) - name = event["app_name"] - if event["kind"] == "app" and name not in apps: - apps.append(name) - elif event["kind"] == "project": - projects.setdefault(name, {}).update(event) - elif event["kind"] == "cleaned": - cleaned.add(name) + for event in self._read_jsonl(self.registry_glob): + name = event["app_name"] + if event["kind"] == "app" and name not in apps: + apps.append(name) + elif event["kind"] == "project": + projects.setdefault(name, {}).update(event) + elif event["kind"] == "cleaned": + cleaned.add(name) return apps, projects, cleaned def teardown(self) -> None: # Each test deletes its own App and stores; this sweeps whatever a test could not, such as # after a crash, and folds the tests' own cleanup results into the evidence. apps, projects, cleaned = self._registry() - self.cleanup_results.extend(self._read_jsonl(self.cleanup_path)) + self.cleanup_results.extend(self._read_jsonl(self.cleanup_glob)) for name in dict.fromkeys([*apps, *projects]): if name not in cleaned: self.cleanup_results.extend( @@ -616,11 +625,10 @@ def teardown(self) -> None: # Evidence - def _read_jsonl(self, path: pathlib.Path) -> list[dict[str, Any]]: - if not path.exists(): - return [] + def _read_jsonl(self, pattern: str) -> list[dict[str, Any]]: return [ json.loads(line) + for path in sorted(self.output.glob(pattern)) for line in path.read_text(encoding="utf-8").splitlines() if line.strip() ] @@ -680,7 +688,7 @@ def write_evidence(self) -> pathlib.Path: "cleanup_required": self.cleanup_required, "cleanup_complete": self.cleanup_complete, "cleanup": self.cleanup_results, - "rows": self._read_jsonl(self.rows_path), + "rows": self._read_jsonl(self.rows_glob), } target = self.output / "evidence.json" temporary = target.with_suffix(".json.tmp") @@ -710,7 +718,8 @@ def execute(self) -> int: self.transcript.write("Resources retained (--keep-resources); skipping teardown.") self.ended_at = now().isoformat() evidence = self.write_evidence() - verified = verify_evidence(evidence) + # A filtered run covers only some cells, so the full-matrix check would always fail. + verified = 0 if self.options.pytest_args else verify_evidence(evidence) return 0 if pytest_returncode == 0 and verified == 0 else 1 @@ -804,6 +813,8 @@ def _verify_tools_evidence(document: dict[str, Any], rows: list[dict[str, Any]]) sys.stdout.write("sandbox evidence: hidden marker provenance is missing or invalid\n") return 1 for tool_kind, marker in sandbox_markers.items(): + if not isinstance(marker, str): + continue # unreachable: validated above matching_rows = [row for row in rows if row.get("tool_kind") == tool_kind] if any( row.get("expected") != marker or marker not in str(row.get("actual", "")) @@ -867,10 +878,10 @@ def verify_evidence(path: pathlib.Path, *, require_cleanup: bool = True) -> int: sys.stdout.write(f"evidence provenance missing: {sorted(missing_provenance)}\n") return 1 suite = document.get("suite") - kinds = MATRIX_KINDS.get(suite) - if kinds is None: + cells = MATRIX_CELLS.get(suite) if isinstance(suite, str) else None + if cells is None: sys.stdout.write( - f"evidence suite {suite!r} is not a known matrix: {sorted(MATRIX_KINDS)}\n" + f"evidence suite {suite!r} is not a known matrix: {sorted(MATRIX_CELLS)}\n" ) return 1 rows = document.get("rows", []) @@ -1007,6 +1018,12 @@ def parse_args() -> argparse.Namespace: help="Skip per-App USE CATALOG grants because catalog access is pre-provisioned.", ) parser.add_argument("--keep-resources", action="store_true") + parser.add_argument( + "--pytest-arg", + action="append", + default=[], + help="Extra pytest argument, e.g. -k test_dev (repeatable). Skips the full-matrix evidence check.", + ) parser.add_argument("--verify-evidence", type=pathlib.Path) args = parser.parse_args() if args.verify_evidence is None and args.profile is None: @@ -1040,6 +1057,7 @@ def main() -> int: app_auth_profile=args.app_auth_profile, preprovisioned_app_catalog_access=args.preprovisioned_app_catalog_access, keep_resources=args.keep_resources, + pytest_args=args.pytest_arg, ) toolchain = Toolchain(options) try: diff --git a/integrations/agentbricks/tests/workspace_client.py b/integrations/agentbricks/tests/workspace_client.py index 1158140e9..5828fe954 100644 --- a/integrations/agentbricks/tests/workspace_client.py +++ b/integrations/agentbricks/tests/workspace_client.py @@ -13,7 +13,7 @@ import subprocess import time from collections.abc import Callable, Sequence -from typing import Any +from typing import Any, cast from common import TOOL_RESOURCE_PREFIX, MatrixError, Transcript, now from databricks.sdk import WorkspaceClient @@ -290,7 +290,10 @@ def _wait_for_app_deleted(self, name: str, timeout: float = 1200) -> None: time.sleep(15) def runtime_store(self, app_name: str) -> dict[str, Any]: - return self.client.api_client.do("GET", f"/api/2.0/agents/runtime-stores/{app_name}") + return cast( + dict[str, Any], + self.client.api_client.do("GET", f"/api/2.0/agents/runtime-stores/{app_name}"), + ) def delete_runtime_store(self, app_name: str) -> None: """Delete the deploy-created Runtime Store; a missing store counts as deleted.""" From f03bcd2d8034fbcd5d9a38a823fa3a6c41d12c73 Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Thu, 8 Oct 2026 21:29:28 +0000 Subject: [PATCH 04/45] [agentbricks] replace e2e suite runner with lazy pytest fixtures and backend seam Co-authored-by: Isaac --- .../agentbricks/tests/agentbricks_cli.py | 265 ++-- integrations/agentbricks/tests/backends.py | 85 ++ integrations/agentbricks/tests/common.py | 345 +----- integrations/agentbricks/tests/conftest.py | 80 +- .../tests/matrices/memory_store/conftest.py | 10 + .../tests/matrices/memory_store/suite.yaml | 18 - .../memory_store/test_memory_store.py | 33 +- .../tests/matrices/session_store/conftest.py | 10 + .../tests/matrices/session_store/suite.yaml | 18 - .../session_store/test_session_store.py | 29 +- .../tests/matrices/tools/conftest.py | 28 +- .../tests/matrices/tools/suite.yaml | 27 - .../tests/matrices/tools/test_deploy.py | 36 +- .../tests/matrices/tools/test_dev.py | 29 +- .../tests/matrices/tools/test_grants.py | 31 +- .../tests/matrices/tools/test_selection.py | 50 +- .../agentbricks/tests/matrices/tools/tools.py | 341 +++--- .../tests/matrices/tracing/conftest.py | 10 + .../tests/matrices/tracing/suite.yaml | 17 - .../tests/matrices/tracing/test_tracing.py | 44 +- .../agentbricks/tests/provisioning.py | 272 +++++ integrations/agentbricks/tests/run_context.py | 152 +++ integrations/agentbricks/tests/run_suites.py | 1079 ----------------- .../agentbricks/tests/shared_state.py | 74 ++ .../agentbricks/tests/workspace_client.py | 139 ++- 25 files changed, 1173 insertions(+), 2049 deletions(-) create mode 100644 integrations/agentbricks/tests/backends.py create mode 100644 integrations/agentbricks/tests/matrices/memory_store/conftest.py delete mode 100644 integrations/agentbricks/tests/matrices/memory_store/suite.yaml create mode 100644 integrations/agentbricks/tests/matrices/session_store/conftest.py delete mode 100644 integrations/agentbricks/tests/matrices/session_store/suite.yaml delete mode 100644 integrations/agentbricks/tests/matrices/tools/suite.yaml create mode 100644 integrations/agentbricks/tests/matrices/tracing/conftest.py delete mode 100644 integrations/agentbricks/tests/matrices/tracing/suite.yaml create mode 100644 integrations/agentbricks/tests/provisioning.py create mode 100644 integrations/agentbricks/tests/run_context.py delete mode 100644 integrations/agentbricks/tests/run_suites.py create mode 100644 integrations/agentbricks/tests/shared_state.py diff --git a/integrations/agentbricks/tests/agentbricks_cli.py b/integrations/agentbricks/tests/agentbricks_cli.py index e7b38c1b5..be58441a3 100644 --- a/integrations/agentbricks/tests/agentbricks_cli.py +++ b/integrations/agentbricks/tests/agentbricks_cli.py @@ -2,7 +2,8 @@ Every call here is a real ``agentbricks`` subprocess, so a regression in the CLI fails the test that makes it. The only non-subprocess part is ``Agent.invoke``, which POSTs to a running agent. -Nothing in this module knows about any particular matrix; tests say what to bind. +Nothing in this module knows about any particular matrix; tests say what to bind. It records the +projects it creates in ``projects`` but deletes nothing: the ``agentbricks_cli`` fixture does. """ from __future__ import annotations @@ -19,6 +20,7 @@ import signal import socket import subprocess +import sys import time import urllib.error import urllib.request @@ -28,18 +30,21 @@ import tomli import tomlkit +from backends import Backend from common import ( + APP_PREFIX, E2E_MODEL, FRAMEWORKS, - Evidence, - Inputs, MatrixError, - Transcript, + RunConfig, + child_env, last_lines, last_nonempty_line, + log, + log_command, now, + project_prefix, ) -from workspace_client import Workspace, cleanup_app _STORE_KINDS = {"memory": ("memory", "stores"), "sessions": ("sessions", "stores")} _DIRECT_HEADER = ( @@ -79,7 +84,6 @@ def __init__( base_url: str, headers: Callable[[], dict[str, str]], log_path: pathlib.Path, - transcript: Transcript, *, app_name: str | None = None, app: dict[str, Any] | None = None, @@ -89,7 +93,6 @@ def __init__( self.app_name = app_name self.app = app self._headers = headers - self._transcript = transcript @property def runtime(self) -> str: @@ -112,15 +115,12 @@ def invoke(self, prompt: str, label: str = "invoke") -> str: for attempt in range(1, 4): try: response = _monitored( - label, - lambda: _http_json(url, body, self._headers()), - self._transcript, - timeout=360, + label, lambda: _http_json(url, body, self._headers()), timeout=360 ) return json.dumps(response, sort_keys=True, default=str) except Exception as exc: last = exc - self._transcript.write(f"attempt {attempt}/3 | {label} | {exc}") + log(f"attempt {attempt}/3 | {label} | {exc}") if attempt < 3: time.sleep(15) raise MatrixError(f"{label} failed after 3 attempts: {last}") @@ -136,50 +136,22 @@ def warm_up(self, deadline_seconds: float = 180.0) -> None: if time.monotonic() - started > deadline_seconds: raise - def curl(self, prompt: str) -> str: - """A redacted curl equivalent of ``invoke``, recorded as the evidence row's command.""" - auth = " -H 'Authorization: Bearer '" if self.app_name else "" - body = json.dumps( - { - "id": "", - "session_id": "", - "input": { - "model": E2E_MODEL, - "messages": [{"role": "user", "content": prompt}], - }, - } - ) - return ( - f"curl -sS -X POST {shlex.quote(f'{self.base_url}/api/invocations')}" - f" -H 'Content-Type: application/json'{auth} --data {shlex.quote(body)}" - ) - class AgentbricksCli: - """Authors, runs and deploys projects with the CLI, and deletes what it created on ``close``.""" - - def __init__( - self, - inputs: Inputs, - workspace: Workspace, - transcript: Transcript, - evidence: Evidence, - ): - self.inputs = inputs - self.workspace = workspace - self.transcript = transcript - self.evidence = evidence - self.logs_dir = inputs.output / "logs" - self._projects: list[Project] = [] + """Authors, runs and deploys projects with the CLI against a backend.""" + + def __init__(self, backend: Backend, run_config: RunConfig): + self.backend = backend + self.run_config = run_config + self.logs_dir = run_config.output / "logs" + self.projects: list[Project] = [] + self._env = child_env(backend.env) self._sequence = 0 # Process plumbing def argv(self, *args: str) -> list[str]: - profile = ( - ["--profile", self.inputs.databricks_profile] if self.inputs.databricks_profile else [] - ) - return [str(self.inputs.agentbricks_bin), *profile, *args] + return [_agentbricks_bin(), *self.backend.cli_args, *args] def cli( self, @@ -189,14 +161,20 @@ def cli( check: bool = True, ) -> subprocess.CompletedProcess[str]: argv = self.argv(*args) - self.transcript.command(argv, cwd) + log_command(argv, cwd) result = subprocess.run( - argv, cwd=cwd, text=True, capture_output=True, timeout=timeout, check=False + argv, + cwd=cwd, + env=self._env, + text=True, + capture_output=True, + timeout=timeout, + check=False, ) if result.stdout.strip(): - self.transcript.write(result.stdout) + log(result.stdout) if result.stderr.strip(): - self.transcript.write(result.stderr) + log(result.stderr) if check and result.returncode != 0: raise MatrixError( f"Command failed ({result.returncode}): {shlex.join(argv)}\n" @@ -204,14 +182,29 @@ def cli( ) return result + def endpoint_invoke(self, agent: Agent, prompt: str) -> str: + """`agentbricks endpoint invoke`, the call the CLI tells users to run after dev/deploy.""" + invocation_id = str(uuid.uuid4()) + body = { + "id": invocation_id, + "session_id": invocation_id, + "input": [{"role": "user", "content": prompt}], + } + target = [agent.app_name] if agent.app_name else ["--url", agent.base_url, "--no-auth"] + result = self.cli( + "endpoint", "invoke", *target, "--path", "/api/invocations", "--json", json.dumps(body) + ) + return result.stdout + def _run_long(self, label: str, argv: Sequence[str], *, timeout: float) -> None: - self.transcript.command(argv) + log_command(argv) log_path = self.logs_dir / f"{label}.log" log_path.parent.mkdir(parents=True, exist_ok=True) started = time.monotonic() with log_path.open("w", encoding="utf-8") as log_file: process = subprocess.Popen( list(argv), + env=self._env, text=True, stdout=log_file, stderr=subprocess.STDOUT, @@ -224,15 +217,13 @@ def _run_long(self, label: str, argv: Sequence[str], *, timeout: float) -> None: terminate_process_group(process, grace=0) raise MatrixError(f"{label} timed out after {timeout:.0f}s; log: {log_path}") if elapsed >= next_tick: - self.transcript.write( - f"tick {now():%H:%M} | {label} | running | {last_nonempty_line(log_path)}" - ) + log(f"tick {now():%H:%M} | {label} | running | {last_nonempty_line(log_path)}") next_tick += 60.0 time.sleep(2) - self.transcript.write(log_path.read_text(encoding="utf-8", errors="replace")) + log(log_path.read_text(encoding="utf-8", errors="replace")) if process.returncode != 0: raise MatrixError(f"{label} failed ({process.returncode}); log: {log_path}") - self.transcript.write(f"tick {now():%H:%M} | {label} | success") + log(f"tick {now():%H:%M} | {label} | success") def _next(self) -> int: self._sequence += 1 @@ -246,21 +237,20 @@ def new_project(self, authoring: str) -> Project: `init` binds a session store, a memory store and a tracing experiment by default, and CLI authoring keeps them. Direct authoring replaces the whole manifest in ``write_manifest``. """ - path = self.inputs.output / "projects" / f"t-{uuid.uuid4().hex[:8]}" + name = f"{project_prefix(self.run_config.run_id)}{uuid.uuid4().hex[:6]}" + path = self.run_config.output / "projects" / name path.parent.mkdir(parents=True, exist_ok=True) - project = Project(authoring, path, f"agent-bricks-{path.name}") - self._projects.append(project) + project = Project(authoring, path, f"{APP_PREFIX}{name}") + self.projects.append(project) self._run_long( f"init-{path.name}", self.argv("init", "--framework", FRAMEWORKS[0], str(path)), timeout=600, ) - if self.inputs.bridge_sha: + if self.run_config.bridge_sha: self._pin_bridge_sources(path) else: self._pin_project_wheel(path) - # Registered before any store exists so cleanup knows the app name if creation fails. - self.evidence.register_project(project.app_name, authoring=authoring) if authoring == "cli": self._sync_manifest(project) return project @@ -293,7 +283,7 @@ def write_file(self, project: Project, relative_path: str, content: str) -> None """Write a user-owned file in the project, as a developer would.""" target = project.path / relative_path target.parent.mkdir(parents=True, exist_ok=True) - self.transcript.file_step(target, "user-owned file") + log(f"# write {target}: user-owned file") target.write_text(content, encoding="utf-8") def write_manifest(self, project: Project, tool_specs: Sequence[ToolSpec]) -> None: @@ -303,10 +293,10 @@ def write_manifest(self, project: Project, tool_specs: Sequence[ToolSpec]) -> No # Follows default_experiment_name's shape so direct authoring carries the same tracing # binding `agentbricks init` gives CLI authoring. slug = re.sub(r"[^a-z0-9-]+", "-", project.path.name.lower()).strip("-") or "agent" - experiment = f"/Shared/agentbricks_traces/{slug}-{self.inputs.run_suffix}" + experiment = f"/Shared/agentbricks_traces/{slug}" sections.append(f'[tracing]\nexperiment_name = "{experiment}"\n') target = project.path / "agent.toml" - self.transcript.file_step(target, "direct authoring; no agentbricks tools command") + log(f"# write {target}: direct authoring; no agentbricks tools command") target.write_text("\n".join(sections), encoding="utf-8") for spec in tool_specs: for relative_path, content in spec.files.items(): @@ -345,15 +335,10 @@ def _sync_manifest(self, project: Project) -> None: ) project.session_resource = json.loads(result.stdout) project.session_store_name = project.session_store_config - self.evidence.register_project( - project.app_name, - authoring=project.authoring, - memory_store_name=project.memory_store_name, - session_store_name=project.session_store_name, - ) def _pin_project_wheel(self, project: pathlib.Path) -> None: - wheel = self.inputs.wheel + wheel = self.run_config.wheel + assert wheel is not None vendor_dir = project / "agentbricks_e2e_wheels" vendor_dir.mkdir() vendored_wheel = vendor_dir / wheel.name @@ -364,14 +349,14 @@ def _pin_project_wheel(self, project: pathlib.Path) -> None: uv = tool.setdefault("uv", {}) sources = uv.setdefault("sources", {}) sources["databricks-agentbricks"] = {"path": vendored_wheel.relative_to(project).as_posix()} - self.transcript.file_step( - pyproject, - f"pin databricks-agentbricks runtime to {vendored_wheel.relative_to(project)}", + log( + f"# write {pyproject}: " + f"pin databricks-agentbricks runtime to {vendored_wheel.relative_to(project)}" ) pyproject.write_text(tomlkit.dumps(document), encoding="utf-8") def _pin_bridge_sources(self, project: pathlib.Path) -> None: - bridge_sha = self.inputs.bridge_sha + bridge_sha = self.run_config.bridge_sha pyproject = project / "pyproject.toml" document = tomlkit.parse(pyproject.read_text(encoding="utf-8")) dependencies = document["project"]["dependencies"] @@ -398,7 +383,7 @@ def _pin_bridge_sources(self, project: pathlib.Path) -> None: } ) document["tool"]["uv"]["sources"][package] = source - self.transcript.file_step(pyproject, f"pin Agent Bricks and LangChain to {bridge_sha}") + log(f"# write {pyproject}: pin Agent Bricks and LangChain to {bridge_sha}") pyproject.write_text(tomlkit.dumps(document), encoding="utf-8") # Running @@ -413,7 +398,7 @@ def dev(self, project: Project) -> Iterator[Agent]: prepare = [] if project.environment_prepared else ["--prepare-environment"] argv = self.argv("dev", "--source", str(project.path), "--app-port", str(port), *prepare) log_path.parent.mkdir(parents=True, exist_ok=True) - self.transcript.command(argv) + log_command(argv) with log_path.open("w", encoding="utf-8") as log_file: process = subprocess.Popen( argv, @@ -422,25 +407,28 @@ def dev(self, project: Project) -> Iterator[Agent]: stderr=subprocess.STDOUT, start_new_session=True, # The template server binds $PORT, so an inherited PORT would override --app-port. - env={**os.environ, "PORT": str(port)}, + env={**self._env, "PORT": str(port)}, ) try: self._wait_for_local(process, port, label, log_path) project.environment_prepared = True - yield Agent(f"http://127.0.0.1:{port}", lambda: {}, log_path, self.transcript) + yield Agent(f"http://127.0.0.1:{port}", lambda: {}, log_path) finally: if process.poll() is None: terminate_process_group(process) - self.transcript.write(f"{label} | stopped") + log(f"{label} | stopped") def deploy(self, project: Project, app: str | None = None) -> Agent: """`agentbricks deploy` the project; deploying again updates the same App.""" + workspace = self.backend.workspace + if workspace is None: + raise MatrixError(f"deploy needs a workspace; the {self.backend.name} backend has none") + workspace.check_app_auth() name = app or project.app_name creating = not project.app_registered if creating: - self.workspace.assert_app_absent(name) - # Deploy can create the App and then fail while waiting for it, so register first. - self.evidence.register_app(name) + workspace.assert_app_absent(name) + # Deploy can create the App and then fail while waiting for it, so mark it first. project.app_registered = True label = f"deploy-{name}-{self._next()}" self._run_long( @@ -448,15 +436,14 @@ def deploy(self, project: Project, app: str | None = None) -> Agent: self.argv("deploy", name, "--source", str(project.path)), timeout=2400, ) - deployed = self.workspace.wait_for_app(name) + deployed = workspace.wait_for_app(name) url = str(deployed.get("url") or "").rstrip("/") if not url: raise MatrixError(f"App {name} has no URL: {deployed}") agent = Agent( url, - lambda: self.workspace.app_headers, + lambda: workspace.app_headers, self.logs_dir / f"{label}.log", - self.transcript, app_name=name, app=deployed, ) @@ -486,84 +473,32 @@ def _wait_for_local( if elapsed > 1200: raise MatrixError(f"{label} did not become reachable: {last_lines(log_path, 30)}") if elapsed >= next_tick: - self.transcript.write( - f"tick {now():%H:%M} | {label} | starting | {last_nonempty_line(log_path)}" - ) + log(f"tick {now():%H:%M} | {label} | starting | {last_nonempty_line(log_path)}") next_tick += 60 time.sleep(5) - # Cleanup + # Stores def delete_store(self, kind: str, name: str) -> subprocess.CompletedProcess[str]: - return delete_store( - self.inputs.agentbricks_bin, self.inputs.databricks_profile, self.transcript, kind, name + """`agentbricks stores delete`.""" + argv = [*self.argv(*_STORE_KINDS[kind]), "delete", name, "--yes"] + log_command(argv) + result = subprocess.run( + argv, env=self._env, text=True, capture_output=True, timeout=600, check=False ) + for text in (result.stdout, result.stderr): + if text.strip(): + log(text) + return result - def close(self) -> None: - """Delete every App, store and role this CLI created; failures are left for the sweep.""" - for project in self._projects: - if project.app_registered: - self._attach_app_logs(project) - info = { - "memory_store_name": project.memory_store_name, - "session_store_name": project.session_store_name, - } - try: - results = cleanup_app( - self.workspace, - self.delete_store, - project.app_name, - info, - has_app=project.app_registered, - ) - except Exception as exc: - self.transcript.write(f"cleanup warning | {project.app_name} | {exc}") - continue - if any(result.get("status") == "failed" for result in results): - self.transcript.write( - f"cleanup incomplete | {project.app_name} | deferred to the runner sweep" - ) - continue - self.evidence.mark_cleaned(project.app_name, results) - - def _attach_app_logs(self, project: Project) -> None: - # The App is deleted next, so a failed cell must capture its logs while it exists. - failed = any( - row["status"] == "fail" and row.get("app_name") == project.app_name - for row in self.evidence.read_rows() - ) - if not failed: - return - log_path = self.workspace.app_logs( - project.app_name, self.logs_dir / f"deploy-runtime-{project.app_name}.log" - ) - if log_path is not None: - self.evidence.attach_artifact_to_failed_rows(project.app_name, log_path) - - -def delete_store( - agentbricks: pathlib.Path | str, - profile: str | None, - transcript: Transcript, - kind: str, - name: str, -) -> subprocess.CompletedProcess[str]: - """`agentbricks stores delete`; also used by the runner's final sweep.""" - command = _STORE_KINDS[kind] - argv = [ - str(agentbricks), - *(["--profile", profile] if profile else []), - *command, - "delete", - name, - "--yes", - ] - transcript.command(argv) - result = subprocess.run(argv, text=True, capture_output=True, timeout=600, check=False) - for text in (result.stdout, result.stderr): - if text.strip(): - transcript.write(text) - return result + +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 MatrixError("The agentbricks CLI is not installed in the test environment.") + return found def terminate_process_group(process: subprocess.Popen[str], *, grace: float = 20) -> None: @@ -589,11 +524,7 @@ def _free_port() -> int: def _monitored( - label: str, - operation: Callable[[], dict[str, Any]], - transcript: Transcript, - *, - timeout: float, + label: str, operation: Callable[[], dict[str, Any]], *, timeout: float ) -> dict[str, Any]: with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor: future = executor.submit(operation) @@ -605,7 +536,7 @@ def _monitored( ) except concurrent.futures.TimeoutError: elapsed = time.monotonic() - started - transcript.write(f"tick {now():%H:%M} | {label} | running | {elapsed:.0f}s") + log(f"tick {now():%H:%M} | {label} | running | {elapsed:.0f}s") if elapsed >= timeout: raise MatrixError(f"{label} timed out after {timeout:.0f}s") from None diff --git a/integrations/agentbricks/tests/backends.py b/integrations/agentbricks/tests/backends.py new file mode 100644 index 000000000..78df6f273 --- /dev/null +++ b/integrations/agentbricks/tests/backends.py @@ -0,0 +1,85 @@ +"""Where the system under test runs: what the CLI subprocess needs, and the workspace tests may touch. + +``live`` is a real Databricks workspace. ``fake`` is an in-process stand-in that has no workspace, so +every requirement that needs one is skipped rather than failed. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from typing import Protocol + +from common import RunConfig +from workspace_client import Workspace + +BACKENDS = ("live", "fake") + + +class BackendUnavailable(Exception): + """The selected backend cannot run in this environment; the suite skips instead of failing.""" + + +class Backend(Protocol): + name: str + + @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 workspace(self) -> Workspace | None: + """SDK access for grants, cleanup and UC setup; None when the backend has no workspace.""" + ... + + +class LiveBackend: + name = "live" + + def __init__(self, profile: str | None, run_config: RunConfig): + self.profile = profile + self._run_config = run_config + self._workspace: 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 workspace(self) -> Workspace: + if self._workspace is None: + self._workspace = Workspace( + self.profile, + app_auth_profile=self._run_config.app_auth_profile, + warehouse_id=self._run_config.warehouse_id, + preprovisioned_app_catalog_access=self._run_config.preprovisioned_app_catalog_access, + ) + return self._workspace + + +class FakeBackend: + """Seam for the in-process fake: DATABRICKS_HOST/DATABRICKS_TOKEN env, no CLI args, no Workspace.""" + + name = "fake" + env: Mapping[str, str] = {} + cli_args: Sequence[str] = () + workspace: Workspace | None = None + + def __init__(self, run_config: RunConfig): + raise BackendUnavailable("fake backend not implemented yet") + + +def make_backend(name: str, profile: str | None, run_config: RunConfig) -> Backend: + if name == "fake": + return FakeBackend(run_config) + return LiveBackend(profile, run_config) diff --git a/integrations/agentbricks/tests/common.py b/integrations/agentbricks/tests/common.py index bc23e24df..241821a07 100644 --- a/integrations/agentbricks/tests/common.py +++ b/integrations/agentbricks/tests/common.py @@ -1,61 +1,25 @@ -"""Matrix-agnostic pieces shared by the test modules and ``run_suites.py``. +"""Constants, ``RunConfig`` and logging/subprocess helpers shared by the matrix tests. -``run_suites.py`` provisions the workspace dependencies and passes their resolved values in a JSON -file named by ``AGENTBRICKS_TEST_INPUTS``. This module loads that file, owns the matrix constants -(so the runner and the tests cannot drift), and records the evidence rows the runner turns into -``evidence.json``. +``run_context.py`` builds one ``RunConfig`` per pytest run; fixtures and tests read it through the +``run_config`` fixture. """ from __future__ import annotations -import contextlib import dataclasses import datetime as dt -import json import os import pathlib import shlex +import subprocess import sys -import threading -import time -from collections.abc import Iterator, Sequence -from typing import Any - -INPUTS_ENV = "AGENTBRICKS_TEST_INPUTS" +from collections.abc import Mapping, Sequence FRAMEWORKS = ("langgraph",) AUTHORING_PATHS = ("cli", "direct") -RUNTIMES = ("dev", "deploy") -TOOL_KINDS = ("sandbox", "mcp", "python", "uc_function", "genie") E2E_MODEL = "system.ai.gpt-5-2" TOOL_RESOURCE_PREFIX = "agentbricks-tool-" - -# The (authoring, runtime, kind) evidence cells each matrix must produce a passing row for. The -# runner's verify step reads this, keyed by suite name. -MATRIX_CELLS: dict[str, tuple[tuple[str, str, str], ...]] = { - "tools": tuple( - (authoring, runtime, kind) - for authoring in AUTHORING_PATHS - for runtime in RUNTIMES - for kind in TOOL_KINDS - ) - + tuple((authoring, "deploy", "selection") for authoring in AUTHORING_PATHS), - "session_store": (("cli", "deploy", "session_store"),), - "memory_store": (("cli", "deploy", "memory_store"),), - "tracing": (("cli", "deploy", "tracing"),), -} - -EXPECTED = { - "sandbox": "the exact hidden marker read from the temporary UC volume file", - "python": "AGENTBRICKS_PYTHON_OK", - "uc_function": "AGENTBRICKS_UC_OK:matrix", - "mcp": "a web-search tool call and a non-empty https result", - "genie": "a genie_ask tool call and a non-empty Genie response", - "selection": "the UC function tool's exact result, and no call to the Python marker tool", - "session_store": "the declared session store exists", - "memory_store": "the declared memory store exists", - "tracing": "the tracing experiment is bound", -} +APP_PREFIX = "agent-bricks-" class MatrixError(RuntimeError): @@ -66,6 +30,21 @@ 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 teardown can sweep the Apps a crashed test leaked.""" + return f"t-{run_id}-" + + +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 "" @@ -79,257 +58,49 @@ def last_nonempty_line(path: pathlib.Path) -> str: return "no output yet" -class Transcript: - def __init__(self, path: pathlib.Path): - self.path = path - self.path.parent.mkdir(parents=True, exist_ok=True) - self._lock = threading.Lock() - - def write(self, text: str) -> None: - line = text.rstrip() + "\n" - with self._lock: - with self.path.open("a", encoding="utf-8") as output: - output.write(line) - sys.stdout.write(line) - sys.stdout.flush() - - def command(self, argv: Sequence[str], cwd: pathlib.Path | None = None) -> None: - prefix = f"cd {shlex.quote(str(cwd))} && " if cwd else "" - self.write(f"$ {prefix}{shlex.join(list(argv))}") - - def file_step(self, path: pathlib.Path, description: str) -> None: - self.write(f"# write {path}: {description}") +def run_command( + argv: Sequence[str], + *, + cwd: pathlib.Path | None = None, + env: Mapping[str, str] | None = None, + timeout: float = 300, +) -> str: + """Run a command, echo its output, and return stdout; a non-zero exit is a ``MatrixError``.""" + log_command(argv, cwd) + result = subprocess.run( + list(argv), + cwd=cwd, + env=None if env is None else dict(env), + text=True, + capture_output=True, + timeout=timeout, + check=False, + ) + for text in (result.stdout, result.stderr): + if text.strip(): + log(text) + if result.returncode != 0: + raise MatrixError( + f"Command failed ({result.returncode}): {shlex.join(argv)}\n" + f"{result.stderr or result.stdout}" + ) + return result.stdout @dataclasses.dataclass(frozen=True) -class Inputs: - """The workspace dependencies ``run_suites.py`` resolved for this suite run.""" +class RunConfig: + """Values shared by every fixture of one pytest run, including all xdist workers.""" - suite: str - databricks_profile: str | None + run_id: str + output: pathlib.Path app_auth_profile: str | None + # Operator override; None lets the workspace pick the first running warehouse. + warehouse_id: str | None preprovisioned_app_catalog_access: bool bridge_sha: str | None - wheel: pathlib.Path - agentbricks_bin: pathlib.Path - output: pathlib.Path - run_suffix: str - catalog: str - scratch_schema: str - uc_function: str - transitive_uc_function: str - uc_volume: str - volume_marker: str - volume_file_path: str - genie_space_id: str - warehouse_id: str - - @classmethod - def from_env(cls) -> Inputs: - path = os.environ.get(INPUTS_ENV) - if not path: - raise MatrixError(f"{INPUTS_ENV} is not set; run these tests through run_suites.py.") - data = json.loads(pathlib.Path(path).read_text(encoding="utf-8")) - return cls( - suite=data["suite"], - databricks_profile=data.get("databricks_profile"), - app_auth_profile=data.get("app_auth_profile") or data.get("databricks_profile"), - preprovisioned_app_catalog_access=bool(data.get("preprovisioned_app_catalog_access")), - bridge_sha=data.get("bridge_sha"), - wheel=pathlib.Path(data["wheel"]), - agentbricks_bin=pathlib.Path(data["agentbricks_bin"]), - output=pathlib.Path(data["output_dir"]), - run_suffix=data["run_suffix"], - catalog=data["catalog"], - scratch_schema=data["schema"], - uc_function=data["uc_function"], - transitive_uc_function=data["transitive_uc_function"], - uc_volume=data["uc_volume"], - volume_marker=data["volume_marker"], - volume_file_path=data["volume_file_path"], - genie_space_id=data["genie_space_id"], - warehouse_id=data["warehouse_id"], - ) - - -@dataclasses.dataclass -class EvidenceRow: - framework: str - authoring: str - runtime: str - tool_kind: str - status: str - command: str - expected: str - actual: str - duration_seconds: float - artifact_paths: list[str] - app_name: str | None = None - app_url: str | None = None - error: str | None = None - - -class Evidence: - """Appends the files under the output directory that the runner turns into ``evidence.json``.""" - - def __init__(self, output: pathlib.Path): - # One set of files per pytest-xdist worker; the runner merges them, so no file is shared - # between processes (attach_artifact_to_failed_rows rewrites its rows file in place). - worker = os.environ.get("PYTEST_XDIST_WORKER", "main") - self.rows_path = output / f"rows.{worker}.jsonl" - self.registry_path = output / f"registry.{worker}.jsonl" - self.cleanup_path = output / f"cleanup.{worker}.jsonl" - self.grant_dir = output / "grant-checks" - self._lock = threading.Lock() - - def register_app(self, name: str) -> None: - self._append(self.registry_path, {"kind": "app", "app_name": name}) - - def register_project(self, app_name: str, **fields: Any) -> None: - self._append(self.registry_path, {"kind": "project", "app_name": app_name, **fields}) - - def mark_cleaned(self, app_name: str, results: Sequence[dict[str, Any]]) -> None: - for result in results: - self._append(self.cleanup_path, result) - self._append(self.registry_path, {"kind": "cleaned", "app_name": app_name}) - - def record_row(self, row: EvidenceRow) -> None: - self._append(self.rows_path, dataclasses.asdict(row)) - - def read_rows(self) -> list[dict[str, Any]]: - if not self.rows_path.exists(): - return [] - with self._lock: - lines = self.rows_path.read_text(encoding="utf-8").splitlines() - return [json.loads(line) for line in lines if line.strip()] - - def record_setup_failure( - self, - authoring: str, - runtime: str, - exc: Exception, - log_path: pathlib.Path | None, - kinds: Sequence[str], - app_name: str | None = None, - ) -> None: - """Fail the cells whose agent could not be set up, so the evidence names them.""" - existing = { - row["tool_kind"] - for row in self.read_rows() - if row["authoring"] == authoring and row["runtime"] == runtime - } - for kind in kinds: - if kind in existing: - continue - self.record_row( - EvidenceRow( - framework=FRAMEWORKS[0], - authoring=authoring, - runtime=runtime, - tool_kind=kind, - status="fail", - command="runtime setup", - expected=EXPECTED[kind], - actual="", - duration_seconds=0.0, - artifact_paths=[str(log_path)] if log_path else [], - app_name=app_name, - error=str(exc), - ) - ) - - def attach_artifact_to_failed_rows(self, app_name: str, artifact: pathlib.Path) -> bool: - rows = self.read_rows() - failed = [ - row for row in rows if row["status"] == "fail" and row.get("app_name") == app_name - ] - for row in failed: - row["artifact_paths"].append(str(artifact)) - if failed: - with self._lock: - temporary = self.rows_path.with_suffix(".jsonl.tmp") - temporary.write_text( - "".join(json.dumps(row) + "\n" for row in rows), encoding="utf-8" - ) - os.replace(temporary, self.rows_path) - return bool(failed) - - def record_grant_check(self, app_name: str, **fields: Any) -> None: - self.grant_dir.mkdir(parents=True, exist_ok=True) - path = self.grant_dir / f"{app_name}.json" - current = json.loads(path.read_text(encoding="utf-8")) if path.exists() else {} - current.update(fields) - current["app_name"] = app_name - path.write_text(json.dumps(current, indent=2), encoding="utf-8") - - def _append(self, path: pathlib.Path, record: dict[str, Any]) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - with self._lock: - with path.open("a", encoding="utf-8") as output: - output.write(json.dumps(record) + "\n") - - -@dataclasses.dataclass -class Outcome: - """What a cell observed; the row's ``actual`` and links are derived from it.""" - - serialized: str = "" - command: str = "" - log_path: pathlib.Path | None = None - app_name: str | None = None - app_url: str | None = None - - -@contextlib.contextmanager -def recorded_row( - evidence: Evidence, - authoring: str, - runtime: str, - kind: str, - *, - command: str, - expected: str, - marker: str | None = None, -) -> Iterator[Outcome]: - """Run one matrix cell and append its pass/fail evidence row whatever happens.""" - outcome = Outcome(command=command) - started = time.monotonic() - error: str | None = None - try: - yield outcome - except Exception as exc: - error = str(exc) or type(exc).__name__ - raise - finally: - evidence.record_row( - EvidenceRow( - framework=FRAMEWORKS[0], - authoring=authoring, - runtime=runtime, - tool_kind=kind, - status="pass" if error is None else "fail", - command=outcome.command, - expected=expected, - actual=evidence_excerpt(outcome.serialized, marker), - duration_seconds=round(time.monotonic() - started, 3), - artifact_paths=[str(outcome.log_path)] if outcome.log_path else [], - app_name=outcome.app_name, - app_url=outcome.app_url, - error=error, - ) - ) + wheel: pathlib.Path | None -def evidence_excerpt(serialized: str, required_marker: str | None) -> str: - limit = 6000 - if len(serialized) <= limit: - return serialized - head_length = limit // 2 - head = serialized[:head_length] - if required_marker is None or required_marker in head: - tail = serialized[-head_length:] - else: - marker_offset = serialized.find(required_marker) - start = max(head_length, marker_offset - head_length // 2) - tail = serialized[start : start + head_length] - return f"{head}\n... response truncated ...\n{tail}" +def child_env(extra: Mapping[str, str] | None = None) -> dict[str, str]: + """The environment for a subprocess: ours plus ``extra``.""" + return {**os.environ, **(extra or {})} diff --git a/integrations/agentbricks/tests/conftest.py b/integrations/agentbricks/tests/conftest.py index fa0d08981..e0109f7a0 100644 --- a/integrations/agentbricks/tests/conftest.py +++ b/integrations/agentbricks/tests/conftest.py @@ -1,7 +1,15 @@ """Fixtures shared by every matrix under ``matrices/``. +``run_context`` (options, backend, run identity) and ``provisioning`` (lazy workspace resources and +their teardown) are plugins. One command runs a matrix, from ``integrations/agentbricks``: + + uv run pytest tests/matrices/tools -n 6 --dist loadgroup \ + --databricks-profile --genie-space-id + +``--backend fake`` runs against no workspace; whatever needs one is skipped. + The CLI wrapper is function-scoped: each test authors, runs and deploys its own projects, and the -wrapper deletes every App and store it created when the test ends. +fixture deletes every App and store they created when the test ends. """ from __future__ import annotations @@ -10,44 +18,58 @@ import pytest from agentbricks_cli import AgentbricksCli -from common import Evidence, Inputs, Transcript -from workspace_client import Workspace - - -@pytest.fixture(scope="session") -def inputs() -> Inputs: - return Inputs.from_env() +from backends import Backend +from common import RunConfig, log +from workspace_client import Workspace, cleanup_app +pytest_plugins = ["run_context", "provisioning"] -@pytest.fixture(scope="session") -def transcript(inputs: Inputs) -> Transcript: - return Transcript(inputs.output / "commands.log") +_failed_tests: set[str] = set() -@pytest.fixture(scope="session") -def evidence(inputs: Inputs) -> Evidence: - return Evidence(inputs.output) +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(scope="session") -def workspace_client(inputs: Inputs, transcript: Transcript) -> Workspace: - workspace = Workspace( - inputs.databricks_profile, - transcript, - app_auth_profile=inputs.app_auth_profile, - warehouse_id=inputs.warehouse_id, - preprovisioned_app_catalog_access=inputs.preprovisioned_app_catalog_access, - ) - workspace.check_app_auth() +def workspace_client(backend: Backend) -> Workspace: + workspace = backend.workspace + if workspace is None: + pytest.skip(f"the {backend.name} backend has no Databricks workspace") return workspace @pytest.fixture def agentbricks_cli( - inputs: Inputs, workspace_client: Workspace, transcript: Transcript, evidence: Evidence + request: pytest.FixtureRequest, backend: Backend, run_config: RunConfig ) -> Iterator[AgentbricksCli]: - cli = AgentbricksCli(inputs, workspace_client, transcript, evidence) - try: - yield cli - finally: - cli.close() + cli = AgentbricksCli(backend, run_config) + yield cli + if cli.projects and backend.workspace is not None: + _delete_projects(cli, backend.workspace, failed=request.node.nodeid in _failed_tests) + + +def _delete_projects(cli: AgentbricksCli, workspace: Workspace, *, failed: bool) -> None: + """Delete every App, store and role the test's projects created; the sweep gets the rest.""" + for project in cli.projects: + if failed and project.app_registered: + # The App is deleted next, so a failing test must capture its logs while it exists. + workspace.app_logs( + project.app_name, cli.logs_dir / f"deploy-runtime-{project.app_name}.log" + ) + try: + failures = cleanup_app( + workspace, + project.app_name, + delete_store=cli.delete_store, + memory_store=project.memory_store_name, + session_store=project.session_store_name, + has_app=project.app_registered, + runtime_store=bool(project.memory_store_name or project.session_store_name), + ) + except Exception as exc: + failures = [str(exc)] + for failure in failures: + log(f"cleanup warning | {project.app_name} | {failure}") diff --git a/integrations/agentbricks/tests/matrices/memory_store/conftest.py b/integrations/agentbricks/tests/matrices/memory_store/conftest.py new file mode 100644 index 000000000..71e05baf6 --- /dev/null +++ b/integrations/agentbricks/tests/matrices/memory_store/conftest.py @@ -0,0 +1,10 @@ +from __future__ import annotations + +import pathlib + +import pytest + + +@pytest.fixture(scope="session") +def scratch_schema(deploy_scratch_schema) -> str: + return deploy_scratch_schema(pathlib.Path(__file__).parent / "bundle") diff --git a/integrations/agentbricks/tests/matrices/memory_store/suite.yaml b/integrations/agentbricks/tests/matrices/memory_store/suite.yaml deleted file mode 100644 index 973a30a98..000000000 --- a/integrations/agentbricks/tests/matrices/memory_store/suite.yaml +++ /dev/null @@ -1,18 +0,0 @@ -# STUB: copy-me manifest for the memory_store matrix. The memory store is created per project by -# the test fixture and deleted by the runner's teardown registry, so there is no workspace -# dependency to declare for it. Add an `existing` requirement here if a real case needs a pre- -# existing resource. -name: memory_store -tests: [matrices/memory_store/] -runs: - nightly: - timeout_minutes: 90 -requires: - - name: catalog - kind: existing - value: ${AGENTBRICKS_E2E_CATALOG:-main} - check: [catalogs, get] - - name: scratch_schema - kind: dab - bundle: bundle - catalog: catalog diff --git a/integrations/agentbricks/tests/matrices/memory_store/test_memory_store.py b/integrations/agentbricks/tests/matrices/memory_store/test_memory_store.py index ac59fbe5f..1a674ec8a 100644 --- a/integrations/agentbricks/tests/matrices/memory_store/test_memory_store.py +++ b/integrations/agentbricks/tests/matrices/memory_store/test_memory_store.py @@ -5,36 +5,21 @@ from __future__ import annotations -import json - import pytest from agentbricks_cli import AgentbricksCli -from common import EXPECTED, Evidence, recorded_row -def test_memory_store_created(agentbricks_cli: AgentbricksCli, evidence: Evidence) -> None: +@pytest.mark.usefixtures("workspace_client") +def test_memory_store_created(agentbricks_cli: AgentbricksCli) -> None: project = agentbricks_cli.new_project("cli") - command = f"agentbricks memory stores create --name {project.memory_store_config}" - with recorded_row( - evidence, - "cli", - "deploy", - "memory_store", - command=command, - expected=EXPECTED["memory_store"], - ) as row: - row.app_name = project.app_name - row.log_path = agentbricks_cli.deploy(project).log_path - resource = project.memory_resource or {} - row.serialized = json.dumps(resource, sort_keys=True) - - assert resource.get("display_name") == project.memory_store_config, resource - assert str(resource.get("name")).startswith("memory-stores/"), resource - assert project.memory_store_name == resource["name"] + agentbricks_cli.deploy(project) + resource = project.memory_resource or {} + + assert resource.get("display_name") == project.memory_store_config, resource + assert str(resource.get("name")).startswith("memory-stores/"), resource + assert project.memory_store_name == resource["name"] @pytest.mark.xfail(reason="STUB: no memory write/recall case written yet", run=False) -def test_memory_recalled_across_invocations( - agentbricks_cli: AgentbricksCli, evidence: Evidence -) -> None: +def test_memory_recalled_across_invocations(agentbricks_cli: AgentbricksCli) -> None: raise NotImplementedError diff --git a/integrations/agentbricks/tests/matrices/session_store/conftest.py b/integrations/agentbricks/tests/matrices/session_store/conftest.py new file mode 100644 index 000000000..71e05baf6 --- /dev/null +++ b/integrations/agentbricks/tests/matrices/session_store/conftest.py @@ -0,0 +1,10 @@ +from __future__ import annotations + +import pathlib + +import pytest + + +@pytest.fixture(scope="session") +def scratch_schema(deploy_scratch_schema) -> str: + return deploy_scratch_schema(pathlib.Path(__file__).parent / "bundle") diff --git a/integrations/agentbricks/tests/matrices/session_store/suite.yaml b/integrations/agentbricks/tests/matrices/session_store/suite.yaml deleted file mode 100644 index c5797d0ad..000000000 --- a/integrations/agentbricks/tests/matrices/session_store/suite.yaml +++ /dev/null @@ -1,18 +0,0 @@ -# STUB: copy-me manifest for the session_store matrix. The session store is created per project by -# the test fixture and deleted by the runner's teardown registry, so there is no workspace -# dependency to declare for it. Add an `existing` requirement here if a real case needs a pre- -# existing resource. -name: session_store -tests: [matrices/session_store/] -runs: - nightly: - timeout_minutes: 90 -requires: - - name: catalog - kind: existing - value: ${AGENTBRICKS_E2E_CATALOG:-main} - check: [catalogs, get] - - name: scratch_schema - kind: dab - bundle: bundle - catalog: catalog diff --git a/integrations/agentbricks/tests/matrices/session_store/test_session_store.py b/integrations/agentbricks/tests/matrices/session_store/test_session_store.py index 7ce71b7b2..2bd4ffa45 100644 --- a/integrations/agentbricks/tests/matrices/session_store/test_session_store.py +++ b/integrations/agentbricks/tests/matrices/session_store/test_session_store.py @@ -6,34 +6,19 @@ from __future__ import annotations -import json - import pytest from agentbricks_cli import AgentbricksCli -from common import EXPECTED, Evidence, recorded_row -def test_session_store_created(agentbricks_cli: AgentbricksCli, evidence: Evidence) -> None: +@pytest.mark.usefixtures("workspace_client") +def test_session_store_created(agentbricks_cli: AgentbricksCli) -> None: project = agentbricks_cli.new_project("cli") - command = f"agentbricks sessions stores create --name {project.session_store_config}" - with recorded_row( - evidence, - "cli", - "deploy", - "session_store", - command=command, - expected=EXPECTED["session_store"], - ) as row: - row.app_name = project.app_name - row.log_path = agentbricks_cli.deploy(project).log_path - resource = project.session_resource or {} - row.serialized = json.dumps(resource, sort_keys=True) - - assert resource.get("session_store_name") == project.session_store_config, resource + agentbricks_cli.deploy(project) + resource = project.session_resource or {} + + assert resource.get("session_store_name") == project.session_store_config, resource @pytest.mark.xfail(reason="STUB: no session-persistence case written yet", run=False) -def test_session_persists_across_invocations( - agentbricks_cli: AgentbricksCli, evidence: Evidence -) -> None: +def test_session_persists_across_invocations(agentbricks_cli: AgentbricksCli) -> None: raise NotImplementedError diff --git a/integrations/agentbricks/tests/matrices/tools/conftest.py b/integrations/agentbricks/tests/matrices/tools/conftest.py index a4de93e51..80e425db1 100644 --- a/integrations/agentbricks/tests/matrices/tools/conftest.py +++ b/integrations/agentbricks/tests/matrices/tools/conftest.py @@ -1,11 +1,35 @@ -"""Tools matrix: every test runs once per authoring path.""" +"""Tools matrix: every test runs once per authoring path, over the tools the backend can support.""" from __future__ import annotations +import pathlib + import pytest -from common import AUTHORING_PATHS +from common import AUTHORING_PATHS, log +from tools import TOOL_DEFINITIONS, Tool @pytest.fixture(params=AUTHORING_PATHS) def authoring(request: pytest.FixtureRequest) -> str: return request.param + + +@pytest.fixture(scope="session") +def scratch_schema(deploy_scratch_schema) -> str: + return deploy_scratch_schema(pathlib.Path(__file__).parent / "bundle") + + +@pytest.fixture(scope="session") +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) diff --git a/integrations/agentbricks/tests/matrices/tools/suite.yaml b/integrations/agentbricks/tests/matrices/tools/suite.yaml deleted file mode 100644 index ad984fb99..000000000 --- a/integrations/agentbricks/tests/matrices/tools/suite.yaml +++ /dev/null @@ -1,27 +0,0 @@ -# Dependency manifest for the LangGraph x CLI/direct x dev/deploy tool matrix. run_suites.py -# provisions these in order; `existing` is only validated, never created or deleted. -name: tools -tests: [matrices/tools/] -runs: - nightly: - timeout_minutes: 60 - # Deploy tests run in parallel; the dev tests share one worker (xdist_group "dev"). - workers: 6 -requires: - - name: catalog - kind: existing - value: ${AGENTBRICKS_E2E_CATALOG:-main} - check: [catalogs, get] - - name: scratch_schema - kind: dab - bundle: bundle - catalog: catalog - # Nested + marker functions and a volume, created imperatively in the schema and swept by its destroy. - - name: uc_function - kind: uc_function - schema: scratch_schema - - name: genie_space - kind: existing - value: ${AGENTBRICKS_E2E_GENIE_SPACE_ID} - check: [genie, get-space] - pattern: "^[0-9a-f]{32}$" diff --git a/integrations/agentbricks/tests/matrices/tools/test_deploy.py b/integrations/agentbricks/tests/matrices/tools/test_deploy.py index ca9ca6a77..0b03d41cc 100644 --- a/integrations/agentbricks/tests/matrices/tools/test_deploy.py +++ b/integrations/agentbricks/tests/matrices/tools/test_deploy.py @@ -9,56 +9,52 @@ from __future__ import annotations from agentbricks_cli import AgentbricksCli -from common import Evidence, Inputs -from tools import bind, invoke_and_check, setup_failures, tool_row, tools_for +from tools import Tool, assert_endpoint_invoke, bind, invoke_and_check from workspace_client import Workspace def test_deploy( agentbricks_cli: AgentbricksCli, workspace_client: Workspace, - inputs: Inputs, - evidence: Evidence, + tools: tuple[Tool, ...], authoring: str, subtests, ) -> None: - tools = tools_for(inputs) project = agentbricks_cli.new_project(authoring) app = project.app_name - def assert_granted(tool) -> None: + def assert_granted(tool: Tool) -> None: if tool.grant is not None: assert workspace_client.granted(app, tool.grant), ( f"{tool.name}: {tool.grant} not granted to {app}: " f"{sorted(workspace_client.granted_tuples(app))}" ) - def grant_transitive_if_needed(tool) -> None: + def grant_transitive_if_needed(tool: Tool) -> None: # The nested function is deliberately outside Agent Bricks' grants; the tool only works # once this manual grant is in place. - if tool.needs_transitive_grant: - workspace_client.grant_transitive(app, inputs.transitive_uc_function) + if tool.transitive_function: + workspace_client.grant_transitive(app, tool.transitive_function) if authoring == "cli": for index, tool in enumerate(tools): - with ( - subtests.test(tool=tool.name), - tool_row(evidence, authoring, "deploy", tool) as row, - ): - row.app_name = app + with subtests.test(tool=tool.name): bind(agentbricks_cli, project, tool) agent = agentbricks_cli.deploy(project, app) grant_transitive_if_needed(tool) - invoke_and_check(row, agent, tool) + invoke_and_check(agent, tool) for added in tools[: index + 1]: assert_granted(added) + with subtests.test(step="endpoint-invoke"): + assert_endpoint_invoke(agentbricks_cli, agent) return - with setup_failures(evidence, authoring, "deploy", tools, app): - agentbricks_cli.write_manifest(project, tools) - agent = agentbricks_cli.deploy(project, app) + agentbricks_cli.write_manifest(project, tools) + agent = agentbricks_cli.deploy(project, app) for tool in tools: - with subtests.test(tool=tool.name), tool_row(evidence, authoring, "deploy", tool) as row: + with subtests.test(tool=tool.name): grant_transitive_if_needed(tool) - invoke_and_check(row, agent, tool) + invoke_and_check(agent, tool) assert_granted(tool) + with subtests.test(step="endpoint-invoke"): + assert_endpoint_invoke(agentbricks_cli, agent) diff --git a/integrations/agentbricks/tests/matrices/tools/test_dev.py b/integrations/agentbricks/tests/matrices/tools/test_dev.py index 32ba2d781..7dd2ec3ae 100644 --- a/integrations/agentbricks/tests/matrices/tools/test_dev.py +++ b/integrations/agentbricks/tests/matrices/tools/test_dev.py @@ -7,18 +7,14 @@ from __future__ import annotations -import contextlib - import pytest from agentbricks_cli import AgentbricksCli -from common import Evidence, Inputs from tools import ( + Tool, + assert_endpoint_invoke, bind, invoke_and_check, reject_unavailable_mcp, - setup_failures, - tool_row, - tools_for, ) # Local dev servers share run-local's proxy port, so all dev runs stay on one worker, in order. @@ -26,25 +22,26 @@ def test_dev( - agentbricks_cli: AgentbricksCli, inputs: Inputs, evidence: Evidence, authoring: str, subtests + agentbricks_cli: AgentbricksCli, tools: tuple[Tool, ...], authoring: str, subtests ) -> None: - tools = tools_for(inputs) project = agentbricks_cli.new_project(authoring) if authoring == "cli": with subtests.test(step="rejects-unavailable-mcp"): reject_unavailable_mcp(agentbricks_cli, project) for tool in tools: - with subtests.test(tool=tool.name), tool_row(evidence, authoring, "dev", tool) as row: + with subtests.test(tool=tool.name): bind(agentbricks_cli, project, tool) with agentbricks_cli.dev(project) as agent: - invoke_and_check(row, agent, tool) + invoke_and_check(agent, tool) + with subtests.test(step="endpoint-invoke"), agentbricks_cli.dev(project) as agent: + assert_endpoint_invoke(agentbricks_cli, agent) return - with contextlib.ExitStack() as stack: - with setup_failures(evidence, authoring, "dev", tools): - agentbricks_cli.write_manifest(project, tools) - agent = stack.enter_context(agentbricks_cli.dev(project)) + agentbricks_cli.write_manifest(project, tools) + with agentbricks_cli.dev(project) as agent: for tool in tools: - with subtests.test(tool=tool.name), tool_row(evidence, authoring, "dev", tool) as row: - invoke_and_check(row, agent, tool) + with subtests.test(tool=tool.name): + invoke_and_check(agent, tool) + with subtests.test(step="endpoint-invoke"): + assert_endpoint_invoke(agentbricks_cli, agent) diff --git a/integrations/agentbricks/tests/matrices/tools/test_grants.py b/integrations/agentbricks/tests/matrices/tools/test_grants.py index 9cc463b5b..8cdcc8dbc 100644 --- a/integrations/agentbricks/tests/matrices/tools/test_grants.py +++ b/integrations/agentbricks/tests/matrices/tools/test_grants.py @@ -3,48 +3,35 @@ from __future__ import annotations from agentbricks_cli import AgentbricksCli -from common import TOOL_RESOURCE_PREFIX, Evidence, Inputs -from tools import bind, tools_for +from common import TOOL_RESOURCE_PREFIX +from provisioning import UcFunction +from tools import bind, uc_function_tool from workspace_client import Workspace def test_grants( agentbricks_cli: AgentbricksCli, workspace_client: Workspace, - inputs: Inputs, - evidence: Evidence, + uc_function: UcFunction, authoring: str, subtests, ) -> None: - uc_function = next(tool for tool in tools_for(inputs) if tool.name == "uc_function") + tool = uc_function_tool(uc_function) # `init` and direct authoring both bind tracing, a resource Agent Bricks does not own. project = agentbricks_cli.new_project(authoring) app = project.app_name if authoring == "cli": - bind(agentbricks_cli, project, uc_function) + bind(agentbricks_cli, project, tool) else: - agentbricks_cli.write_manifest(project, [uc_function]) + agentbricks_cli.write_manifest(project, [tool]) agentbricks_cli.deploy(project, app) - state: dict = {} with subtests.test(check="transitive-function-not-auto-granted"): # Raises if the hidden nested function is already granted or listed as an App resource. - state.update(workspace_client.transitive_state(app, inputs.transitive_uc_function)) + workspace_client.transitive_state(app, uc_function.nested_function) with subtests.test(check="manual-grant-is-direct-and-effective"): - granted = workspace_client.grant_transitive(app, inputs.transitive_uc_function) - if state: - evidence.record_grant_check( - app, - authoring=authoring, - service_principal_client_id=state["service_principal_client_id"], - initial={ - "transitive_direct_privileges": state["direct_privileges"], - "transitive_effective_privileges": state["effective_privileges"], - }, - post_manual_transitive_grant=granted, - manual_transitive_grant_applied=True, - ) + granted = workspace_client.grant_transitive(app, uc_function.nested_function) assert "EXECUTE" in granted["direct_privileges"] assert "EXECUTE" in granted["effective_privileges"] diff --git a/integrations/agentbricks/tests/matrices/tools/test_selection.py b/integrations/agentbricks/tests/matrices/tools/test_selection.py index 16dbd9d3d..01f2f807c 100644 --- a/integrations/agentbricks/tests/matrices/tools/test_selection.py +++ b/integrations/agentbricks/tests/matrices/tools/test_selection.py @@ -7,8 +7,8 @@ from __future__ import annotations from agentbricks_cli import AgentbricksCli -from common import EXPECTED, Evidence, Inputs, recorded_row -from tools import PYTHON_MARKER, UC_MARKER, bind, tools_for +from provisioning import UcFunction +from tools import PYTHON_MARKER, UC_MARKER, Tool, bind from workspace_client import Workspace # Names the Unity Catalog function without its tool id, and the Python helper only to exclude it. @@ -21,39 +21,23 @@ def test_selects_uc_function_not_python_marker( agentbricks_cli: AgentbricksCli, workspace_client: Workspace, - inputs: Inputs, - evidence: Evidence, + tools: tuple[Tool, ...], + uc_function: UcFunction, authoring: str, ) -> None: - tools = tools_for(inputs) project = agentbricks_cli.new_project(authoring) app = project.app_name - with recorded_row( - evidence, - authoring, - "deploy", - "selection", - command="agentbricks deploy", - expected=EXPECTED["selection"], - marker=UC_MARKER, - ) as row: - row.app_name = app - if authoring == "cli": - for tool in tools: - bind(agentbricks_cli, project, tool) - else: - agentbricks_cli.write_manifest(project, tools) - agent = agentbricks_cli.deploy(project, app) - # The nested function is deliberately outside Agent Bricks' grants. - workspace_client.grant_transitive(app, inputs.transitive_uc_function) - row.command = agent.curl(PROMPT) - row.log_path = agent.log_path - row.app_url = agent.base_url - row.serialized = agent.invoke(PROMPT, "deploy-selection") + if authoring == "cli": + for tool in tools: + bind(agentbricks_cli, project, tool) + else: + agentbricks_cli.write_manifest(project, tools) + agent = agentbricks_cli.deploy(project, app) + # The nested function is deliberately outside Agent Bricks' grants. + workspace_client.grant_transitive(app, uc_function.nested_function) + serialized = agent.invoke(PROMPT, "deploy-selection") - assert UC_MARKER in row.serialized, ( - f"The UC function tool was not called: {row.serialized[:2000]}" - ) - assert PYTHON_MARKER not in row.serialized, ( - f"The Python marker tool was called too: {row.serialized[:2000]}" - ) + assert UC_MARKER in serialized, f"The UC function tool was not called: {serialized[:2000]}" + assert PYTHON_MARKER not in serialized, ( + f"The Python marker tool was called too: {serialized[:2000]}" + ) diff --git a/integrations/agentbricks/tests/matrices/tools/tools.py b/integrations/agentbricks/tests/matrices/tools/tools.py index 8bcd49e70..58f5422d9 100644 --- a/integrations/agentbricks/tests/matrices/tools/tools.py +++ b/integrations/agentbricks/tests/matrices/tools/tools.py @@ -2,13 +2,12 @@ from __future__ import annotations -import contextlib import dataclasses import json -from collections.abc import Iterator, Mapping, Sequence +from collections.abc import Callable, Mapping from agentbricks_cli import Agent, AgentbricksCli, Project -from common import EXPECTED, Evidence, Inputs, MatrixError, Outcome, recorded_row +from provisioning import UcFunction from workspace_client import GrantTuple PYTHON_MARKER_FILE = "agent/tools/matrix_marker.py" @@ -23,12 +22,15 @@ PYTHON_MARKER = "AGENTBRICKS_PYTHON_OK" +Check = Callable[[str], None] + + @dataclasses.dataclass(frozen=True) class Tool: name: str prompt: str - # The exact text the response must contain; None for tools checked by their tool-call evidence. - marker: str | None + # Asserts the serialized response proves the tool ran. + check: Check # `agentbricks tools add` arguments; None when the tool is a user-owned file instead. add_args: tuple[str, ...] | None # The [[tools]] entry for direct authoring, and any user-owned files the tool needs. @@ -39,126 +41,155 @@ class Tool: app_auth: bool # The App resource a deploy grants for it; None when it needs no UC grant. grant: GrantTuple | None = None - # The UC function wraps a nested one that Agent Bricks deliberately does not grant. - needs_transitive_grant: bool = False - - def check(self, serialized: str) -> None: - lowered = serialized.lower() - if self.marker is not None: - if self.marker not in serialized: - raise AssertionError(f"Missing exact marker {self.marker!r}: {serialized[:2000]}") - elif self.name == "mcp": - tool_evidence = any(v in lowered for v in ("web_search", "web search", "search")) - if not tool_evidence or "https" not in lowered or len(serialized) < 80: - raise AssertionError( - f"Missing web-search execution/result evidence: {serialized[:2000]}" - ) - elif self.name == "genie": - if ( - not any(v in lowered for v in ("genie_ask", "genie", "conversation_id")) - or len(serialized) < 80 - ): - raise AssertionError( - f"Missing Genie execution/result evidence: {serialized[:2000]}" - ) - else: - raise MatrixError(f"No semantic assertion is defined for tool {self.name!r}.") - - -def tools_for(inputs: Inputs) -> tuple[Tool, ...]: - """The TOOLS table, resolved against this run's scratch volume, function and Genie space.""" - uc_tool_name = inputs.uc_function.replace(".", "__") - return ( - Tool( - name="sandbox", - prompt=( - "You must call the sandbox tool. In the sandbox, use Python to read the " - f"entire text file at {inputs.volume_file_path}. Return only its exact contents; " - "do not fabricate them." - ), - marker=inputs.volume_marker, - add_args=("sandbox", "--scope", f"volume:{inputs.uc_volume}", "--auth", "app"), - toml=( - '[[tools]]\nid = "sandbox"\n' - 'source = { kind = "sandbox", service = "system.ai.sandbox" }\n' - "policy = { downscope = [\n" - f' {{ resource = "volume:{inputs.uc_volume}", permission = "read_only" }},\n' - "] }\n" - ), - files={}, - tool_id="sandbox", - app_auth=True, - grant=("uc_securable", inputs.uc_volume, "VOLUME", "READ_VOLUME"), + # A function the tool calls that Agent Bricks deliberately does not grant; it works only after + # a manual grant. + transitive_function: str | None = None + + +@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 contains(marker: str) -> Check: + def check(serialized: str) -> None: + assert marker in serialized, f"Missing exact marker {marker!r}: {serialized[:2000]}" + + return check + + +def check_web_search(serialized: str) -> None: + lowered = serialized.lower() + tool_evidence = any(v in lowered for v in ("web_search", "web search", "search")) + assert tool_evidence and "https" in lowered and len(serialized) >= 80, ( + f"Missing web-search execution/result evidence: {serialized[:2000]}" + ) + + +def check_genie(serialized: str) -> None: + lowered = serialized.lower() + assert ( + any(v in lowered for v in ("genie_ask", "genie", "conversation_id")) + and len(serialized) >= 80 + ), f"Missing Genie execution/result evidence: {serialized[:2000]}" + + +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." + ), + check=contains(uc_function.volume_marker), + add_args=("sandbox", "--scope", f"volume:{uc_function.volume}", "--auth", "app"), + toml=( + '[[tools]]\nid = "sandbox"\n' + 'source = { kind = "sandbox", service = "system.ai.sandbox" }\n' + "policy = { downscope = [\n" + f' {{ resource = "volume:{uc_function.volume}", permission = "read_only" }},\n' + "] }\n" + ), + files={}, + tool_id="sandbox", + app_auth=True, + grant=("uc_securable", uc_function.volume, "VOLUME", "READ_VOLUME"), + ) + + +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." ), - 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." - ), - marker=None, - add_args=("mcp", "system.ai.web_search", "--auth", "app"), - toml=( - '[[tools]]\nid = "web_search"\n' - 'source = { kind = "mcp", service = "system.ai.web_search" }\n' - ), - files={}, - tool_id="web_search", - app_auth=True, + check=check_web_search, + add_args=("mcp", "system.ai.web_search", "--auth", "app"), + toml=( + '[[tools]]\nid = "web_search"\n' + 'source = { kind = "mcp", service = "system.ai.web_search" }\n' ), - Tool( - name="python", - prompt=( - "You must call the matrix_marker Python tool with value 'matrix'. " - "Return its exact result." - ), - marker=PYTHON_MARKER, - add_args=None, - toml="", - files={PYTHON_MARKER_FILE: PYTHON_MARKER_SOURCE}, - tool_id=None, - app_auth=False, + files={}, + tool_id="web_search", + app_auth=True, + ) + + +def python_tool() -> Tool: + return Tool( + name="python", + prompt=( + "You must call the matrix_marker Python tool with value 'matrix'. " + "Return its exact result." ), - Tool( - name="uc_function", - prompt=( - f"You must call the tool named {uc_tool_name} with value 'matrix'. " - "Return the called tool's exact result." - ), - marker=UC_MARKER, - add_args=("uc-function", inputs.uc_function, "--name", "agentbricks_uc_marker"), - toml=( - '[[tools]]\nid = "agentbricks_uc_marker"\n' - f'source = {{ kind = "uc_function", function = "{inputs.uc_function}" }}\n' - ), - files={}, - tool_id="agentbricks_uc_marker", - app_auth=False, - grant=("uc_securable", inputs.uc_function, "FUNCTION", "EXECUTE"), - needs_transitive_grant=True, + check=contains(PYTHON_MARKER), + add_args=None, + toml="", + files={PYTHON_MARKER_FILE: PYTHON_MARKER_SOURCE}, + tool_id=None, + app_auth=False, + ) + + +def uc_function_tool(uc_function: UcFunction) -> Tool: + return Tool( + name="uc_function", + prompt=( + f"You must call the tool named {uc_function.function.replace('.', '__')} with value " + "'matrix'. Return the called tool's exact result." + ), + check=contains(UC_MARKER), + add_args=("uc-function", uc_function.function, "--name", "agentbricks_uc_marker"), + toml=( + '[[tools]]\nid = "agentbricks_uc_marker"\n' + f'source = {{ kind = "uc_function", function = "{uc_function.function}" }}\n' + ), + files={}, + tool_id="agentbricks_uc_marker", + app_auth=False, + grant=("uc_securable", uc_function.function, "FUNCTION", "EXECUTE"), + transitive_function=uc_function.nested_function, + ) + + +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." ), - 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." - ), - marker=None, - add_args=("genie-agent", inputs.genie_space_id, "--name", "genie", "--auth", "app"), - toml=( - '[[tools]]\nid = "genie"\nauth = "app"\n' - f'source = {{ kind = "genie_agent", space_id = "{inputs.genie_space_id}" }}\n' - ), - files={}, - tool_id="genie", - app_auth=True, - grant=("genie_space", inputs.genie_space_id, "GENIE_SPACE", "CAN_RUN"), + check=check_genie, + add_args=("genie-agent", genie_space, "--name", "genie", "--auth", "app"), + toml=( + '[[tools]]\nid = "genie"\nauth = "app"\n' + f'source = {{ kind = "genie_agent", space_id = "{genie_space}" }}\n' ), + files={}, + tool_id="genie", + app_auth=True, + grant=("genie_space", genie_space, "GENIE_SPACE", "CAN_RUN"), ) +# 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), +) + + def bind(cli: AgentbricksCli, project: Project, tool: Tool) -> None: """CLI authoring: add one tool and confirm agent.toml declares it as intended.""" if tool.add_args is not None: @@ -168,10 +199,13 @@ def bind(cli: AgentbricksCli, project: Project, tool: Tool) -> None: if tool.tool_id is None: return declared = {t["id"]: t for t in cli.manifest(project).get("tools", [])} - if tool.tool_id not in declared: - raise MatrixError(f"tools add did not declare {tool.tool_id!r}: {sorted(declared)}") - if tool.app_auth and declared[tool.tool_id].get("auth") != "app": - raise MatrixError(f"CLI-authored {tool.tool_id} binding is not App-auth.") + assert tool.tool_id in declared, ( + f"tools add did not declare {tool.tool_id!r}: {sorted(declared)}" + ) + if tool.app_auth: + assert declared[tool.tool_id].get("auth") == "app", ( + f"CLI-authored {tool.tool_id} binding is not App-auth." + ) def reject_unavailable_mcp(cli: AgentbricksCli, project: Project) -> None: @@ -187,55 +221,24 @@ def reject_unavailable_mcp(cli: AgentbricksCli, project: Project) -> None: check=False, json_output=True, ) - if rejected.returncode == 0 or manifest.read_bytes() != before: - raise MatrixError( - "agentbricks tools add accepted an unavailable MCP service or changed agent.toml" - ) - if json.loads(rejected.stderr).get("error", {}).get("code") not in { - "NOT_FOUND", - "RESOURCE_DOES_NOT_EXIST", - }: - raise MatrixError(f"Unexpected MCP validation error: {rejected.stderr}") + assert rejected.returncode != 0 and manifest.read_bytes() == before, ( + "agentbricks tools add accepted an unavailable MCP service or changed agent.toml" + ) + code = json.loads(rejected.stderr).get("error", {}).get("code") + assert code in {"NOT_FOUND", "RESOURCE_DOES_NOT_EXIST"}, ( + f"Unexpected MCP validation error: {rejected.stderr}" + ) cli.tools_remove(project, "mcp", "system.ai.missing_service") - if any(tool["id"] == "broken_mcp" for tool in cli.manifest(project).get("tools", [])): - raise MatrixError("agentbricks tools remove left the broken MCP binding in agent.toml") - - -def tool_row(evidence: Evidence, authoring: str, runtime: str, tool: Tool): - """One evidence row for one tool in one authoring x runtime cell.""" - return recorded_row( - evidence, - authoring, - runtime, - tool.name, - command=f"agentbricks {'dev' if runtime == 'dev' else 'deploy'}", - expected=tool.marker or EXPECTED[tool.name], - marker=tool.marker, + assert not any(tool["id"] == "broken_mcp" for tool in cli.manifest(project).get("tools", [])), ( + "agentbricks tools remove left the broken MCP binding in agent.toml" ) -def invoke_and_check(outcome: Outcome, agent: Agent, tool: Tool) -> None: - outcome.command = agent.curl(tool.prompt) - outcome.log_path = agent.log_path - outcome.app_name = agent.app_name - outcome.app_url = agent.base_url if agent.app_name else None - outcome.serialized = agent.invoke(tool.prompt, f"{agent.runtime}-{tool.name}") - tool.check(outcome.serialized) - - -@contextlib.contextmanager -def setup_failures( - evidence: Evidence, - authoring: str, - runtime: str, - tools: Sequence[Tool], - app_name: str | None = None, -) -> Iterator[None]: - """Fail every tool's cell when the shared agent they all need cannot be set up.""" - try: - yield - except Exception as exc: - evidence.record_setup_failure( - authoring, runtime, exc, None, [tool.name for tool in tools], app_name - ) - raise +def invoke_and_check(agent: Agent, tool: Tool) -> None: + tool.check(agent.invoke(tool.prompt, f"{agent.runtime}-{tool.name}")) + + +def assert_endpoint_invoke(agentbricks_cli: AgentbricksCli, agent: Agent) -> None: + """The CLI's own invoke command reaches the running agent and gets a reply.""" + output = agentbricks_cli.endpoint_invoke(agent, "Reply with the single word OK.") + assert "ok" in output.lower(), f"endpoint invoke returned no reply: {output[:500]}" diff --git a/integrations/agentbricks/tests/matrices/tracing/conftest.py b/integrations/agentbricks/tests/matrices/tracing/conftest.py new file mode 100644 index 000000000..71e05baf6 --- /dev/null +++ b/integrations/agentbricks/tests/matrices/tracing/conftest.py @@ -0,0 +1,10 @@ +from __future__ import annotations + +import pathlib + +import pytest + + +@pytest.fixture(scope="session") +def scratch_schema(deploy_scratch_schema) -> str: + return deploy_scratch_schema(pathlib.Path(__file__).parent / "bundle") diff --git a/integrations/agentbricks/tests/matrices/tracing/suite.yaml b/integrations/agentbricks/tests/matrices/tracing/suite.yaml deleted file mode 100644 index d86ed9f2e..000000000 --- a/integrations/agentbricks/tests/matrices/tracing/suite.yaml +++ /dev/null @@ -1,17 +0,0 @@ -# STUB: copy-me manifest for the tracing matrix. The tracing experiment is get-or-created by -# deploy under /Shared/agentbricks_traces, so there is no workspace dependency to declare for it. -# Add an `existing` requirement here if a real case needs a pre-existing resource. -name: tracing -tests: [matrices/tracing/] -runs: - nightly: - timeout_minutes: 90 -requires: - - name: catalog - kind: existing - value: ${AGENTBRICKS_E2E_CATALOG:-main} - check: [catalogs, get] - - name: scratch_schema - kind: dab - bundle: bundle - catalog: catalog diff --git a/integrations/agentbricks/tests/matrices/tracing/test_tracing.py b/integrations/agentbricks/tests/matrices/tracing/test_tracing.py index 5dbdd25c2..54a1b3663 100644 --- a/integrations/agentbricks/tests/matrices/tracing/test_tracing.py +++ b/integrations/agentbricks/tests/matrices/tracing/test_tracing.py @@ -5,43 +5,31 @@ from __future__ import annotations -import json - import pytest from agentbricks_cli import AgentbricksCli -from common import EXPECTED, TOOL_RESOURCE_PREFIX, Evidence, recorded_row +from common import TOOL_RESOURCE_PREFIX from workspace_client import Workspace def test_tracing_experiment_bound( - agentbricks_cli: AgentbricksCli, workspace_client: Workspace, evidence: Evidence + agentbricks_cli: AgentbricksCli, workspace_client: Workspace ) -> None: project = agentbricks_cli.new_project("cli") - with recorded_row( - evidence, - "cli", - "deploy", - "tracing", - command="agentbricks deploy (agent.toml [tracing] experiment_name)", - expected=EXPECTED["tracing"], - ) as row: - row.app_name = project.app_name - row.log_path = agentbricks_cli.deploy(project).log_path - assert project.experiment_name, "agent.toml binds no tracing experiment" - row.serialized = json.dumps({"experiment_name": project.experiment_name}) - - # Deploy attaches the experiment as an App resource that Agent Bricks does not own. - experiments = [ - r for r in workspace_client.non_tool_resources(project.app_name) if "experiment" in r - ] - assert experiments, "No experiment App resource among the non-tool resources" - assert not any(r["name"].startswith(TOOL_RESOURCE_PREFIX) for r in experiments) - expected_id = workspace_client.experiment_id(project.experiment_name) - assert any(str(r["experiment"].get("experiment_id")) == expected_id for r in experiments), ( - f"{project.experiment_name} ({expected_id}) not in {experiments}" - ) + agentbricks_cli.deploy(project) + assert project.experiment_name, "agent.toml binds no tracing experiment" + + # Deploy attaches the experiment as an App resource that Agent Bricks does not own. + experiments = [ + r for r in workspace_client.non_tool_resources(project.app_name) if "experiment" in r + ] + assert experiments, "No experiment App resource among the non-tool resources" + assert not any(r["name"].startswith(TOOL_RESOURCE_PREFIX) for r in experiments) + expected_id = workspace_client.experiment_id(project.experiment_name) + assert any(str(r["experiment"].get("experiment_id")) == expected_id for r in experiments), ( + f"{project.experiment_name} ({expected_id}) not in {experiments}" + ) @pytest.mark.xfail(reason="STUB: no trace-export case written yet", run=False) -def test_invocation_exports_trace(agentbricks_cli: AgentbricksCli, evidence: Evidence) -> None: +def test_invocation_exports_trace(agentbricks_cli: AgentbricksCli) -> None: raise NotImplementedError diff --git a/integrations/agentbricks/tests/provisioning.py b/integrations/agentbricks/tests/provisioning.py new file mode 100644 index 000000000..e6fef7c35 --- /dev/null +++ b/integrations/agentbricks/tests/provisioning.py @@ -0,0 +1,272 @@ +"""pytest plugin: the workspace resources tests ask for, as lazy session fixtures, and their teardown. + +A fixture exists per requirement (``catalog``, ``scratch_schema``, ``uc_function``, +``genie_space``) and is only created when a selected test needs it, so a run that never asks for +``uc_function`` never deploys a bundle or creates a function. Each fixture skips when the backend +cannot provide it or the operator did not configure it. + +Creation happens once per run across xdist workers (``shared_state``). Teardown happens once, in +``pytest_sessionfinish`` on the controller (the only process that sees every worker finish): it +destroys whatever the creators tracked and sweeps Apps a crashed test leaked. +""" + +from __future__ import annotations + +import dataclasses +import json +import re +import uuid +from collections.abc import Callable +from functools import partial +from pathlib import Path + +import pytest +from backends import Backend, BackendUnavailable +from common import ( + APP_PREFIX, + MatrixError, + RunConfig, + child_env, + log, + project_prefix, + run_command, +) +from run_context import base_run_config, build_backend, started_run_config +from shared_state import SharedResources, Track +from workspace_client import Workspace, cleanup_app + +BUNDLE_TARGET = "nightly" +SCHEMA_RESOURCE = "scratch_schema" + + +@dataclasses.dataclass(frozen=True) +class UcFunction: + """The scratch UC objects the tools use.""" + + function: str + # Called by ``function`` and deliberately never declared in agent.toml. + nested_function: str + volume: str + volume_file_path: str + volume_marker: str + + +def _require_workspace(backend: Backend, what: str) -> Workspace: + workspace = backend.workspace + if workspace is None: + pytest.skip(f"{what} needs a Databricks workspace; the {backend.name} backend has none") + return workspace + + +@pytest.fixture(scope="session") +def catalog(request: pytest.FixtureRequest, backend: Backend) -> str: + workspace = _require_workspace(backend, "a UC catalog") + name = request.config.getoption("catalog") + workspace.require_catalog(name) + return name + + +@pytest.fixture(scope="session") +def genie_space(request: pytest.FixtureRequest, backend: Backend) -> str: + space_id = request.config.getoption("genie_space_id") + if not space_id: + pytest.skip("no Genie space; pass --genie-space-id or set AGENTBRICKS_E2E_GENIE_SPACE_ID") + if re.fullmatch(r"[0-9a-f]{32}", space_id) is None: + raise MatrixError(f"Genie space id {space_id!r} is not 32 lowercase hex characters.") + _require_workspace(backend, "a Genie space").require_genie_space(space_id) + return space_id + + +@pytest.fixture(scope="session") +def deploy_scratch_schema( + catalog: str, + backend: Backend, + shared_resources: SharedResources, + request: pytest.FixtureRequest, +) -> Callable[[Path], str]: + """Deploy a matrix's bundle once per run and return its ``catalog.schema``. + + Every bundle names its schema ``ab_e2e_``, so one run can deploy only one of them. + """ + run = base_run_config(request.config) + + def deploy(bundle_dir: Path) -> str: + value = shared_resources.get_or_create( + "scratch_schema", + lambda track: _deploy_bundle(backend, run, catalog, bundle_dir, track), + ) + if value["bundle_dir"] != str(bundle_dir.resolve()): + raise MatrixError( + f"This run already deployed {value['bundle_dir']}, whose schema has the same name " + f"as the one in {bundle_dir}; run one bundle-backed matrix per invocation." + ) + return value["schema"] + + return deploy + + +@pytest.fixture(scope="session") +def uc_function( + scratch_schema: str, backend: Backend, shared_resources: SharedResources +) -> UcFunction: + """Marker functions and a marker volume in ``scratch_schema`` (a fixture each matrix defines).""" + workspace = _require_workspace(backend, "a UC function") + value = shared_resources.get_or_create( + "uc_function", lambda track: _create_uc_function(workspace, scratch_schema, track) + ) + return UcFunction(**value) + + +# Creation + + +def _bundle_env(backend: Backend, catalog: str, run_id: str) -> dict[str, str]: + return child_env( + {**backend.env, "BUNDLE_VAR_catalog_name": catalog, "BUNDLE_VAR_run_id": run_id} + ) + + +def _bundle( + backend: Backend, bundle_dir: Path, env: dict[str, str], *args: str, timeout: float +) -> str: + argv = ["databricks", "bundle", *args, "-t", BUNDLE_TARGET, *backend.cli_args] + return run_command(argv, cwd=bundle_dir, env=env, timeout=timeout) + + +def _deploy_bundle( + backend: Backend, run: RunConfig, catalog: str, bundle_dir: Path, track: Track +) -> dict[str, str]: + bundle_dir = bundle_dir.resolve() + env = _bundle_env(backend, catalog, run.run_id) + # Tracked first: a deploy that fails midway still leaves a schema to destroy. + track(bundle_dir=str(bundle_dir), catalog=catalog, run_id=run.run_id) + _bundle(backend, bundle_dir, env, "deploy", timeout=900) + summary = json.loads( + _bundle(backend, bundle_dir, env, "summary", "--output", "json", timeout=120) + ) + resource = summary.get("resources", {}).get("schemas", {}).get(SCHEMA_RESOURCE) + if not isinstance(resource, dict): + raise MatrixError(f"bundle summary has no schemas.{SCHEMA_RESOURCE}: {summary}") + resolved_id = str(resource.get("id") or "") + if resolved_id.count(".") == 1: + schema = resolved_id + else: + schema = f"{resource.get('catalog_name') or catalog}.{resource.get('name')}" + if schema.count(".") != 1 or "None" in schema: + raise MatrixError(f"Could not resolve the deployed schema from {resource}.") + return {"schema": schema, "bundle_dir": str(bundle_dir)} + + +def _create_uc_function(workspace: Workspace, schema: str, track: Track) -> dict[str, str]: + catalog, _, schema_name = schema.partition(".") + suffix = uuid.uuid4().hex[:8] + nested_name = f"agentbricks_nested_{suffix}" + nested_function = f"{catalog}.{schema_name}.{nested_name}" + track(nested_function=nested_function) + workspace.sql( + f"CREATE OR REPLACE FUNCTION `{catalog}`.`{schema_name}`.`{nested_name}`" + "(value STRING) RETURNS STRING " + "COMMENT 'Transitive Agent Bricks E2E marker; never declared in agent.toml' " + "RETURN concat('AGENTBRICKS_UC_OK:', value)" + ) + # Leave room for catalog and schema in the 64-character MCP tool name. + function_name = f"ab_uc_{suffix}" + function = f"{catalog}.{schema_name}.{function_name}" + volume_name = f"agentbricks_volume_{suffix}" + volume = f"{catalog}.{schema_name}.{volume_name}" + track(volume=volume) + workspace.sql(f"CREATE VOLUME `{catalog}`.`{schema_name}`.`{volume_name}`") + marker = f"AGENTBRICKS_VOLUME_{uuid.uuid4().hex}" + file_path = f"/Volumes/{catalog}/{schema_name}/{volume_name}/marker.txt" + track(volume_file_path=file_path) + workspace.upload(file_path, marker.encode()) + exposed_tool_name = function.replace(".", "__") + if len(exposed_tool_name) > 64: + raise MatrixError( + "The UC function's MCP tool name would exceed 64 characters: " + f"{exposed_tool_name!r}. Use a shorter catalog or run id." + ) + track(function=function) + workspace.sql( + f"CREATE OR REPLACE FUNCTION `{catalog}`.`{schema_name}`.`{function_name}`" + "(value STRING) RETURNS STRING " + "COMMENT 'Deterministic Agent Bricks E2E marker tool' " + f"RETURN `{catalog}`.`{schema_name}`.`{nested_name}`(value)" + ) + return { + "function": function, + "nested_function": nested_function, + "volume": volume, + "volume_file_path": file_path, + "volume_marker": marker, + } + + +# Teardown + + +@pytest.hookimpl(trylast=True) +def pytest_sessionfinish(session: pytest.Session) -> None: + config = session.config + # Workers must not tear down: the controller outlives them and sees the last one finish. + if hasattr(config, "workerinput"): + return + run = started_run_config(config) + if run is None: + return + tracked = SharedResources(run.output / "state").tracked() + if not tracked and not (run.output / "projects").exists(): + return + if config.getoption("keep_resources"): + log("Resources retained (--keep-resources); skipping teardown.") + return + try: + backend = build_backend(config) + except BackendUnavailable: + return + if backend.workspace is not None and _teardown(backend, backend.workspace, run, tracked): + if session.exitstatus == 0: + session.exitstatus = pytest.ExitCode.TESTS_FAILED + + +def _teardown( + backend: Backend, workspace: Workspace, run: RunConfig, tracked: dict[str, dict[str, str]] +) -> list[str]: + """Remove what the run created and return a description of each failure.""" + failures: list[str] = [] + + def attempt(label: str, action: Callable[[], object]) -> None: + try: + action() + except Exception as exc: + log(f"cleanup warning | {label} | {exc}") + failures.append(f"{label}: {exc}") + + attempt("leaked Apps", lambda: failures.extend(_sweep_apps(workspace, run.run_id))) + functions = tracked.get("uc_function", {}) + for key in ("function", "nested_function"): + if name := functions.get(key): + attempt(f"function {name}", partial(workspace.drop, "FUNCTION", name)) + if path := functions.get("volume_file_path"): + attempt(f"file {path}", partial(workspace.delete_file, path)) + if volume := functions.get("volume"): + attempt(f"volume {volume}", partial(workspace.drop, "VOLUME", volume)) + if bundle := tracked.get("scratch_schema"): + env = _bundle_env(backend, bundle["catalog"], bundle["run_id"]) + bundle_dir = Path(bundle["bundle_dir"]) + attempt( + "bundle destroy", + lambda: _bundle(backend, bundle_dir, env, "destroy", "--auto-approve", timeout=900), + ) + return failures + + +def _sweep_apps(workspace: Workspace, run_id: str) -> list[str]: + """Delete Apps whose per-test cleanup failed or never ran, found by the run's name prefix.""" + prefix = f"{APP_PREFIX}{project_prefix(run_id)}" + failures: list[str] = [] + for app in list(workspace.client.apps.list()): + if app.name and app.name.startswith(prefix): + log(f"sweeping leaked App {app.name}") + failures += cleanup_app(workspace, app.name, runtime_store=True) + return failures diff --git a/integrations/agentbricks/tests/run_context.py b/integrations/agentbricks/tests/run_context.py new file mode 100644 index 000000000..a248b2dbc --- /dev/null +++ b/integrations/agentbricks/tests/run_context.py @@ -0,0 +1,152 @@ +"""pytest plugin: command-line options, the run's identity, and the backend the suite runs against. + +Under xdist the controller picks the run id and output directory and hands them to each worker in +``workerinput``, so every process names, finds and later tears down the same resources. +""" + +from __future__ import annotations + +import dataclasses +import os +import pathlib +import re +import tempfile +import uuid +from typing import Any + +import pytest +from backends import BACKENDS, Backend, BackendUnavailable, make_backend +from common import MatrixError, RunConfig, log, run_command +from shared_state import SharedResources + +PACKAGE_DIR = pathlib.Path(__file__).resolve().parent.parent +_RUN_KEY = pytest.StashKey[RunConfig]() +_HANDOFF = "agentbricks_run" + + +def pytest_addoption(parser: pytest.Parser) -> None: + group = parser.getgroup("agentbricks", "Agent Bricks E2E matrices") + group.addoption( + "--backend", + choices=BACKENDS, + default="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( + "--app-auth-profile", + help="OAuth profile for deployed App /api calls; defaults to --databricks-profile.", + ) + group.addoption( + "--catalog", + default=os.environ.get("AGENTBRICKS_E2E_CATALOG", "main"), + help="Existing catalog for scratch schemas (AGENTBRICKS_E2E_CATALOG, default main).", + ) + group.addoption( + "--genie-space-id", + default=os.environ.get("AGENTBRICKS_E2E_GENIE_SPACE_ID"), + help="Genie space for the genie_space fixture (AGENTBRICKS_E2E_GENIE_SPACE_ID).", + ) + group.addoption("--bridge-sha", help="Immutable bridge commit for generated App dependencies.") + 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( + "--preprovisioned-app-catalog-access", + action="store_true", + help="Skip per-App USE CATALOG grants because catalog access is pre-provisioned.", + ) + group.addoption( + "--keep-resources", action="store_true", help="Skip teardown of provisioned resources." + ) + + +def pytest_configure(config: pytest.Config) -> None: + sha = config.getoption("bridge_sha") + if sha is not None and re.fullmatch(r"[0-9a-f]{40}", sha) is None: + raise pytest.UsageError("--bridge-sha must be a 40-character lowercase Git commit SHA.") + + +@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 = uuid.uuid4().hex[:6] + output = pathlib.Path(tempfile.gettempdir()) / f"agentbricks-e2e-{run_id}" + option = config.getoption + run = RunConfig( + run_id=run_id, + output=output, + app_auth_profile=option("app_auth_profile") or option("databricks_profile"), + warehouse_id=option("warehouse_id"), + preprovisioned_app_catalog_access=bool(option("preprovisioned_app_catalog_access")), + bridge_sha=option("bridge_sha"), + wheel=option("wheel"), + ) + config.stash[_RUN_KEY] = run + return run + + +def started_run_config(config: pytest.Config) -> RunConfig | None: + """The run's config if this process ever created one; None means no test needed it.""" + return config.stash.get(_RUN_KEY, None) + + +def build_backend(config: pytest.Config) -> Backend: + return make_backend( + config.getoption("backend"), + config.getoption("databricks_profile"), + base_run_config(config), + ) + + +@pytest.fixture(scope="session") +def backend(request: pytest.FixtureRequest) -> Backend: + try: + return build_backend(request.config) + except BackendUnavailable as exc: + pytest.skip(str(exc)) + + +@pytest.fixture(scope="session") +def shared_resources(request: pytest.FixtureRequest) -> SharedResources: + return SharedResources(base_run_config(request.config).output / "state") + + +@pytest.fixture(scope="session") +def run_config(request: pytest.FixtureRequest, shared_resources: SharedResources) -> 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 or run.bridge_sha: + return run + + def build(_track: object) -> dict[str, Any]: + dist = run.output / "dist" + run_command( + ["uv", "build", "--wheel", "--out-dir", str(dist)], cwd=PACKAGE_DIR, timeout=600 + ) + wheels = sorted(dist.glob("*.whl")) + if not wheels: + raise MatrixError("uv build produced no wheel") + return {"path": str(wheels[-1])} + + built = shared_resources.get_or_create("wheel", build) + return dataclasses.replace(run, wheel=pathlib.Path(built["path"])) diff --git a/integrations/agentbricks/tests/run_suites.py b/integrations/agentbricks/tests/run_suites.py deleted file mode 100644 index 73591a0b8..000000000 --- a/integrations/agentbricks/tests/run_suites.py +++ /dev/null @@ -1,1079 +0,0 @@ -#!/usr/bin/env python3 -"""Provision a suite's workspace dependencies, run its pytest tests, and tear everything down. - -One entry point for local and CI runs. Auth comes from the environment (ambient Databricks -credentials, or DATABRICKS_CONFIG_PROFILE locally): - - uv run python tests/run_suites.py --profile nightly - -Each ``matrices//suite.yaml`` manifest declares the dependencies (``requires``) its pytest directory -needs. Three kinds exist: ``existing`` (validated, never created or deleted), ``dab`` (a bundle -deployed before and destroyed after the tests), and ``uc_function`` (imperative SQL inside a -``dab`` schema, swept by the schema's destroy). The runner hands the resolved values to pytest -through the JSON file named by ``AGENTBRICKS_TEST_INPUTS`` and owns the evidence file. -""" - -from __future__ import annotations - -import argparse -import dataclasses -import functools -import hashlib -import json -import os -import pathlib -import re -import shlex -import signal -import subprocess -import sys -import tempfile -import time -import uuid -import zipfile -from collections.abc import Sequence -from typing import Any - -import yaml -from agentbricks_cli import delete_store -from common import ( - AUTHORING_PATHS, - FRAMEWORKS, - INPUTS_ENV, - MATRIX_CELLS, - MatrixError, - Transcript, - last_nonempty_line, - now, -) -from databricks.sdk import WorkspaceClient -from workspace_client import Workspace, cleanup_app - -TESTS_DIR = pathlib.Path(__file__).resolve().parent -MATRICES_DIR = TESTS_DIR / "matrices" -PACKAGE_DIR = TESTS_DIR.parent -REPOSITORY_ROOT = TESTS_DIR.parents[2] -PROFILES = ("ci", "nightly") -REQUIREMENT_KINDS = ("existing", "dab", "uc_function") -# Time held back from the suite timeout so teardown can finish after pytest is stopped. -TEARDOWN_RESERVE_SECONDS = 180 - -_WHEEL_SOURCE_FILES = ( - "databricks_agentkit/_api_client.py", - "databricks_agentbricks/tool_access.py", - "databricks_agentbricks/app_resources.py", - "databricks_agentbricks/cli/deploy.py", -) -_ENV_REFERENCE = re.compile(r"\$\{([A-Za-z_][A-Za-z0-9_]*)(?::-([^}]*))?\}") - - -@dataclasses.dataclass -class Suite: - name: str - path: pathlib.Path - tests: list[str] - runs: dict[str, dict[str, Any]] - requires: list[dict[str, Any]] - - -@dataclasses.dataclass -class Options: - run_profile: str - output: pathlib.Path - wheel: pathlib.Path | None - commit_sha: str | None - source_root: pathlib.Path - genie_space_id: str | None - bridge_sha: str | None - warehouse_id: str | None - databricks_profile: str | None - app_auth_profile: str | None - preprovisioned_app_catalog_access: bool - keep_resources: bool - pytest_args: list[str] = dataclasses.field(default_factory=list) - - -class Shell: - """Subprocess and Databricks CLI wrappers that log every command.""" - - def __init__(self, profile: str | None, output: pathlib.Path): - self.profile = profile - self.output = output - self.transcript = Transcript(output / "commands.log") - self.warehouse_id: str | None = None - - def profile_args(self) -> list[str]: - return ["--profile", self.profile] if self.profile else [] - - def auth_label(self) -> str: - return f"profile {self.profile!r}" if self.profile else "ambient environment credentials" - - def run( - self, - argv: Sequence[str], - *, - cwd: pathlib.Path | None = None, - env: dict[str, str] | None = None, - timeout: float = 300, - log: bool = True, - check: bool = True, - ) -> subprocess.CompletedProcess[str]: - if log: - self.transcript.command(argv, cwd) - result = subprocess.run( - list(argv), - cwd=cwd, - env=env, - text=True, - capture_output=True, - timeout=timeout, - check=False, - ) - if log and result.stdout.strip(): - self.transcript.write(result.stdout) - if log and result.stderr.strip(): - self.transcript.write(result.stderr) - if check and result.returncode != 0: - raise MatrixError( - f"Command failed ({result.returncode}): {shlex.join(list(argv))}\n" - f"{result.stderr or result.stdout}" - ) - return result - - def run_long( - self, - label: str, - argv: Sequence[str], - *, - cwd: pathlib.Path | None = None, - env: dict[str, str] | None = None, - timeout: float = 1800, - ) -> str: - self.transcript.command(argv, cwd) - log_path = self.output / "logs" / f"{label}.log" - log_path.parent.mkdir(parents=True, exist_ok=True) - started = time.monotonic() - with log_path.open("w", encoding="utf-8") as log_file: - process = subprocess.Popen( - list(argv), - cwd=cwd, - env=env, - text=True, - stdout=log_file, - stderr=subprocess.STDOUT, - start_new_session=True, - ) - next_tick = 60.0 - while process.poll() is None: - elapsed = time.monotonic() - started - if elapsed >= timeout: - os.killpg(process.pid, signal.SIGTERM) - raise MatrixError(f"{label} timed out after {timeout:.0f}s; log: {log_path}") - if elapsed >= next_tick: - self.transcript.write( - f"tick {now():%H:%M} | {label} | running | {last_nonempty_line(log_path)}" - ) - next_tick += 60.0 - time.sleep(2) - output = log_path.read_text(encoding="utf-8", errors="replace") - self.transcript.write(output) - if process.returncode != 0: - raise MatrixError(f"{label} failed ({process.returncode}); log: {log_path}") - self.transcript.write(f"tick {now():%H:%M} | {label} | success") - return output - - def databricks(self, args: Sequence[str], *, timeout: float = 300) -> Any: - result = self.run( - ["databricks", *args, *self.profile_args(), "--output", "json"], - timeout=timeout, - ) - try: - return json.loads(result.stdout or "{}") - except json.JSONDecodeError as exc: - raise MatrixError(f"Databricks CLI returned invalid JSON: {result.stdout}") from exc - - -class Toolchain(Shell): - """The wheel under test, its venv, and the workspace identity shared by every suite.""" - - def __init__(self, options: Options): - super().__init__(options.databricks_profile, options.output) - self.options = options - self.app_auth_profile = options.app_auth_profile or options.databricks_profile - self.wheel = options.wheel.resolve() if options.wheel else None - self.commit_sha = options.commit_sha.lower() if options.commit_sha else None - self.source_provenance: dict[str, Any] = {} - self.versions: dict[str, str] = {} - self.host: str | None = None - self.runner_venv = options.output / "runner-venv" - self.agentbricks = self.runner_venv / "bin" / "agentbricks" - - def _app_auth_label(self) -> str: - return ( - f"App auth profile {self.app_auth_profile!r}" - if self.app_auth_profile - else "ambient environment credentials" - ) - - def bootstrap(self) -> None: - self.output.mkdir(parents=True, exist_ok=True) - source_root = self.options.source_root.resolve() - if self.commit_sha is None: - self.commit_sha = ( - self.run(["git", "rev-parse", "HEAD"], cwd=source_root).stdout.strip().lower() - ) - if re.fullmatch(r"[0-9a-f]{40}", self.commit_sha) is None: - raise MatrixError("--commit-sha must be a full 40-character hexadecimal Git SHA") - if self.options.bridge_sha is not None and ( - re.fullmatch(r"[0-9a-f]{40}", self.options.bridge_sha) is None - ): - raise MatrixError("--bridge-sha must be a 40-character lowercase Git commit SHA.") - if self.wheel is None: - dist = self.output / "dist" - self.run( - ["uv", "build", "--wheel", "--out-dir", str(dist)], cwd=PACKAGE_DIR, timeout=600 - ) - wheels = sorted(dist.glob("*.whl")) - if not wheels: - raise MatrixError("uv build produced no wheel") - self.wheel = wheels[-1].resolve() - self.source_provenance = _source_provenance(source_root, self.commit_sha, self.wheel) - self.run(["uv", "venv", str(self.runner_venv)], timeout=300) - self.run( - [ - "uv", - "pip", - "install", - "--python", - str(self.runner_venv / "bin" / "python"), - str(self.wheel), - ], - timeout=600, - ) - self.run([str(self.agentbricks), "tools", "--help"]) - version_commands = { - "agentbricks": [str(self.agentbricks), "--version"], - "databricks": ["databricks", "version"], - "uv": ["uv", "--version"], - "python": [str(self.runner_venv / "bin" / "python"), "--version"], - } - for name, command in version_commands.items(): - result = self.run(command, log=False, check=False) - version = (result.stdout or result.stderr).strip() - if result.returncode != 0 or not version: - raise MatrixError(f"Could not record {name} version: {version or 'no output'}") - self.versions[name] = version - workspace_client = WorkspaceClient(profile=self.profile) - app_auth_client = WorkspaceClient(profile=self.app_auth_profile) - if not workspace_client.config.host: - raise MatrixError(f"Could not resolve a host from {self.auth_label()}.") - if not app_auth_client.config.host: - raise MatrixError(f"Could not resolve a host from {self._app_auth_label()}.") - self.host = workspace_client.config.host.rstrip("/") - app_auth_host = app_auth_client.config.host.rstrip("/") - if app_auth_host != self.host: - raise MatrixError(f"{self._app_auth_label()} targets {app_auth_host}, not {self.host}.") - if app_auth_client.config.auth_type == "pat": - raise MatrixError( - f"{self._app_auth_label()} uses a PAT. " - "Databricks Apps /api routes require OAuth; run `databricks auth login` " - "for a profile on the same workspace." - ) - if not app_auth_client.config.authenticate().get("Authorization"): - raise MatrixError(f"Could not resolve credentials from {self._app_auth_label()}.") - - -class SuiteRunner(Shell): - def __init__(self, suite: Suite, toolchain: Toolchain, options: Options): - super().__init__(toolchain.profile, options.output / suite.name) - self.suite = suite - self.toolchain = toolchain - self.options = options - self.run_id = uuid.uuid4().hex[:8] - self.run_suffix = self.run_id[:6] - self.cleanup_required = not options.keep_resources - self.cleanup_complete = False - self.cleanup_results: list[dict[str, Any]] = [] - self.started_at = now().isoformat() - self.ended_at: str | None = None - self.resolved: dict[str, str] = {} - self.catalog: str | None = None - self.schema: str | None = None - self.uc_function: str | None = None - self.transitive_uc_function: str | None = None - self.uc_volume: str | None = None - self.volume_marker: str | None = None - self.volume_file_path: str | None = None - self.genie_space_id: str | None = None - self.bundle_dir: pathlib.Path | None = None - self.bundle_env: dict[str, str] = {} - self.bundle_target = "nightly" if options.run_profile == "nightly" else "dev" - self.inputs_path = self.output / "inputs.json" - # Per-worker files written by common.Evidence; glob patterns merge them. - self.rows_glob = "rows.*.jsonl" - self.registry_glob = "registry.*.jsonl" - self.cleanup_glob = "cleanup.*.jsonl" - self.grant_dir = self.output / "grant-checks" - self._workspace: Workspace | None = None - - @property - def workspace(self) -> Workspace: - if self._workspace is None: - self._workspace = Workspace(self.profile, self.transcript) - return self._workspace - - # Provisioning - - def provision(self) -> None: - for requirement in self.suite.requires: - kind = requirement.get("kind") - if kind == "existing": - self._provision_existing(requirement) - elif kind == "dab": - self._provision_dab(requirement) - elif kind == "uc_function": - self._provision_uc_function(requirement) - else: - raise MatrixError(f"Unknown requirement kind {kind!r} in {self.suite.name}.") - self.genie_space_id = self.resolved.get("genie_space") - - def _provision_existing(self, requirement: dict[str, Any]) -> None: - name = requirement["name"] - value = _expand(str(requirement["value"]), self._overrides()) - pattern = requirement.get("pattern") - if pattern and re.fullmatch(pattern, value) is None: - raise MatrixError(f"Requirement {name!r} value {value!r} does not match {pattern!r}.") - check = requirement.get("check") - if check: - self.databricks([*check, value]) - self.resolved[name] = value - if name == "catalog": - self.catalog = value - - def _provision_dab(self, requirement: dict[str, Any]) -> None: - name = requirement["name"] - catalog = self.resolved.get(requirement["catalog"]) - if not catalog: - raise MatrixError(f"Requirement {name!r} needs {requirement['catalog']!r} first.") - self.bundle_dir = (self.suite.path.parent / requirement["bundle"]).resolve() - self.bundle_env = { - **os.environ, - "BUNDLE_VAR_catalog_name": catalog, - "BUNDLE_VAR_run_id": self.run_id, - } - bundle_args = ["-t", self.bundle_target, *self.profile_args()] - self.run_long( - "bundle-deploy", - ["databricks", "bundle", "deploy", *bundle_args], - cwd=self.bundle_dir, - env=self.bundle_env, - timeout=900, - ) - summary = json.loads( - self.run( - ["databricks", "bundle", "summary", *bundle_args, "--output", "json"], - cwd=self.bundle_dir, - env=self.bundle_env, - timeout=120, - ).stdout - ) - resource = summary.get("resources", {}).get("schemas", {}).get(name) - if not isinstance(resource, dict): - raise MatrixError(f"bundle summary has no schemas.{name}: {summary}") - resolved_id = str(resource.get("id") or "") - if resolved_id.count(".") == 1: - schema = resolved_id - else: - schema = f"{resource.get('catalog_name') or catalog}.{resource.get('name')}" - if schema.count(".") != 1 or "None" in schema: - raise MatrixError(f"Could not resolve the deployed schema from {resource}.") - self.schema = schema - self.resolved[name] = schema - - def _provision_uc_function(self, requirement: dict[str, Any]) -> None: - schema = self.resolved.get(requirement["schema"]) - if not schema: - raise MatrixError(f"Requirement {requirement['name']!r} needs a deployed schema first.") - self.warehouse_id = self.workspace.start_warehouse(self.options.warehouse_id) - catalog, _, schema_name = schema.partition(".") - suffix = uuid.uuid4().hex[:8] - nested_function_name = f"agentbricks_nested_{suffix}" - self.transitive_uc_function = f"{catalog}.{schema_name}.{nested_function_name}" - self.workspace.sql( - f"CREATE OR REPLACE FUNCTION `{catalog}`.`{schema_name}`.`{nested_function_name}`" - "(value STRING) RETURNS STRING " - "COMMENT 'Transitive Agent Bricks E2E marker; never declared in agent.toml' " - "RETURN concat('AGENTBRICKS_UC_OK:', value)" - ) - # Leave room for catalog and schema in the 64-character MCP tool name. - function_name = f"ab_uc_{suffix}" - self.uc_function = f"{catalog}.{schema_name}.{function_name}" - volume_name = f"agentbricks_volume_{suffix}" - self.uc_volume = f"{catalog}.{schema_name}.{volume_name}" - self.workspace.sql(f"CREATE VOLUME `{catalog}`.`{schema_name}`.`{volume_name}`") - self.volume_marker = f"AGENTBRICKS_VOLUME_{uuid.uuid4().hex}" - self.volume_file_path = f"/Volumes/{catalog}/{schema_name}/{volume_name}/marker.txt" - self.workspace.upload(self.volume_file_path, self.volume_marker.encode()) - exposed_tool_name = self.uc_function.replace(".", "__") - if len(exposed_tool_name) > 64: - raise MatrixError( - "The UC function's MCP tool name would exceed 64 characters: " - f"{exposed_tool_name!r}. Use a shorter catalog or run id." - ) - self.workspace.sql( - f"CREATE OR REPLACE FUNCTION `{catalog}`.`{schema_name}`.`{function_name}`" - "(value STRING) RETURNS STRING " - "COMMENT 'Deterministic Agent Bricks E2E marker tool' " - f"RETURN `{catalog}`.`{schema_name}`.`{nested_function_name}`(value)" - ) - self.resolved[requirement["name"]] = self.uc_function - - def _overrides(self) -> dict[str, str]: - return ( - {"AGENTBRICKS_E2E_GENIE_SPACE_ID": self.options.genie_space_id} - if self.options.genie_space_id - else {} - ) - - def write_inputs(self) -> None: - toolchain = self.toolchain - inputs = { - "suite": self.suite.name, - "run_profile": self.options.run_profile, - "databricks_profile": self.profile, - "app_auth_profile": toolchain.app_auth_profile, - "preprovisioned_app_catalog_access": self.options.preprovisioned_app_catalog_access, - "bridge_sha": self.options.bridge_sha, - "wheel": str(toolchain.wheel), - "agentbricks_bin": str(toolchain.agentbricks), - "output_dir": str(self.output), - "run_id": self.run_id, - "run_suffix": self.run_suffix, - "host": toolchain.host, - "catalog": self.catalog, - "schema": self.schema, - "uc_function": self.uc_function, - "transitive_uc_function": self.transitive_uc_function, - "uc_volume": self.uc_volume, - "volume_marker": self.volume_marker, - "volume_file_path": self.volume_file_path, - "genie_space_id": self.genie_space_id, - "warehouse_id": self.warehouse_id, - "requirements": self.resolved, - } - self.inputs_path.parent.mkdir(parents=True, exist_ok=True) - self.inputs_path.write_text(json.dumps(inputs, indent=2), encoding="utf-8") - self.inputs_path.chmod(0o600) - - # Test execution - - def _parallel_args(self) -> list[str]: - # loadgroup keeps tests sharing an xdist_group (the local dev runs) on one worker. - workers = int(self.suite.runs[self.options.run_profile].get("workers", 1)) - return ["-n", str(workers), "--dist", "loadgroup"] if workers > 1 else [] - - def run_pytest(self, budget_seconds: float) -> int: - tests = [str(TESTS_DIR / relative) for relative in self.suite.tests] - argv = [ - sys.executable, - "-m", - "pytest", - *tests, - "-v", - "-ra", - "--capture=tee-sys", - "-p", - "no:cacheprovider", - "--junitxml", - str(self.output / "junit.xml"), - *self._parallel_args(), - *self.options.pytest_args, - ] - env = {**os.environ, INPUTS_ENV: str(self.inputs_path)} - if self.options.databricks_profile: - # The CLI auth token lookup is by host; pin the profile when several share one host. - env["DATABRICKS_CONFIG_PROFILE"] = self.options.databricks_profile - self.transcript.command(argv) - process = subprocess.Popen(argv, cwd=TESTS_DIR, env=env, start_new_session=True) - try: - return process.wait(timeout=budget_seconds) - except subprocess.TimeoutExpired: - self.transcript.write(f"pytest exceeded its {budget_seconds:.0f}s budget; stopping it.") - # SIGINT lets pytest unwind its fixtures (the dev servers) before teardown starts. - _signal_group(process, signal.SIGINT) - try: - process.wait(timeout=120) - except subprocess.TimeoutExpired: - _signal_group(process, signal.SIGKILL) - process.wait() - return 124 - - # Teardown - - def _registry(self) -> tuple[list[str], dict[str, dict[str, Any]], set[str]]: - apps: list[str] = [] - projects: dict[str, dict[str, Any]] = {} - cleaned: set[str] = set() - for event in self._read_jsonl(self.registry_glob): - name = event["app_name"] - if event["kind"] == "app" and name not in apps: - apps.append(name) - elif event["kind"] == "project": - projects.setdefault(name, {}).update(event) - elif event["kind"] == "cleaned": - cleaned.add(name) - return apps, projects, cleaned - - def teardown(self) -> None: - # Each test deletes its own App and stores; this sweeps whatever a test could not, such as - # after a crash, and folds the tests' own cleanup results into the evidence. - apps, projects, cleaned = self._registry() - self.cleanup_results.extend(self._read_jsonl(self.cleanup_glob)) - for name in dict.fromkeys([*apps, *projects]): - if name not in cleaned: - self.cleanup_results.extend( - cleanup_app( - self.workspace, - functools.partial( - delete_store, - self.toolchain.agentbricks, - self.profile, - self.transcript, - ), - name, - projects.get(name, {}), - has_app=name in apps, - ) - ) - for label, function in ( - ("UC function", self.uc_function), - ("transitive UC function", self.transitive_uc_function), - ): - if not function: - continue - catalog, schema, function_name = function.split(".") - try: - self.workspace.sql( - f"DROP FUNCTION IF EXISTS `{catalog}`.`{schema}`.`{function_name}`" - ) - self.cleanup_results.append( - {"resource": f"function:{function}", "status": "deleted"} - ) - except Exception as exc: - self.transcript.write(f"cleanup warning | {label} | {exc}") - self.cleanup_results.append( - {"resource": f"function:{function}", "status": "failed", "detail": str(exc)} - ) - if self.uc_volume: - catalog, schema, volume_name = self.uc_volume.split(".") - if self.volume_file_path: - try: - self.workspace.delete_file(self.volume_file_path) - self.cleanup_results.append( - {"resource": f"file:{self.volume_file_path}", "status": "deleted"} - ) - except Exception as exc: - self.transcript.write(f"cleanup warning | UC volume file | {exc}") - self.cleanup_results.append( - { - "resource": f"file:{self.volume_file_path}", - "status": "failed", - "detail": str(exc), - } - ) - try: - self.workspace.sql(f"DROP VOLUME IF EXISTS `{catalog}`.`{schema}`.`{volume_name}`") - self.cleanup_results.append( - {"resource": f"volume:{self.uc_volume}", "status": "deleted"} - ) - except Exception as exc: - self.transcript.write(f"cleanup warning | UC volume | {exc}") - self.cleanup_results.append( - { - "resource": f"volume:{self.uc_volume}", - "status": "failed", - "detail": str(exc), - } - ) - if self.bundle_dir is not None: - resource_name = f"bundle:{self.schema or self.suite.name}" - try: - self.run_long( - "bundle-destroy", - [ - "databricks", - "bundle", - "destroy", - "--auto-approve", - "-t", - self.bundle_target, - *self.profile_args(), - ], - cwd=self.bundle_dir, - env=self.bundle_env, - timeout=900, - ) - self.cleanup_results.append({"resource": resource_name, "status": "deleted"}) - except Exception as exc: - self.transcript.write(f"cleanup warning | bundle destroy | {exc}") - self.cleanup_results.append( - {"resource": resource_name, "status": "failed", "detail": str(exc)} - ) - self.cleanup_complete = not any( - result.get("status") == "failed" for result in self.cleanup_results - ) - - # Evidence - - def _read_jsonl(self, pattern: str) -> list[dict[str, Any]]: - return [ - json.loads(line) - for path in sorted(self.output.glob(pattern)) - for line in path.read_text(encoding="utf-8").splitlines() - if line.strip() - ] - - def _grant_checks(self) -> list[dict[str, Any]]: - checks = [] - for path in sorted(self.grant_dir.glob("*.json")) if self.grant_dir.exists() else (): - recorded = json.loads(path.read_text(encoding="utf-8")) - checks.append( - { - "app_name": recorded.get("app_name"), - "authoring": recorded.get("authoring"), - "service_principal_client_id": recorded.get("service_principal_client_id"), - "initial": recorded.get("initial", {}), - "post_manual_transitive_grant": recorded.get( - "post_manual_transitive_grant", {} - ), - "manual_transitive_grant_applied": recorded.get( - "manual_transitive_grant_applied", False - ), - } - ) - return checks - - def write_evidence(self) -> pathlib.Path: - toolchain = self.toolchain - wheel = toolchain.wheel - payload = { - "schema_version": 1, - "suite": self.suite.name, - "run_profile": self.options.run_profile, - "run_id": self.run_id, - "commit_sha": toolchain.commit_sha, - "source_provenance": toolchain.source_provenance, - "template_repo": "wheel://databricks-agentbricks", - "template_ref": _sha256(wheel) if wheel else None, - "workspace_host": toolchain.host, - "started_at": self.started_at, - "ended_at": self.ended_at, - "versions": toolchain.versions, - "profile": self.profile, - "app_auth_profile": toolchain.app_auth_profile, - "bridge_sha": self.options.bridge_sha, - "preprovisioned_app_catalog_access": self.options.preprovisioned_app_catalog_access, - "generated_at": now().isoformat(), - "wheel": str(wheel) if wheel else None, - "wheel_sha256": _sha256(wheel) if wheel else None, - "schema": self.schema, - "uc_function": self.uc_function, - "transitive_uc_function": self.transitive_uc_function, - "uc_volume": self.uc_volume, - "volume_marker": self.volume_marker, - "volume_file_path": self.volume_file_path, - "genie_space_id": self.genie_space_id, - "warehouse_id": self.warehouse_id, - "grant_checks": self._grant_checks(), - "cleanup_required": self.cleanup_required, - "cleanup_complete": self.cleanup_complete, - "cleanup": self.cleanup_results, - "rows": self._read_jsonl(self.rows_glob), - } - target = self.output / "evidence.json" - temporary = target.with_suffix(".json.tmp") - temporary.write_text(json.dumps(payload, indent=2), encoding="utf-8") - os.replace(temporary, target) - return target - - # Orchestration - - def execute(self) -> int: - timeout_minutes = self.suite.runs[self.options.run_profile]["timeout_minutes"] - started = time.monotonic() - pytest_returncode: int | None = None - try: - self.provision() - self.write_inputs() - remaining = ( - timeout_minutes * 60 - TEARDOWN_RESERVE_SECONDS - (time.monotonic() - started) - ) - pytest_returncode = self.run_pytest(max(60.0, remaining)) - except Exception as exc: - self.transcript.write(f"suite {self.suite.name} failed before tests finished: {exc}") - finally: - if self.cleanup_required: - self.teardown() - else: - self.transcript.write("Resources retained (--keep-resources); skipping teardown.") - self.ended_at = now().isoformat() - evidence = self.write_evidence() - # A filtered run covers only some cells, so the full-matrix check would always fail. - verified = 0 if self.options.pytest_args else verify_evidence(evidence) - return 0 if pytest_returncode == 0 and verified == 0 else 1 - - -def _signal_group(process: subprocess.Popen[Any], sig: signal.Signals) -> None: - try: - os.killpg(process.pid, sig) - except ProcessLookupError: - pass - - -def _expand(value: str, overrides: dict[str, str]) -> str: - def replace(match: re.Match[str]) -> str: - name, default = match.group(1), match.group(2) - resolved = overrides.get(name) or os.environ.get(name) or default - if not resolved: - raise MatrixError(f"Required environment variable {name} is not set.") - return resolved - - return _ENV_REFERENCE.sub(replace, value) - - -def _sha256(path: pathlib.Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as input_file: - for chunk in iter(lambda: input_file.read(1024 * 1024), b""): - digest.update(chunk) - return digest.hexdigest() - - -def _source_provenance( - source_root: pathlib.Path, commit_sha: str, wheel: pathlib.Path -) -> dict[str, Any]: - def git(*args: str, text: bool = True) -> subprocess.CompletedProcess: - result = subprocess.run( - ["git", *args], - cwd=source_root, - capture_output=True, - text=text, - check=False, - ) - if result.returncode != 0: - detail = result.stderr if text else result.stderr.decode(errors="replace") - raise MatrixError(f"Could not inspect source provenance: {detail.strip()}") - return result - - head = git("rev-parse", "HEAD").stdout.strip().lower() - if head != commit_sha.lower(): - raise MatrixError( - f"Claimed commit {commit_sha} does not match source checkout HEAD {head}." - ) - status = git("status", "--porcelain", "--untracked-files=no").stdout.strip() - source_hashes: dict[str, str] = {} - try: - with zipfile.ZipFile(wheel) as archive: - for member in _WHEEL_SOURCE_FILES: - source_path = source_root / "integrations" / "agentbricks" / "src" / member - source_bytes = source_path.read_bytes() - try: - wheel_bytes = archive.read(member) - except KeyError as exc: - raise MatrixError(f"Built wheel is missing source module {member}.") from exc - if wheel_bytes != source_bytes: - raise MatrixError( - f"Built wheel module {member} does not match source checkout." - ) - source_hashes[member] = hashlib.sha256(wheel_bytes).hexdigest() - except zipfile.BadZipFile as exc: - raise MatrixError(f"Built wheel is not a readable zip archive: {wheel}") from exc - diff = git("diff", "--binary", "HEAD", text=False).stdout - return { - "source_head_sha": head, - "source_dirty": bool(status), - "source_diff_sha256": hashlib.sha256(diff).hexdigest(), - "wheel_source_matches": True, - "wheel_source_sha256": source_hashes, - } - - -def _verify_tools_evidence(document: dict[str, Any], rows: list[dict[str, Any]]) -> int: - """Sandbox-marker and grant proofs, which only the tools matrix produces.""" - grant_checks = document.get("grant_checks", []) - sandbox_markers = { - "sandbox": document.get("volume_marker"), - } - volume_file_path = document.get("volume_file_path") - if ( - any(not isinstance(marker, str) or not marker for marker in sandbox_markers.values()) - or not isinstance(volume_file_path, str) - or not volume_file_path.startswith("/Volumes/") - ): - sys.stdout.write("sandbox evidence: hidden marker provenance is missing or invalid\n") - return 1 - for tool_kind, marker in sandbox_markers.items(): - if not isinstance(marker, str): - continue # unreachable: validated above - matching_rows = [row for row in rows if row.get("tool_kind") == tool_kind] - if any( - row.get("expected") != marker or marker not in str(row.get("actual", "")) - for row in matching_rows - ): - sys.stdout.write( - f"sandbox evidence: exact hidden marker proof failed for {tool_kind}\n" - ) - return 1 - # One deploy per authoring path proves the transitive exclusion and the manual grant. Repeat-deploy - # idempotency is proven by the incremental CLI deploy rows, which re-check earlier tools' grants. - transitive = [check for check in grant_checks if check.get("manual_transitive_grant_applied")] - expected_transitive = len(FRAMEWORKS) * len(AUTHORING_PATHS) - if len(transitive) != expected_transitive or any( - "EXECUTE" in check.get("initial", {}).get("transitive_direct_privileges", []) - or "EXECUTE" in check.get("initial", {}).get("transitive_effective_privileges", []) - or "EXECUTE" - not in check.get("post_manual_transitive_grant", {}).get("direct_privileges", []) - or "EXECUTE" - not in check.get("post_manual_transitive_grant", {}).get("effective_privileges", []) - for check in transitive - ): - sys.stdout.write( - f"grant evidence: expected {expected_transitive} transitive exclusion/manual-grant " - f"proofs, got {len(transitive)} (or one failed)\n" - ) - return 1 - return 0 - - -def verify_evidence(path: pathlib.Path, *, require_cleanup: bool = True) -> int: - document = json.loads(path.read_text(encoding="utf-8")) - provenance_fields = ("commit_sha", "workspace_host", "started_at", "ended_at") - missing_provenance = [field for field in provenance_fields if not document.get(field)] - template_repo = document.get("template_repo") - template_ref = document.get("template_ref") - if ( - not isinstance(template_repo, str) - or not template_repo - or not isinstance(template_ref, str) - or not template_ref - ): - missing_provenance.append("template_repo/template_ref") - elif template_repo.startswith("wheel://") and template_ref != document.get("wheel_sha256"): - missing_provenance.append("wheel_template_ref") - source_provenance = document.get("source_provenance") - if ( - not isinstance(source_provenance, dict) - or source_provenance.get("source_head_sha") != document.get("commit_sha") - or source_provenance.get("wheel_source_matches") is not True - or not isinstance(source_provenance.get("wheel_source_sha256"), dict) - or not source_provenance.get("wheel_source_sha256") - ): - missing_provenance.append("source_provenance") - versions = document.get("versions") - if not isinstance(versions, dict) or any( - not versions.get(name) for name in ("agentbricks", "databricks", "uv", "python") - ): - missing_provenance.append("versions") - if missing_provenance: - sys.stdout.write(f"evidence provenance missing: {sorted(missing_provenance)}\n") - return 1 - suite = document.get("suite") - cells = MATRIX_CELLS.get(suite) if isinstance(suite, str) else None - if cells is None: - sys.stdout.write( - f"evidence suite {suite!r} is not a known matrix: {sorted(MATRIX_CELLS)}\n" - ) - return 1 - rows = document.get("rows", []) - expected = { - (framework, authoring, runtime, tool) - for framework in FRAMEWORKS - for authoring, runtime, tool in cells - } - actual = { - (row["framework"], row["authoring"], row["runtime"], row["tool_kind"]) for row in rows - } - duplicates = len(rows) - len(actual) - passed = sum(row.get("status") == "pass" for row in rows) - failed = sum(row.get("status") == "fail" for row in rows) - skipped = len(expected - actual) - sys.stdout.write(f"{passed} passed, {failed} failed, {skipped} skipped\n") - if actual != expected or duplicates or passed != len(expected): - if expected - actual: - sys.stdout.write(f"missing cells: {sorted(expected - actual)}\n") - if duplicates: - sys.stdout.write(f"duplicate rows: {duplicates}\n") - return 1 - if suite == "tools" and _verify_tools_evidence(document, rows): - return 1 - cleanup_required = document.get("cleanup_required") - if not isinstance(cleanup_required, bool): - sys.stdout.write("cleanup evidence: cleanup_required was not recorded\n") - return 1 - cleanup = document.get("cleanup", []) - if ( - require_cleanup - and cleanup_required - and ( - document.get("cleanup_complete") is not True - or not isinstance(cleanup, list) - or any(result.get("status") == "failed" for result in cleanup) - or any( - result.get("resource", "").startswith("app:") - and result.get("status") == "deleted" - and not result.get("confirmed_absent_at") - for result in cleanup - ) - ) - ): - sys.stdout.write("cleanup evidence: required cleanup is incomplete or failed\n") - if isinstance(cleanup, list): - for result in cleanup: - if result.get("status") == "failed": - sys.stdout.write( - f"cleanup failed | {result.get('resource')} | {result.get('detail')}\n" - ) - # teardown() always stamps confirmed_absent_at on deleted apps, so - # this only fires for evidence from an older runner or a hand-edited file. - elif ( - result.get("resource", "").startswith("app:") - and result.get("status") == "deleted" - and not result.get("confirmed_absent_at") - ): - sys.stdout.write( - f"cleanup unconfirmed | {result.get('resource')} |" - " missing confirmed_absent_at\n" - ) - return 1 - if suite == "tools": - grant_checks = document.get("grant_checks", []) - transitive = [c for c in grant_checks if c.get("manual_transitive_grant_applied")] - sys.stdout.write(f"{len(transitive)} transitive grant proofs passed\n") - return 0 - - -def load_suites() -> list[Suite]: - suites = [] - for path in sorted(MATRICES_DIR.glob("*/suite.yaml")): - document = yaml.safe_load(path.read_text(encoding="utf-8")) - for requirement in document.get("requires", []): - if requirement.get("kind") not in REQUIREMENT_KINDS: - raise MatrixError(f"{path.parent.name}: unknown requirement kind in {requirement}.") - suites.append( - Suite( - name=document["name"], - path=path, - tests=list(document["tests"]), - runs=document.get("runs", {}), - requires=list(document.get("requires", [])), - ) - ) - return suites - - -def select_suites(profile: str, names: Sequence[str]) -> list[Suite]: - available = load_suites() - if names: - by_name = {suite.name: suite for suite in available} - unknown = sorted(set(names) - set(by_name)) - if unknown: - raise MatrixError(f"Unknown suite(s) {unknown}; available: {sorted(by_name)}") - selected = [by_name[name] for name in dict.fromkeys(names)] - missing = [suite.name for suite in selected if profile not in suite.runs] - if missing: - raise MatrixError(f"Suite(s) {missing} do not run under the {profile!r} profile.") - return selected - return [suite for suite in available if profile in suite.runs] - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--profile", choices=PROFILES, help="Run profile selecting suites.") - parser.add_argument("--suite", action="append", default=[], help="Suite name; repeatable.") - parser.add_argument( - "--output", type=pathlib.Path, help="Output directory (default: a tempdir)." - ) - parser.add_argument("--wheel", type=pathlib.Path, help="Prebuilt wheel (default: build one).") - parser.add_argument("--commit-sha", help="Source commit (default: HEAD of --source-root).") - parser.add_argument("--source-root", type=pathlib.Path, default=REPOSITORY_ROOT) - parser.add_argument( - "--genie-space-id", - help="Overrides AGENTBRICKS_E2E_GENIE_SPACE_ID for the genie_space requirement.", - ) - parser.add_argument( - "--bridge-sha", help="Immutable bridge commit for generated App dependencies." - ) - parser.add_argument("--warehouse-id") - parser.add_argument( - "--databricks-profile", - help="Databricks CLI profile; default is ambient credentials (DATABRICKS_CONFIG_PROFILE).", - ) - parser.add_argument( - "--app-auth-profile", - help="OAuth profile for deployed App /api calls; defaults to --databricks-profile.", - ) - parser.add_argument( - "--preprovisioned-app-catalog-access", - action="store_true", - help="Skip per-App USE CATALOG grants because catalog access is pre-provisioned.", - ) - parser.add_argument("--keep-resources", action="store_true") - parser.add_argument( - "--pytest-arg", - action="append", - default=[], - help="Extra pytest argument, e.g. -k test_dev (repeatable). Skips the full-matrix evidence check.", - ) - parser.add_argument("--verify-evidence", type=pathlib.Path) - args = parser.parse_args() - if args.verify_evidence is None and args.profile is None: - parser.error("--profile is required unless --verify-evidence is used") - return args - - -def main() -> int: - args = parse_args() - if args.verify_evidence: - return verify_evidence(args.verify_evidence) - try: - suites = select_suites(args.profile, args.suite) - except MatrixError as exc: - sys.stderr.write(f"{exc}\n") - return 2 - if not suites: - sys.stderr.write(f"No suites run under the {args.profile!r} profile.\n") - return 2 - output = args.output or pathlib.Path(tempfile.mkdtemp(prefix="agentbricks-suites-")) - options = Options( - run_profile=args.profile, - output=output.expanduser().resolve(), - wheel=args.wheel, - commit_sha=args.commit_sha, - source_root=args.source_root, - genie_space_id=args.genie_space_id or os.environ.get("AGENTBRICKS_E2E_GENIE_SPACE_ID"), - bridge_sha=args.bridge_sha, - warehouse_id=args.warehouse_id, - databricks_profile=args.databricks_profile, - app_auth_profile=args.app_auth_profile, - preprovisioned_app_catalog_access=args.preprovisioned_app_catalog_access, - keep_resources=args.keep_resources, - pytest_args=args.pytest_arg, - ) - toolchain = Toolchain(options) - try: - toolchain.bootstrap() - except MatrixError as exc: - toolchain.transcript.write(f"bootstrap failed: {exc}") - return 1 - failed = [ - suite.name for suite in suites if SuiteRunner(suite, toolchain, options).execute() != 0 - ] - if failed: - sys.stdout.write(f"suites failed: {failed}; output: {options.output}\n") - return 1 - sys.stdout.write(f"suites passed: {[suite.name for suite in suites]}\n") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/integrations/agentbricks/tests/shared_state.py b/integrations/agentbricks/tests/shared_state.py new file mode 100644 index 000000000..952e529ab --- /dev/null +++ b/integrations/agentbricks/tests/shared_state.py @@ -0,0 +1,74 @@ +"""Create-once, share-everywhere resources for a pytest run, including all xdist workers. + +Each resource is one JSON file in a directory every worker can reach, guarded by a per-resource +file lock (the pattern pytest-xdist documents). The first process to ask creates it while the others +wait; later processes read the cached value. Whatever a creator ``track``s lands in the same file, +which is how the controller finds what to destroy after every worker has finished. +""" + +from __future__ import annotations + +import contextlib +import fcntl +import json +import os +import pathlib +from collections.abc import Callable, Iterator +from typing import Any + +from common import MatrixError + +Track = Callable[..., None] + + +@contextlib.contextmanager +def _locked(path: pathlib.Path) -> Iterator[None]: + descriptor = os.open(path, os.O_CREAT | os.O_RDWR) + try: + fcntl.flock(descriptor, fcntl.LOCK_EX) + yield + finally: + os.close(descriptor) + + +def _read(path: pathlib.Path) -> dict[str, Any]: + return json.loads(path.read_text(encoding="utf-8")) if path.exists() else {} + + +class SharedResources: + def __init__(self, directory: pathlib.Path): + self._directory = directory + + def get_or_create(self, key: str, create: Callable[[Track], dict[str, Any]]) -> dict[str, Any]: + """The cached value for ``key``, calling ``create(track)`` in exactly one process. + + ``track(**fields)`` records what ``create`` has made so far, before it can fail, for + ``tracked``. A failure is cached too, so other workers fail fast instead of retrying a + deploy that takes minutes. + """ + self._directory.mkdir(parents=True, exist_ok=True) + path = self._directory / f"{key}.json" + with _locked(self._directory / f"{key}.lock"): + entry = _read(path) + if "error" in entry: + raise MatrixError(f"{key} setup already failed: {entry['error']}") + if "value" in entry: + return entry["value"] + + def track(**fields: Any) -> None: + entry.setdefault("cleanup", {}).update(fields) + path.write_text(json.dumps(entry), encoding="utf-8") + + try: + entry["value"] = create(track) + except Exception as exc: + entry["error"] = str(exc) + path.write_text(json.dumps(entry), encoding="utf-8") + raise + path.write_text(json.dumps(entry), encoding="utf-8") + return entry["value"] + + def tracked(self) -> dict[str, dict[str, Any]]: + """What each resource's creator recorded, for teardown.""" + entries = {path.stem: _read(path) for path in sorted(self._directory.glob("*.json"))} + return {key: entry["cleanup"] for key, entry in entries.items() if entry.get("cleanup")} diff --git a/integrations/agentbricks/tests/workspace_client.py b/integrations/agentbricks/tests/workspace_client.py index 5828fe954..096875506 100644 --- a/integrations/agentbricks/tests/workspace_client.py +++ b/integrations/agentbricks/tests/workspace_client.py @@ -1,7 +1,6 @@ -"""Databricks SDK helpers for the matrix tests and runner: the workspace, which is not under test. +"""Databricks SDK helpers for the matrix tests and fixtures: the workspace, which is not under test. -Everything here goes through the SDK except ``app_logs``, since the SDK has no Apps log API, and -bundles, which the runner drives with the Databricks CLI. +Everything here goes through the SDK except ``app_logs``, since the SDK has no Apps log API. """ from __future__ import annotations @@ -15,7 +14,7 @@ from collections.abc import Callable, Sequence from typing import Any, cast -from common import TOOL_RESOURCE_PREFIX, MatrixError, Transcript, now +from common import TOOL_RESOURCE_PREFIX, MatrixError, log, now from databricks.sdk import WorkspaceClient from databricks.sdk.errors import DatabricksError, NotFound from databricks.sdk.service.catalog import SecurableType @@ -28,7 +27,6 @@ class Workspace: def __init__( self, profile: str | None, - transcript: Transcript, *, app_auth_profile: str | None = None, warehouse_id: str | None = None, @@ -36,11 +34,12 @@ def __init__( ): self.profile = profile self.app_auth_profile = app_auth_profile or profile - self.transcript = transcript self.warehouse_id = warehouse_id self.preprovisioned_app_catalog_access = preprovisioned_app_catalog_access self.client = WorkspaceClient(profile=profile) self._app_auth_client: WorkspaceClient | None = None + self._app_auth_checked = False + self._warehouse_started = False # App auth: Databricks Apps /api routes need OAuth, so a PAT is rejected. @@ -57,12 +56,15 @@ def _app_client(self) -> WorkspaceClient: return self._app_auth_client def check_app_auth(self) -> None: + if self._app_auth_checked: + return if self._app_client().config.auth_type == "pat": raise MatrixError( f"{self._app_auth_label()} uses a PAT. " "Databricks Apps /api routes require OAuth; run `databricks auth login` " "for a profile on the same workspace." ) + self._app_auth_checked = True @property def app_headers(self) -> dict[str, str]: @@ -72,6 +74,20 @@ def app_headers(self) -> dict[str, str]: raise MatrixError(f"Could not resolve credentials from {self._app_auth_label()}.") return {"Authorization": authorization} + # Existing resources the suite depends on + + def require_catalog(self, name: str) -> None: + try: + self.client.catalogs.get(name) + except DatabricksError as exc: + raise MatrixError(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 MatrixError(f"Genie space {space_id!r} is not accessible: {exc}") from exc + # SQL, warehouse, files def start_warehouse(self, override: str | None = None) -> str: @@ -85,17 +101,24 @@ def start_warehouse(self, override: str | None = None) -> str: (item for item in warehouses if item.state == State.RUNNING), warehouses[0] ) self.warehouse_id = str(running.id) - self.transcript.write(f"# start warehouse {self.warehouse_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: - if self.warehouse_id is None: - raise MatrixError("SQL warehouse was not selected.") - self.transcript.write(f"$ sql: {statement}") + warehouse_id = self._ensure_warehouse() + log(f"$ sql: {statement}") response = self.client.statement_execution.execute_statement( statement=statement, - warehouse_id=self.warehouse_id, + warehouse_id=warehouse_id, wait_timeout="30s", on_wait_timeout=ExecuteStatementRequestOnWaitTimeout.CONTINUE, ) @@ -114,9 +137,13 @@ def sql(self, statement: str, *, timeout: float = 600) -> None: if not response.status or response.status.state != StatementState.SUCCEEDED: raise MatrixError(f"SQL failed: {response.as_dict()}") + def drop(self, kind: str, full_name: str) -> None: + catalog, schema, name = full_name.split(".") + self.sql(f"DROP {kind} IF EXISTS `{catalog}`.`{schema}`.`{name}`") + def upload(self, path: str, data: bytes) -> None: self.client.files.upload(path, io.BytesIO(data), overwrite=True) - self.transcript.write(f"# uploaded {path}") + log(f"# uploaded {path}") def delete_file(self, path: str) -> None: try: @@ -150,7 +177,7 @@ def wait_for_app(self, name: str) -> dict[str, Any]: return app elapsed = time.monotonic() - started if elapsed >= next_tick: - self.transcript.write(f"tick {now():%H:%M} | app-{name} | {state or 'UNKNOWN'}") + log(f"tick {now():%H:%M} | app-{name} | {state or 'UNKNOWN'}") next_tick += 60 time.sleep(15) raise MatrixError(f"App {name} did not become ACTIVE.") @@ -169,9 +196,9 @@ def app_logs(self, name: str, log_path: pathlib.Path) -> pathlib.Path | None: log_path.parent.mkdir(parents=True, exist_ok=True) log_path.write_text(content, encoding="utf-8") except OSError as exc: - self.transcript.write(f"App runtime log capture warning for {name}: {exc}") + log(f"App runtime log capture warning for {name}: {exc}") return None - self.transcript.write(f"App runtime logs captured: {log_path}") + log(f"App runtime logs captured: {log_path}") return log_path # Grants @@ -275,7 +302,7 @@ def _wait_for_app_deleted(self, name: str, timeout: float = 1200) -> None: try: self.client.apps.get(name) except NotFound: - self.transcript.write(f"tick {now():%H:%M} | delete-{name} | absent") + log(f"tick {now():%H:%M} | delete-{name} | absent") return except DatabricksError as exc: raise MatrixError(f"Could not verify deletion of App {name!r}: {exc}") from exc @@ -285,7 +312,7 @@ def _wait_for_app_deleted(self, name: str, timeout: float = 1200) -> None: f"App {name!r} still existed {timeout:.0f}s after delete returned." ) if elapsed >= next_tick: - self.transcript.write(f"tick {now():%H:%M} | delete-{name} | deleting") + log(f"tick {now():%H:%M} | delete-{name} | deleting") next_tick += 60 time.sleep(15) @@ -317,9 +344,7 @@ def app_role_target(self, app_name: str) -> str | None: or not isinstance(branch, str) or not re.fullmatch(r"projects/[^/]+/branches/[^/]+", branch) ): - self.transcript.write( - f"cleanup warning | Lakebase role for {app_name} | ownership not verified" - ) + log(f"cleanup warning | Lakebase role for {app_name} | ownership not verified") return None for role in self.client.postgres.list_roles(parent=branch): data = role.as_dict() @@ -333,7 +358,7 @@ def app_role_target(self, app_name: str) -> str | None: ): return name except Exception as exc: - self.transcript.write(f"cleanup warning | Lakebase role lookup for {app_name} | {exc}") + log(f"cleanup warning | Lakebase role lookup for {app_name} | {exc}") return None def delete_role(self, role_name: str) -> None: @@ -345,82 +370,54 @@ def delete_role(self, role_name: str) -> None: def cleanup_app( workspace: Workspace, - delete_store: Callable[[str, str], subprocess.CompletedProcess[str]], name: str, - info: dict[str, Any], *, - has_app: bool, -) -> list[dict[str, Any]]: - """Delete one project's stores, and its App, runtime store and Lakebase role if it has an App. + delete_store: Callable[[str, str], subprocess.CompletedProcess[str]] | None = None, + memory_store: str | None = None, + session_store: str | None = None, + has_app: bool = True, + runtime_store: bool = False, +) -> list[str]: + """Delete one project's stores and, if it has an App, its runtime store, App and Lakebase role. ``delete_store(kind, store)`` is the agentbricks CLI's store delete, which has no SDK surface. + Returns a description of each failure; an empty list means everything was deleted. """ - transcript = workspace.transcript - results: list[dict[str, Any]] = [] - has_stores = bool(info.get("memory_store_name") or info.get("session_store_name")) + failures: list[str] = [] # Looked up before anything is deleted, since it needs the App and its runtime store. - role_target = workspace.app_role_target(name) if has_app and has_stores else None - stores_ok = True - for label, kind, store in ( - ("memory store", "memory", info.get("memory_store_name")), - ("session store", "sessions", info.get("session_store_name")), - ): - if not store: + role_target = workspace.app_role_target(name) if has_app and runtime_store else None + for kind, store in (("memory", memory_store), ("sessions", session_store)): + if not (store and delete_store): continue try: result = delete_store(kind, store) if result.returncode != 0: - stores_ok = False - results.append( - { - "resource": f"{label}:{store}", - "status": "failed", - "detail": (result.stderr or result.stdout).strip(), - } - ) - else: - results.append({"resource": f"{label}:{store}", "status": "deleted"}) + failures.append(f"{kind} store {store}: {(result.stderr or result.stdout).strip()}") except Exception as exc: - stores_ok = False - transcript.write(f"cleanup warning | {label} {store} | {exc}") - results.append({"resource": f"{label}:{store}", "status": "failed", "detail": str(exc)}) + failures.append(f"{kind} store {store}: {exc}") if not has_app: - return results + return failures runtime_store_deleted = False - if has_stores: + if runtime_store: # The Runtime Store owns a dedicated Lakebase database; it must be deleted before the # role, or the role delete fails on database ownership. try: workspace.delete_runtime_store(name) runtime_store_deleted = True - results.append({"resource": f"runtime-store:{name}", "status": "deleted"}) except Exception as exc: - results.append( - {"resource": f"runtime-store:{name}", "status": "failed", "detail": str(exc)} - ) + failures.append(f"runtime store {name}: {exc}") try: workspace.delete_app(name) except Exception as exc: - results.append({"resource": f"app:{name}", "status": "failed", "detail": str(exc)}) - return results - results.append( - {"resource": f"app:{name}", "status": "deleted", "confirmed_absent_at": now().isoformat()} - ) - if role_target and stores_ok and runtime_store_deleted: + failures.append(f"app {name}: {exc}") + return failures + if role_target and not failures and runtime_store_deleted: try: workspace.delete_role(role_target) - results.append({"resource": f"lakebase-role:{role_target}", "status": "deleted"}) except Exception as exc: - transcript.write(f"cleanup warning | Lakebase role {role_target} | {exc}") - results.append( - { - "resource": f"lakebase-role:{role_target}", - "status": "failed", - "detail": str(exc), - } - ) - return results + failures.append(f"Lakebase role {role_target}: {exc}") + return failures # Grant helpers From fa058bad261c4b33503e0fe5bde8a501a65c820e Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Thu, 8 Oct 2026 23:24:25 +0000 Subject: [PATCH 05/45] [agentbricks] rename e2e backend to target workspace, scope genie space to tools matrix Co-authored-by: Isaac --- .../agentbricks/tests/agentbricks_cli.py | 18 +++-- integrations/agentbricks/tests/conftest.py | 22 +++--- .../tests/matrices/tools/conftest.py | 20 ++++- .../agentbricks/tests/provisioning.py | 79 +++++++++---------- integrations/agentbricks/tests/run_context.py | 30 +++---- .../{backends.py => target_workspace.py} | 42 +++++----- 6 files changed, 114 insertions(+), 97 deletions(-) rename integrations/agentbricks/tests/{backends.py => target_workspace.py} (59%) diff --git a/integrations/agentbricks/tests/agentbricks_cli.py b/integrations/agentbricks/tests/agentbricks_cli.py index be58441a3..37de9a022 100644 --- a/integrations/agentbricks/tests/agentbricks_cli.py +++ b/integrations/agentbricks/tests/agentbricks_cli.py @@ -30,7 +30,6 @@ import tomli import tomlkit -from backends import Backend from common import ( APP_PREFIX, E2E_MODEL, @@ -45,6 +44,7 @@ now, project_prefix, ) +from target_workspace import TargetWorkspace _STORE_KINDS = {"memory": ("memory", "stores"), "sessions": ("sessions", "stores")} _DIRECT_HEADER = ( @@ -138,20 +138,20 @@ def warm_up(self, deadline_seconds: float = 180.0) -> None: class AgentbricksCli: - """Authors, runs and deploys projects with the CLI against a backend.""" + """Authors, runs and deploys projects with the CLI against a target workspace.""" - def __init__(self, backend: Backend, run_config: RunConfig): - self.backend = backend + def __init__(self, target_workspace: TargetWorkspace, run_config: RunConfig): + self.target_workspace = target_workspace self.run_config = run_config self.logs_dir = run_config.output / "logs" self.projects: list[Project] = [] - self._env = child_env(backend.env) + self._env = child_env(target_workspace.env) self._sequence = 0 # Process plumbing def argv(self, *args: str) -> list[str]: - return [_agentbricks_bin(), *self.backend.cli_args, *args] + return [_agentbricks_bin(), *self.target_workspace.cli_args, *args] def cli( self, @@ -420,9 +420,11 @@ def dev(self, project: Project) -> Iterator[Agent]: def deploy(self, project: Project, app: str | None = None) -> Agent: """`agentbricks deploy` the project; deploying again updates the same App.""" - workspace = self.backend.workspace + workspace = self.target_workspace.client if workspace is None: - raise MatrixError(f"deploy needs a workspace; the {self.backend.name} backend has none") + raise MatrixError( + f"deploy needs a live workspace; got --workspace {self.target_workspace.name}" + ) workspace.check_app_auth() name = app or project.app_name creating = not project.app_registered diff --git a/integrations/agentbricks/tests/conftest.py b/integrations/agentbricks/tests/conftest.py index e0109f7a0..a2f535814 100644 --- a/integrations/agentbricks/tests/conftest.py +++ b/integrations/agentbricks/tests/conftest.py @@ -1,12 +1,12 @@ """Fixtures shared by every matrix under ``matrices/``. -``run_context`` (options, backend, run identity) and ``provisioning`` (lazy workspace resources and +``run_context`` (options, target workspace, run identity) and ``provisioning`` (lazy workspace resources and their teardown) are plugins. One command runs a matrix, from ``integrations/agentbricks``: uv run pytest tests/matrices/tools -n 6 --dist loadgroup \ - --databricks-profile --genie-space-id + --databricks-profile -``--backend fake`` runs against no workspace; whatever needs one is skipped. +``--workspace fake`` runs against no workspace; whatever needs one is skipped. The CLI wrapper is function-scoped: each test authors, runs and deploys its own projects, and the fixture deletes every App and store they created when the test ends. @@ -18,8 +18,8 @@ import pytest from agentbricks_cli import AgentbricksCli -from backends import Backend from common import RunConfig, log +from target_workspace import TargetWorkspace from workspace_client import Workspace, cleanup_app pytest_plugins = ["run_context", "provisioning"] @@ -34,21 +34,21 @@ def pytest_runtest_logreport(report: pytest.TestReport) -> None: @pytest.fixture(scope="session") -def workspace_client(backend: Backend) -> Workspace: - workspace = backend.workspace +def workspace_client(target_workspace: TargetWorkspace) -> Workspace: + workspace = target_workspace.client if workspace is None: - pytest.skip(f"the {backend.name} backend has no Databricks workspace") + pytest.skip(f"--workspace {target_workspace.name} has no Databricks SDK access") return workspace @pytest.fixture def agentbricks_cli( - request: pytest.FixtureRequest, backend: Backend, run_config: RunConfig + request: pytest.FixtureRequest, target_workspace: TargetWorkspace, run_config: RunConfig ) -> Iterator[AgentbricksCli]: - cli = AgentbricksCli(backend, run_config) + cli = AgentbricksCli(target_workspace, run_config) yield cli - if cli.projects and backend.workspace is not None: - _delete_projects(cli, backend.workspace, failed=request.node.nodeid in _failed_tests) + if cli.projects and target_workspace.client is not None: + _delete_projects(cli, target_workspace.client, failed=request.node.nodeid in _failed_tests) def _delete_projects(cli: AgentbricksCli, workspace: Workspace, *, failed: bool) -> None: diff --git a/integrations/agentbricks/tests/matrices/tools/conftest.py b/integrations/agentbricks/tests/matrices/tools/conftest.py index 80e425db1..06c5db193 100644 --- a/integrations/agentbricks/tests/matrices/tools/conftest.py +++ b/integrations/agentbricks/tests/matrices/tools/conftest.py @@ -1,11 +1,15 @@ -"""Tools matrix: every test runs once per authoring path, over the tools the backend can support.""" +"""Tools matrix: every test runs once per authoring path, over the tools the target workspace can support.""" from __future__ import annotations +import os import pathlib +import re import pytest -from common import AUTHORING_PATHS, log +from common import AUTHORING_PATHS, MatrixError, log +from provisioning import require_workspace +from target_workspace import TargetWorkspace from tools import TOOL_DEFINITIONS, Tool @@ -19,6 +23,18 @@ def scratch_schema(deploy_scratch_schema) -> str: return deploy_scratch_schema(pathlib.Path(__file__).parent / "bundle") +@pytest.fixture(scope="session") +def genie_space(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 = os.environ.get("AGENTBRICKS_E2E_GENIE_SPACE_ID") + if not space_id: + pytest.skip("no Genie space; set AGENTBRICKS_E2E_GENIE_SPACE_ID") + if re.fullmatch(r"[0-9a-f]{32}", space_id) is None: + raise MatrixError(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(scope="session") def tools(request: pytest.FixtureRequest) -> tuple[Tool, ...]: """The tools whose fixtures are all available; the rest are logged and left out.""" diff --git a/integrations/agentbricks/tests/provisioning.py b/integrations/agentbricks/tests/provisioning.py index e6fef7c35..88e9b6cc2 100644 --- a/integrations/agentbricks/tests/provisioning.py +++ b/integrations/agentbricks/tests/provisioning.py @@ -1,8 +1,8 @@ """pytest plugin: the workspace resources tests ask for, as lazy session fixtures, and their teardown. -A fixture exists per requirement (``catalog``, ``scratch_schema``, ``uc_function``, -``genie_space``) and is only created when a selected test needs it, so a run that never asks for -``uc_function`` never deploys a bundle or creates a function. Each fixture skips when the backend +A fixture exists per requirement (``catalog``, ``scratch_schema``, ``uc_function``) and is only +created when a selected test needs it, so a run that never asks for ``uc_function`` never deploys a +bundle or creates a function. Each fixture skips when the target workspace cannot provide it or the operator did not configure it. Creation happens once per run across xdist workers (``shared_state``). Teardown happens once, in @@ -14,14 +14,12 @@ import dataclasses import json -import re import uuid from collections.abc import Callable from functools import partial from pathlib import Path import pytest -from backends import Backend, BackendUnavailable from common import ( APP_PREFIX, MatrixError, @@ -31,8 +29,9 @@ project_prefix, run_command, ) -from run_context import base_run_config, build_backend, started_run_config +from run_context import base_run_config, build_target_workspace, started_run_config from shared_state import SharedResources, Track +from target_workspace import TargetWorkspace, TargetWorkspaceUnavailable from workspace_client import Workspace, cleanup_app BUNDLE_TARGET = "nightly" @@ -51,36 +50,25 @@ class UcFunction: volume_marker: str -def _require_workspace(backend: Backend, what: str) -> Workspace: - workspace = backend.workspace +def require_workspace(target_workspace: TargetWorkspace, what: str) -> Workspace: + workspace = target_workspace.client if workspace is None: - pytest.skip(f"{what} needs a Databricks workspace; the {backend.name} backend has none") + pytest.skip(f"{what} needs a live workspace; got --workspace {target_workspace.name}") return workspace @pytest.fixture(scope="session") -def catalog(request: pytest.FixtureRequest, backend: Backend) -> str: - workspace = _require_workspace(backend, "a UC catalog") +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 -@pytest.fixture(scope="session") -def genie_space(request: pytest.FixtureRequest, backend: Backend) -> str: - space_id = request.config.getoption("genie_space_id") - if not space_id: - pytest.skip("no Genie space; pass --genie-space-id or set AGENTBRICKS_E2E_GENIE_SPACE_ID") - if re.fullmatch(r"[0-9a-f]{32}", space_id) is None: - raise MatrixError(f"Genie space id {space_id!r} is not 32 lowercase hex characters.") - _require_workspace(backend, "a Genie space").require_genie_space(space_id) - return space_id - - @pytest.fixture(scope="session") def deploy_scratch_schema( catalog: str, - backend: Backend, + target_workspace: TargetWorkspace, shared_resources: SharedResources, request: pytest.FixtureRequest, ) -> Callable[[Path], str]: @@ -93,7 +81,7 @@ def deploy_scratch_schema( def deploy(bundle_dir: Path) -> str: value = shared_resources.get_or_create( "scratch_schema", - lambda track: _deploy_bundle(backend, run, catalog, bundle_dir, track), + lambda track: _deploy_bundle(target_workspace, run, catalog, bundle_dir, track), ) if value["bundle_dir"] != str(bundle_dir.resolve()): raise MatrixError( @@ -107,10 +95,10 @@ def deploy(bundle_dir: Path) -> str: @pytest.fixture(scope="session") def uc_function( - scratch_schema: str, backend: Backend, shared_resources: SharedResources + scratch_schema: str, target_workspace: TargetWorkspace, shared_resources: SharedResources ) -> UcFunction: """Marker functions and a marker volume in ``scratch_schema`` (a fixture each matrix defines).""" - workspace = _require_workspace(backend, "a UC function") + workspace = require_workspace(target_workspace, "a UC function") value = shared_resources.get_or_create( "uc_function", lambda track: _create_uc_function(workspace, scratch_schema, track) ) @@ -120,29 +108,33 @@ def uc_function( # Creation -def _bundle_env(backend: Backend, catalog: str, run_id: str) -> dict[str, str]: +def _bundle_env(target_workspace: TargetWorkspace, catalog: str, run_id: str) -> dict[str, str]: return child_env( - {**backend.env, "BUNDLE_VAR_catalog_name": catalog, "BUNDLE_VAR_run_id": run_id} + {**target_workspace.env, "BUNDLE_VAR_catalog_name": catalog, "BUNDLE_VAR_run_id": run_id} ) def _bundle( - backend: Backend, bundle_dir: Path, env: dict[str, str], *args: str, timeout: float + target_workspace: TargetWorkspace, + bundle_dir: Path, + env: dict[str, str], + *args: str, + timeout: float, ) -> str: - argv = ["databricks", "bundle", *args, "-t", BUNDLE_TARGET, *backend.cli_args] + argv = ["databricks", "bundle", *args, "-t", BUNDLE_TARGET, *target_workspace.cli_args] return run_command(argv, cwd=bundle_dir, env=env, timeout=timeout) def _deploy_bundle( - backend: Backend, run: RunConfig, catalog: str, bundle_dir: Path, track: Track + target_workspace: TargetWorkspace, run: RunConfig, catalog: str, bundle_dir: Path, track: Track ) -> dict[str, str]: bundle_dir = bundle_dir.resolve() - env = _bundle_env(backend, catalog, run.run_id) + env = _bundle_env(target_workspace, catalog, run.run_id) # Tracked first: a deploy that fails midway still leaves a schema to destroy. track(bundle_dir=str(bundle_dir), catalog=catalog, run_id=run.run_id) - _bundle(backend, bundle_dir, env, "deploy", timeout=900) + _bundle(target_workspace, bundle_dir, env, "deploy", timeout=900) summary = json.loads( - _bundle(backend, bundle_dir, env, "summary", "--output", "json", timeout=120) + _bundle(target_workspace, bundle_dir, env, "summary", "--output", "json", timeout=120) ) resource = summary.get("resources", {}).get("schemas", {}).get(SCHEMA_RESOURCE) if not isinstance(resource, dict): @@ -221,16 +213,21 @@ def pytest_sessionfinish(session: pytest.Session) -> None: log("Resources retained (--keep-resources); skipping teardown.") return try: - backend = build_backend(config) - except BackendUnavailable: + target_workspace = build_target_workspace(config) + except TargetWorkspaceUnavailable: return - if backend.workspace is not None and _teardown(backend, backend.workspace, run, tracked): + if target_workspace.client is not None and _teardown( + target_workspace, target_workspace.client, run, tracked + ): if session.exitstatus == 0: session.exitstatus = pytest.ExitCode.TESTS_FAILED def _teardown( - backend: Backend, workspace: Workspace, run: RunConfig, tracked: dict[str, dict[str, str]] + target_workspace: TargetWorkspace, + workspace: Workspace, + run: RunConfig, + tracked: dict[str, dict[str, str]], ) -> list[str]: """Remove what the run created and return a description of each failure.""" failures: list[str] = [] @@ -252,11 +249,13 @@ def attempt(label: str, action: Callable[[], object]) -> None: if volume := functions.get("volume"): attempt(f"volume {volume}", partial(workspace.drop, "VOLUME", volume)) if bundle := tracked.get("scratch_schema"): - env = _bundle_env(backend, bundle["catalog"], bundle["run_id"]) + env = _bundle_env(target_workspace, bundle["catalog"], bundle["run_id"]) bundle_dir = Path(bundle["bundle_dir"]) attempt( "bundle destroy", - lambda: _bundle(backend, bundle_dir, env, "destroy", "--auto-approve", timeout=900), + lambda: _bundle( + target_workspace, bundle_dir, env, "destroy", "--auto-approve", timeout=900 + ), ) return failures diff --git a/integrations/agentbricks/tests/run_context.py b/integrations/agentbricks/tests/run_context.py index a248b2dbc..aedefbbef 100644 --- a/integrations/agentbricks/tests/run_context.py +++ b/integrations/agentbricks/tests/run_context.py @@ -1,4 +1,4 @@ -"""pytest plugin: command-line options, the run's identity, and the backend the suite runs against. +"""pytest plugin: command-line options, the run's identity, and the workspace the suite runs against. Under xdist the controller picks the run id and output directory and hands them to each worker in ``workerinput``, so every process names, finds and later tears down the same resources. @@ -15,9 +15,14 @@ from typing import Any import pytest -from backends import BACKENDS, Backend, BackendUnavailable, make_backend from common import MatrixError, RunConfig, log, run_command from shared_state import SharedResources +from target_workspace import ( + WORKSPACE_KINDS, + TargetWorkspace, + TargetWorkspaceUnavailable, + make_target_workspace, +) PACKAGE_DIR = pathlib.Path(__file__).resolve().parent.parent _RUN_KEY = pytest.StashKey[RunConfig]() @@ -27,8 +32,8 @@ def pytest_addoption(parser: pytest.Parser) -> None: group = parser.getgroup("agentbricks", "Agent Bricks E2E matrices") group.addoption( - "--backend", - choices=BACKENDS, + "--workspace", + choices=WORKSPACE_KINDS, default="live", help="live runs against a Databricks workspace; fake skips whatever needs one.", ) @@ -45,11 +50,6 @@ def pytest_addoption(parser: pytest.Parser) -> None: default=os.environ.get("AGENTBRICKS_E2E_CATALOG", "main"), help="Existing catalog for scratch schemas (AGENTBRICKS_E2E_CATALOG, default main).", ) - group.addoption( - "--genie-space-id", - default=os.environ.get("AGENTBRICKS_E2E_GENIE_SPACE_ID"), - help="Genie space for the genie_space fixture (AGENTBRICKS_E2E_GENIE_SPACE_ID).", - ) group.addoption("--bridge-sha", help="Immutable bridge commit for generated App dependencies.") group.addoption( "--wheel", @@ -109,19 +109,19 @@ def started_run_config(config: pytest.Config) -> RunConfig | None: return config.stash.get(_RUN_KEY, None) -def build_backend(config: pytest.Config) -> Backend: - return make_backend( - config.getoption("backend"), +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 backend(request: pytest.FixtureRequest) -> Backend: +def target_workspace(request: pytest.FixtureRequest) -> TargetWorkspace: try: - return build_backend(request.config) - except BackendUnavailable as exc: + return build_target_workspace(request.config) + except TargetWorkspaceUnavailable as exc: pytest.skip(str(exc)) diff --git a/integrations/agentbricks/tests/backends.py b/integrations/agentbricks/tests/target_workspace.py similarity index 59% rename from integrations/agentbricks/tests/backends.py rename to integrations/agentbricks/tests/target_workspace.py index 78df6f273..079a1b636 100644 --- a/integrations/agentbricks/tests/backends.py +++ b/integrations/agentbricks/tests/target_workspace.py @@ -1,7 +1,7 @@ -"""Where the system under test runs: what the CLI subprocess needs, and the workspace tests may touch. +"""The Databricks workspace the system under test talks to (``--workspace live|fake``). -``live`` is a real Databricks workspace. ``fake`` is an in-process stand-in that has no workspace, so -every requirement that needs one is skipped rather than failed. +``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 @@ -12,14 +12,14 @@ from common import RunConfig from workspace_client import Workspace -BACKENDS = ("live", "fake") +WORKSPACE_KINDS = ("live", "fake") -class BackendUnavailable(Exception): - """The selected backend cannot run in this environment; the suite skips instead of failing.""" +class TargetWorkspaceUnavailable(Exception): + """The selected workspace cannot run in this environment; the suite skips instead of failing.""" -class Backend(Protocol): +class TargetWorkspace(Protocol): name: str @property @@ -33,18 +33,18 @@ def cli_args(self) -> Sequence[str]: ... @property - def workspace(self) -> Workspace | None: - """SDK access for grants, cleanup and UC setup; None when the backend has no workspace.""" + def client(self) -> Workspace | None: + """SDK access for grants, cleanup and UC setup; None for the fake workspace.""" ... -class LiveBackend: +class LiveWorkspace: name = "live" def __init__(self, profile: str | None, run_config: RunConfig): self.profile = profile self._run_config = run_config - self._workspace: Workspace | None = None + self._client: Workspace | None = None @property def env(self) -> Mapping[str, str]: @@ -56,30 +56,30 @@ def cli_args(self) -> Sequence[str]: return ["--profile", self.profile] if self.profile else [] @property - def workspace(self) -> Workspace: - if self._workspace is None: - self._workspace = Workspace( + def client(self) -> Workspace: + if self._client is None: + self._client = Workspace( self.profile, app_auth_profile=self._run_config.app_auth_profile, warehouse_id=self._run_config.warehouse_id, preprovisioned_app_catalog_access=self._run_config.preprovisioned_app_catalog_access, ) - return self._workspace + return self._client -class FakeBackend: +class FakeWorkspace: """Seam for the in-process fake: DATABRICKS_HOST/DATABRICKS_TOKEN env, no CLI args, no Workspace.""" name = "fake" env: Mapping[str, str] = {} cli_args: Sequence[str] = () - workspace: Workspace | None = None + client: Workspace | None = None def __init__(self, run_config: RunConfig): - raise BackendUnavailable("fake backend not implemented yet") + raise TargetWorkspaceUnavailable("fake workspace not implemented yet") -def make_backend(name: str, profile: str | None, run_config: RunConfig) -> Backend: +def make_target_workspace(name: str, profile: str | None, run_config: RunConfig) -> TargetWorkspace: if name == "fake": - return FakeBackend(run_config) - return LiveBackend(profile, run_config) + return FakeWorkspace(run_config) + return LiveWorkspace(profile, run_config) From ffd4bda1fab4083c0a959bc66b42084e59b31acb Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Thu, 8 Oct 2026 23:40:03 +0000 Subject: [PATCH 06/45] [agentbricks] drop unused scratch-schema bundles from store and tracing matrices Co-authored-by: Isaac --- .../memory_store/bundle/databricks.yml | 26 ------------------- .../tests/matrices/memory_store/conftest.py | 10 ------- .../session_store/bundle/databricks.yml | 26 ------------------- .../tests/matrices/session_store/conftest.py | 10 ------- .../matrices/tracing/bundle/databricks.yml | 26 ------------------- .../tests/matrices/tracing/conftest.py | 10 ------- .../agentbricks/tests/provisioning.py | 2 +- 7 files changed, 1 insertion(+), 109 deletions(-) delete mode 100644 integrations/agentbricks/tests/matrices/memory_store/bundle/databricks.yml delete mode 100644 integrations/agentbricks/tests/matrices/memory_store/conftest.py delete mode 100644 integrations/agentbricks/tests/matrices/session_store/bundle/databricks.yml delete mode 100644 integrations/agentbricks/tests/matrices/session_store/conftest.py delete mode 100644 integrations/agentbricks/tests/matrices/tracing/bundle/databricks.yml delete mode 100644 integrations/agentbricks/tests/matrices/tracing/conftest.py diff --git a/integrations/agentbricks/tests/matrices/memory_store/bundle/databricks.yml b/integrations/agentbricks/tests/matrices/memory_store/bundle/databricks.yml deleted file mode 100644 index ae61d6a26..000000000 --- a/integrations/agentbricks/tests/matrices/memory_store/bundle/databricks.yml +++ /dev/null @@ -1,26 +0,0 @@ -# Per-run scratch schema for this matrix; the UC function and volume live inside it. -bundle: - name: agentbricks-memory-store-scratch - -variables: - catalog_name: {} - run_id: {} - -targets: - nightly: - default: true - variables: - catalog_name: main - dev: - variables: - catalog_name: main - -workspace: - # State is per run so concurrent or repeated runs never rename each other's schema. - root_path: ~/.bundle/${bundle.name}/${bundle.target}/${var.run_id} - -resources: - schemas: - scratch_schema: - catalog_name: ${var.catalog_name} - name: ab_e2e_${var.run_id} diff --git a/integrations/agentbricks/tests/matrices/memory_store/conftest.py b/integrations/agentbricks/tests/matrices/memory_store/conftest.py deleted file mode 100644 index 71e05baf6..000000000 --- a/integrations/agentbricks/tests/matrices/memory_store/conftest.py +++ /dev/null @@ -1,10 +0,0 @@ -from __future__ import annotations - -import pathlib - -import pytest - - -@pytest.fixture(scope="session") -def scratch_schema(deploy_scratch_schema) -> str: - return deploy_scratch_schema(pathlib.Path(__file__).parent / "bundle") diff --git a/integrations/agentbricks/tests/matrices/session_store/bundle/databricks.yml b/integrations/agentbricks/tests/matrices/session_store/bundle/databricks.yml deleted file mode 100644 index ce606e4aa..000000000 --- a/integrations/agentbricks/tests/matrices/session_store/bundle/databricks.yml +++ /dev/null @@ -1,26 +0,0 @@ -# Per-run scratch schema for this matrix; the UC function and volume live inside it. -bundle: - name: agentbricks-session-store-scratch - -variables: - catalog_name: {} - run_id: {} - -targets: - nightly: - default: true - variables: - catalog_name: main - dev: - variables: - catalog_name: main - -workspace: - # State is per run so concurrent or repeated runs never rename each other's schema. - root_path: ~/.bundle/${bundle.name}/${bundle.target}/${var.run_id} - -resources: - schemas: - scratch_schema: - catalog_name: ${var.catalog_name} - name: ab_e2e_${var.run_id} diff --git a/integrations/agentbricks/tests/matrices/session_store/conftest.py b/integrations/agentbricks/tests/matrices/session_store/conftest.py deleted file mode 100644 index 71e05baf6..000000000 --- a/integrations/agentbricks/tests/matrices/session_store/conftest.py +++ /dev/null @@ -1,10 +0,0 @@ -from __future__ import annotations - -import pathlib - -import pytest - - -@pytest.fixture(scope="session") -def scratch_schema(deploy_scratch_schema) -> str: - return deploy_scratch_schema(pathlib.Path(__file__).parent / "bundle") diff --git a/integrations/agentbricks/tests/matrices/tracing/bundle/databricks.yml b/integrations/agentbricks/tests/matrices/tracing/bundle/databricks.yml deleted file mode 100644 index f4c86aeac..000000000 --- a/integrations/agentbricks/tests/matrices/tracing/bundle/databricks.yml +++ /dev/null @@ -1,26 +0,0 @@ -# Per-run scratch schema for this matrix; the UC function and volume live inside it. -bundle: - name: agentbricks-tracing-scratch - -variables: - catalog_name: {} - run_id: {} - -targets: - nightly: - default: true - variables: - catalog_name: main - dev: - variables: - catalog_name: main - -workspace: - # State is per run so concurrent or repeated runs never rename each other's schema. - root_path: ~/.bundle/${bundle.name}/${bundle.target}/${var.run_id} - -resources: - schemas: - scratch_schema: - catalog_name: ${var.catalog_name} - name: ab_e2e_${var.run_id} diff --git a/integrations/agentbricks/tests/matrices/tracing/conftest.py b/integrations/agentbricks/tests/matrices/tracing/conftest.py deleted file mode 100644 index 71e05baf6..000000000 --- a/integrations/agentbricks/tests/matrices/tracing/conftest.py +++ /dev/null @@ -1,10 +0,0 @@ -from __future__ import annotations - -import pathlib - -import pytest - - -@pytest.fixture(scope="session") -def scratch_schema(deploy_scratch_schema) -> str: - return deploy_scratch_schema(pathlib.Path(__file__).parent / "bundle") diff --git a/integrations/agentbricks/tests/provisioning.py b/integrations/agentbricks/tests/provisioning.py index 88e9b6cc2..4817e74b1 100644 --- a/integrations/agentbricks/tests/provisioning.py +++ b/integrations/agentbricks/tests/provisioning.py @@ -97,7 +97,7 @@ def deploy(bundle_dir: Path) -> str: def uc_function( scratch_schema: str, target_workspace: TargetWorkspace, shared_resources: SharedResources ) -> UcFunction: - """Marker functions and a marker volume in ``scratch_schema`` (a fixture each matrix defines).""" + """Marker functions and a marker volume in ``scratch_schema`` (defined by the matrix that needs one).""" workspace = require_workspace(target_workspace, "a UC function") value = shared_resources.get_or_create( "uc_function", lambda track: _create_uc_function(workspace, scratch_schema, track) From 4ef55a7d1e493ee3a7715ff556a3086c5f8f07c3 Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Fri, 9 Oct 2026 20:32:00 +0000 Subject: [PATCH 07/45] [agentbricks] per-test temporary schemas and full Lakebase cleanup in e2e matrices Co-authored-by: Isaac --- .../agentbricks/tests/agentbricks_cli.py | 29 +- integrations/agentbricks/tests/common.py | 11 +- integrations/agentbricks/tests/conftest.py | 10 +- .../matrices/tools/bundle/databricks.yml | 26 -- .../tests/matrices/tools/conftest.py | 48 ++- .../tests/matrices/tools/test_grants.py | 2 +- .../tests/matrices/tools/test_selection.py | 2 +- .../agentbricks/tests/matrices/tools/tools.py | 2 +- .../tests/matrices/tools/uc_objects.py | 62 ++++ .../agentbricks/tests/provisioning.py | 312 +++++----------- integrations/agentbricks/tests/run_context.py | 34 +- .../agentbricks/tests/shared_state.py | 74 ---- .../agentbricks/tests/workspace_client.py | 336 ++++++++++++++---- 13 files changed, 509 insertions(+), 439 deletions(-) delete mode 100644 integrations/agentbricks/tests/matrices/tools/bundle/databricks.yml create mode 100644 integrations/agentbricks/tests/matrices/tools/uc_objects.py delete mode 100644 integrations/agentbricks/tests/shared_state.py diff --git a/integrations/agentbricks/tests/agentbricks_cli.py b/integrations/agentbricks/tests/agentbricks_cli.py index 37de9a022..c1f01a723 100644 --- a/integrations/agentbricks/tests/agentbricks_cli.py +++ b/integrations/agentbricks/tests/agentbricks_cli.py @@ -36,7 +36,6 @@ FRAMEWORKS, MatrixError, RunConfig, - child_env, last_lines, last_nonempty_line, log, @@ -45,6 +44,7 @@ project_prefix, ) from target_workspace import TargetWorkspace +from workspace_client import AppIdentity, Workspace, record_app_identity _STORE_KINDS = {"memory": ("memory", "stores"), "sessions": ("sessions", "stores")} _DIRECT_HEADER = ( @@ -74,6 +74,8 @@ class Project: # Set only once this project registered the App, so cleanup never deletes one it did not create. app_registered: bool = False environment_prepared: bool = False + # Recorded at deploy time: cleanup cannot look this up once the App or its store is gone. + identity: AppIdentity = dataclasses.field(default_factory=AppIdentity) class Agent: @@ -145,7 +147,7 @@ def __init__(self, target_workspace: TargetWorkspace, run_config: RunConfig): self.run_config = run_config self.logs_dir = run_config.output / "logs" self.projects: list[Project] = [] - self._env = child_env(target_workspace.env) + self._env = {**os.environ, **target_workspace.env} self._sequence = 0 # Process plumbing @@ -433,12 +435,16 @@ def deploy(self, project: Project, app: str | None = None) -> Agent: # Deploy can create the App and then fail while waiting for it, so mark it first. project.app_registered = True label = f"deploy-{name}-{self._next()}" - self._run_long( - label, - self.argv("deploy", name, "--source", str(project.path)), - timeout=2400, - ) - deployed = workspace.wait_for_app(name) + try: + self._run_long( + label, + self.argv("deploy", name, "--source", str(project.path)), + timeout=2400, + ) + deployed = workspace.wait_for_app(name) + finally: + # Also after a failed deploy, which can leave an App, store and role behind. + self._record_identity(project, name, workspace) url = str(deployed.get("url") or "").rstrip("/") if not url: raise MatrixError(f"App {name} has no URL: {deployed}") @@ -453,6 +459,13 @@ def deploy(self, project: Project, app: str | None = None) -> Agent: agent.warm_up() return agent + def _record_identity(self, project: Project, app: str, workspace: Workspace) -> None: + """Keep the App's principal and Lakebase branch for cleanup, and for the controller's sweep.""" + identity = project.identity.merged(workspace.app_identity(app)) + if identity != project.identity: + project.identity = identity + record_app_identity(self.run_config.output, app, identity) + def _wait_for_local( self, process: subprocess.Popen[str], port: int, label: str, log_path: pathlib.Path ) -> None: diff --git a/integrations/agentbricks/tests/common.py b/integrations/agentbricks/tests/common.py index 241821a07..71d8f1ef4 100644 --- a/integrations/agentbricks/tests/common.py +++ b/integrations/agentbricks/tests/common.py @@ -8,7 +8,6 @@ import dataclasses import datetime as dt -import os import pathlib import shlex import subprocess @@ -35,6 +34,11 @@ def project_prefix(run_id: str) -> str: return f"t-{run_id}-" +def scratch_schema_prefix(run_id: str) -> str: + """Names every temporary schema of a run, so the controller can sweep ones a test left behind.""" + return f"ab_e2e_{run_id}_" + + def log(text: str) -> None: sys.stdout.write(text.rstrip() + "\n") sys.stdout.flush() @@ -99,8 +103,3 @@ class RunConfig: preprovisioned_app_catalog_access: bool bridge_sha: str | None wheel: pathlib.Path | None - - -def child_env(extra: Mapping[str, str] | None = None) -> dict[str, str]: - """The environment for a subprocess: ours plus ``extra``.""" - return {**os.environ, **(extra or {})} diff --git a/integrations/agentbricks/tests/conftest.py b/integrations/agentbricks/tests/conftest.py index a2f535814..4a355babb 100644 --- a/integrations/agentbricks/tests/conftest.py +++ b/integrations/agentbricks/tests/conftest.py @@ -1,7 +1,7 @@ """Fixtures shared by every matrix under ``matrices/``. -``run_context`` (options, target workspace, run identity) and ``provisioning`` (lazy workspace resources and -their teardown) are plugins. One command runs a matrix, from ``integrations/agentbricks``: +``run_context`` (options, target workspace, run identity) and ``provisioning`` (the catalog +requirement and the end-of-run cleanup) are plugins. One command runs a matrix, from ``integrations/agentbricks``: uv run pytest tests/matrices/tools -n 6 --dist loadgroup \ --databricks-profile @@ -52,7 +52,7 @@ def agentbricks_cli( def _delete_projects(cli: AgentbricksCli, workspace: Workspace, *, failed: bool) -> None: - """Delete every App, store and role the test's projects created; the sweep gets the rest.""" + """Delete every App, store, source folder and Lakebase leftover the test's projects created.""" for project in cli.projects: if failed and project.app_registered: # The App is deleted next, so a failing test must capture its logs while it exists. @@ -67,7 +67,9 @@ def _delete_projects(cli: AgentbricksCli, workspace: Workspace, *, failed: bool) memory_store=project.memory_store_name, session_store=project.session_store_name, has_app=project.app_registered, - runtime_store=bool(project.memory_store_name or project.session_store_name), + # Every managed deploy creates a Runtime Store, with or without memory/session stores. + runtime_store=project.app_registered, + identity=project.identity, ) except Exception as exc: failures = [str(exc)] diff --git a/integrations/agentbricks/tests/matrices/tools/bundle/databricks.yml b/integrations/agentbricks/tests/matrices/tools/bundle/databricks.yml deleted file mode 100644 index ca46ec79e..000000000 --- a/integrations/agentbricks/tests/matrices/tools/bundle/databricks.yml +++ /dev/null @@ -1,26 +0,0 @@ -# Per-run scratch schema for this matrix; the UC function and volume live inside it. -bundle: - name: agentbricks-tools-scratch - -variables: - catalog_name: {} - run_id: {} - -targets: - nightly: - default: true - variables: - catalog_name: main - dev: - variables: - catalog_name: main - -workspace: - # State is per run so concurrent or repeated runs never rename each other's schema. - root_path: ~/.bundle/${bundle.name}/${bundle.target}/${var.run_id} - -resources: - schemas: - scratch_schema: - catalog_name: ${var.catalog_name} - name: ab_e2e_${var.run_id} diff --git a/integrations/agentbricks/tests/matrices/tools/conftest.py b/integrations/agentbricks/tests/matrices/tools/conftest.py index 06c5db193..e4f33e4af 100644 --- a/integrations/agentbricks/tests/matrices/tools/conftest.py +++ b/integrations/agentbricks/tests/matrices/tools/conftest.py @@ -1,16 +1,31 @@ -"""Tools matrix: every test runs once per authoring path, over the tools the target workspace can support.""" +"""Tools matrix: every test runs once per authoring path, over the tools the target workspace can support. + +Each test gets its own temporary UC schema (``ab_e2e__``) holding the function and +volume its tools use, and deletes it, grants included, in its own teardown. The controller sweeps +any a crashed test left behind by the run's prefix (see ``provisioning``). +""" from __future__ import annotations +import datetime as dt import os -import pathlib import re +import uuid +from collections.abc import Iterator import pytest -from common import AUTHORING_PATHS, MatrixError, log +from common import ( + AUTHORING_PATHS, + MatrixError, + RunConfig, + log, + now, + scratch_schema_prefix, +) from provisioning import require_workspace from target_workspace import TargetWorkspace from tools import TOOL_DEFINITIONS, Tool +from uc_objects import UcFunction, create_uc_function @pytest.fixture(params=AUTHORING_PATHS) @@ -18,9 +33,28 @@ def authoring(request: pytest.FixtureRequest) -> str: return request.param -@pytest.fixture(scope="session") -def scratch_schema(deploy_scratch_schema) -> str: - return deploy_scratch_schema(pathlib.Path(__file__).parent / "bundle") +@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 separate sweeper 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 functions and volume in ``scratch_schema``.""" + return create_uc_function(require_workspace(target_workspace, "a UC function"), scratch_schema) @pytest.fixture(scope="session") @@ -35,7 +69,7 @@ def genie_space(target_workspace: TargetWorkspace) -> str: return space_id -@pytest.fixture(scope="session") +@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] = [] diff --git a/integrations/agentbricks/tests/matrices/tools/test_grants.py b/integrations/agentbricks/tests/matrices/tools/test_grants.py index 8cdcc8dbc..eb4af2c3d 100644 --- a/integrations/agentbricks/tests/matrices/tools/test_grants.py +++ b/integrations/agentbricks/tests/matrices/tools/test_grants.py @@ -4,8 +4,8 @@ from agentbricks_cli import AgentbricksCli from common import TOOL_RESOURCE_PREFIX -from provisioning import UcFunction from tools import bind, uc_function_tool +from uc_objects import UcFunction from workspace_client import Workspace diff --git a/integrations/agentbricks/tests/matrices/tools/test_selection.py b/integrations/agentbricks/tests/matrices/tools/test_selection.py index 01f2f807c..1dd4f49b1 100644 --- a/integrations/agentbricks/tests/matrices/tools/test_selection.py +++ b/integrations/agentbricks/tests/matrices/tools/test_selection.py @@ -7,8 +7,8 @@ from __future__ import annotations from agentbricks_cli import AgentbricksCli -from provisioning import UcFunction from tools import PYTHON_MARKER, UC_MARKER, Tool, bind +from uc_objects import UcFunction from workspace_client import Workspace # Names the Unity Catalog function without its tool id, and the Python helper only to exclude it. diff --git a/integrations/agentbricks/tests/matrices/tools/tools.py b/integrations/agentbricks/tests/matrices/tools/tools.py index 58f5422d9..264e7b67e 100644 --- a/integrations/agentbricks/tests/matrices/tools/tools.py +++ b/integrations/agentbricks/tests/matrices/tools/tools.py @@ -7,7 +7,7 @@ from collections.abc import Callable, Mapping from agentbricks_cli import Agent, AgentbricksCli, Project -from provisioning import UcFunction +from uc_objects import UcFunction from workspace_client import GrantTuple PYTHON_MARKER_FILE = "agent/tools/matrix_marker.py" diff --git a/integrations/agentbricks/tests/matrices/tools/uc_objects.py b/integrations/agentbricks/tests/matrices/tools/uc_objects.py new file mode 100644 index 000000000..bfd9234c9 --- /dev/null +++ b/integrations/agentbricks/tests/matrices/tools/uc_objects.py @@ -0,0 +1,62 @@ +"""The UC objects the tools matrix uses, created fresh in each test's temporary schema.""" + +from __future__ import annotations + +import dataclasses +import uuid + +from common import MatrixError +from workspace_client import Workspace + +NESTED_FUNCTION = "ab_nested_marker" +FUNCTION = "ab_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 + # Called by ``function`` and deliberately never declared in agent.toml. + nested_function: str + volume: str + volume_file_path: str + volume_marker: str + + +def create_uc_function(workspace: Workspace, scratch_schema: str) -> UcFunction: + """Create the marker functions 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 MatrixError( + f"The UC function's MCP tool name would exceed {MAX_TOOL_NAME_LENGTH} characters: " + f"{exposed_tool_name!r}. Use a shorter catalog." + ) + # The nested function first: the outer one resolves it when created. + workspace.sql( + f"CREATE FUNCTION `{catalog}`.`{schema}`.`{NESTED_FUNCTION}`" + "(value STRING) RETURNS STRING " + "COMMENT 'Transitive Agent Bricks E2E marker; never declared in agent.toml' " + "RETURN concat('AGENTBRICKS_UC_OK:', value)" + ) + workspace.sql( + f"CREATE FUNCTION `{catalog}`.`{schema}`.`{FUNCTION}`" + "(value STRING) RETURNS STRING " + "COMMENT 'Deterministic Agent Bricks E2E marker tool' " + f"RETURN `{catalog}`.`{schema}`.`{NESTED_FUNCTION}`(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}", + nested_function=f"{scratch_schema}.{NESTED_FUNCTION}", + volume=f"{scratch_schema}.{VOLUME}", + volume_file_path=file_path, + volume_marker=marker, + ) diff --git a/integrations/agentbricks/tests/provisioning.py b/integrations/agentbricks/tests/provisioning.py index 4817e74b1..dcbd5bb67 100644 --- a/integrations/agentbricks/tests/provisioning.py +++ b/integrations/agentbricks/tests/provisioning.py @@ -1,53 +1,26 @@ -"""pytest plugin: the workspace resources tests ask for, as lazy session fixtures, and their teardown. +"""pytest plugin: the workspace requirements shared by every matrix, and the end-of-run cleanup. -A fixture exists per requirement (``catalog``, ``scratch_schema``, ``uc_function``) and is only -created when a selected test needs it, so a run that never asks for ``uc_function`` never deploys a -bundle or creates a function. Each fixture skips when the target workspace -cannot provide it or the operator did not configure it. +``catalog`` is the one fixture here; a matrix that needs more (the tools matrix's UC schema) builds +it from ``require_workspace`` in its own conftest, and tears it down in the test's own teardown. -Creation happens once per run across xdist workers (``shared_state``). Teardown happens once, in -``pytest_sessionfinish`` on the controller (the only process that sees every worker finish): it -destroys whatever the creators tracked and sweeps Apps a crashed test leaked. +Cleanup of leftovers happens once, in ``pytest_sessionfinish`` on the controller (the only process +that sees every worker finish): it sweeps what a crashed test or a failed per-test cleanup leaked, +found by the run's name prefix (Apps, source folders, schemas) and by the App identities the workers +recorded in ``apps.jsonl`` (Lakebase databases and roles, which outlive their App). """ from __future__ import annotations -import dataclasses -import json -import uuid -from collections.abc import Callable -from functools import partial -from pathlib import Path - import pytest -from common import ( - APP_PREFIX, - MatrixError, - RunConfig, - child_env, - log, - project_prefix, - run_command, -) -from run_context import base_run_config, build_target_workspace, started_run_config -from shared_state import SharedResources, Track +from common import APP_PREFIX, RunConfig, log, project_prefix, scratch_schema_prefix +from run_context import build_target_workspace, started_run_config from target_workspace import TargetWorkspace, TargetWorkspaceUnavailable -from workspace_client import Workspace, cleanup_app - -BUNDLE_TARGET = "nightly" -SCHEMA_RESOURCE = "scratch_schema" - - -@dataclasses.dataclass(frozen=True) -class UcFunction: - """The scratch UC objects the tools use.""" - - function: str - # Called by ``function`` and deliberately never declared in agent.toml. - nested_function: str - volume: str - volume_file_path: str - volume_marker: str +from workspace_client import ( + Workspace, + cleanup_app, + delete_lakebase_leftovers, + recorded_app_identities, +) def require_workspace(target_workspace: TargetWorkspace, what: str) -> Workspace: @@ -65,207 +38,106 @@ def catalog(request: pytest.FixtureRequest, target_workspace: TargetWorkspace) - return name -@pytest.fixture(scope="session") -def deploy_scratch_schema( - catalog: str, - target_workspace: TargetWorkspace, - shared_resources: SharedResources, - request: pytest.FixtureRequest, -) -> Callable[[Path], str]: - """Deploy a matrix's bundle once per run and return its ``catalog.schema``. - - Every bundle names its schema ``ab_e2e_``, so one run can deploy only one of them. - """ - run = base_run_config(request.config) - - def deploy(bundle_dir: Path) -> str: - value = shared_resources.get_or_create( - "scratch_schema", - lambda track: _deploy_bundle(target_workspace, run, catalog, bundle_dir, track), - ) - if value["bundle_dir"] != str(bundle_dir.resolve()): - raise MatrixError( - f"This run already deployed {value['bundle_dir']}, whose schema has the same name " - f"as the one in {bundle_dir}; run one bundle-backed matrix per invocation." - ) - return value["schema"] - - return deploy - - -@pytest.fixture(scope="session") -def uc_function( - scratch_schema: str, target_workspace: TargetWorkspace, shared_resources: SharedResources -) -> UcFunction: - """Marker functions and a marker volume in ``scratch_schema`` (defined by the matrix that needs one).""" - workspace = require_workspace(target_workspace, "a UC function") - value = shared_resources.get_or_create( - "uc_function", lambda track: _create_uc_function(workspace, scratch_schema, track) - ) - return UcFunction(**value) - - -# Creation - - -def _bundle_env(target_workspace: TargetWorkspace, catalog: str, run_id: str) -> dict[str, str]: - return child_env( - {**target_workspace.env, "BUNDLE_VAR_catalog_name": catalog, "BUNDLE_VAR_run_id": run_id} - ) - - -def _bundle( - target_workspace: TargetWorkspace, - bundle_dir: Path, - env: dict[str, str], - *args: str, - timeout: float, -) -> str: - argv = ["databricks", "bundle", *args, "-t", BUNDLE_TARGET, *target_workspace.cli_args] - return run_command(argv, cwd=bundle_dir, env=env, timeout=timeout) - - -def _deploy_bundle( - target_workspace: TargetWorkspace, run: RunConfig, catalog: str, bundle_dir: Path, track: Track -) -> dict[str, str]: - bundle_dir = bundle_dir.resolve() - env = _bundle_env(target_workspace, catalog, run.run_id) - # Tracked first: a deploy that fails midway still leaves a schema to destroy. - track(bundle_dir=str(bundle_dir), catalog=catalog, run_id=run.run_id) - _bundle(target_workspace, bundle_dir, env, "deploy", timeout=900) - summary = json.loads( - _bundle(target_workspace, bundle_dir, env, "summary", "--output", "json", timeout=120) - ) - resource = summary.get("resources", {}).get("schemas", {}).get(SCHEMA_RESOURCE) - if not isinstance(resource, dict): - raise MatrixError(f"bundle summary has no schemas.{SCHEMA_RESOURCE}: {summary}") - resolved_id = str(resource.get("id") or "") - if resolved_id.count(".") == 1: - schema = resolved_id - else: - schema = f"{resource.get('catalog_name') or catalog}.{resource.get('name')}" - if schema.count(".") != 1 or "None" in schema: - raise MatrixError(f"Could not resolve the deployed schema from {resource}.") - return {"schema": schema, "bundle_dir": str(bundle_dir)} - - -def _create_uc_function(workspace: Workspace, schema: str, track: Track) -> dict[str, str]: - catalog, _, schema_name = schema.partition(".") - suffix = uuid.uuid4().hex[:8] - nested_name = f"agentbricks_nested_{suffix}" - nested_function = f"{catalog}.{schema_name}.{nested_name}" - track(nested_function=nested_function) - workspace.sql( - f"CREATE OR REPLACE FUNCTION `{catalog}`.`{schema_name}`.`{nested_name}`" - "(value STRING) RETURNS STRING " - "COMMENT 'Transitive Agent Bricks E2E marker; never declared in agent.toml' " - "RETURN concat('AGENTBRICKS_UC_OK:', value)" - ) - # Leave room for catalog and schema in the 64-character MCP tool name. - function_name = f"ab_uc_{suffix}" - function = f"{catalog}.{schema_name}.{function_name}" - volume_name = f"agentbricks_volume_{suffix}" - volume = f"{catalog}.{schema_name}.{volume_name}" - track(volume=volume) - workspace.sql(f"CREATE VOLUME `{catalog}`.`{schema_name}`.`{volume_name}`") - marker = f"AGENTBRICKS_VOLUME_{uuid.uuid4().hex}" - file_path = f"/Volumes/{catalog}/{schema_name}/{volume_name}/marker.txt" - track(volume_file_path=file_path) - workspace.upload(file_path, marker.encode()) - exposed_tool_name = function.replace(".", "__") - if len(exposed_tool_name) > 64: - raise MatrixError( - "The UC function's MCP tool name would exceed 64 characters: " - f"{exposed_tool_name!r}. Use a shorter catalog or run id." - ) - track(function=function) - workspace.sql( - f"CREATE OR REPLACE FUNCTION `{catalog}`.`{schema_name}`.`{function_name}`" - "(value STRING) RETURNS STRING " - "COMMENT 'Deterministic Agent Bricks E2E marker tool' " - f"RETURN `{catalog}`.`{schema_name}`.`{nested_name}`(value)" - ) - return { - "function": function, - "nested_function": nested_function, - "volume": volume, - "volume_file_path": file_path, - "volume_marker": marker, - } - - -# Teardown - - @pytest.hookimpl(trylast=True) def pytest_sessionfinish(session: pytest.Session) -> None: config = session.config - # Workers must not tear down: the controller outlives them and sees the last one finish. + # Workers must not clean up: the controller outlives them and sees the last one finish. if hasattr(config, "workerinput"): return run = started_run_config(config) - if run is None: - return - tracked = SharedResources(run.output / "state").tracked() - if not tracked and not (run.output / "projects").exists(): + # The output dir exists once any test has run, so its absence means nothing was provisioned. + if run is None or not run.output.exists(): return if config.getoption("keep_resources"): - log("Resources retained (--keep-resources); skipping teardown.") + log("Resources retained (--keep-resources); skipping cleanup.") return try: target_workspace = build_target_workspace(config) except TargetWorkspaceUnavailable: return - if target_workspace.client is not None and _teardown( - target_workspace, target_workspace.client, run, tracked - ): - if session.exitstatus == 0: - session.exitstatus = pytest.ExitCode.TESTS_FAILED - - -def _teardown( - target_workspace: TargetWorkspace, - workspace: Workspace, - run: RunConfig, - tracked: dict[str, dict[str, str]], -) -> list[str]: - """Remove what the run created and return a description of each failure.""" + workspace = target_workspace.client + if workspace is None: + return failures: list[str] = [] - - def attempt(label: str, action: Callable[[], object]) -> None: + sweeps = ( + ("leaked Apps", lambda: _sweep_apps(workspace, run)), + ("leaked Lakebase databases and roles", lambda: _sweep_lakebase(workspace, run)), + ("leaked source folders", lambda: _sweep_deployment_sources(workspace, run.run_id)), + ( + "leaked schemas", + lambda: _sweep_schemas(workspace, config.getoption("catalog"), run.run_id), + ), + ) + for label, action in sweeps: try: - action() + failures.extend(action()) except Exception as exc: - log(f"cleanup warning | {label} | {exc}") failures.append(f"{label}: {exc}") + for failure in failures: + log(f"cleanup warning | {failure}") + if failures and session.exitstatus == 0: + session.exitstatus = pytest.ExitCode.TESTS_FAILED - attempt("leaked Apps", lambda: failures.extend(_sweep_apps(workspace, run.run_id))) - functions = tracked.get("uc_function", {}) - for key in ("function", "nested_function"): - if name := functions.get(key): - attempt(f"function {name}", partial(workspace.drop, "FUNCTION", name)) - if path := functions.get("volume_file_path"): - attempt(f"file {path}", partial(workspace.delete_file, path)) - if volume := functions.get("volume"): - attempt(f"volume {volume}", partial(workspace.drop, "VOLUME", volume)) - if bundle := tracked.get("scratch_schema"): - env = _bundle_env(target_workspace, bundle["catalog"], bundle["run_id"]) - bundle_dir = Path(bundle["bundle_dir"]) - attempt( - "bundle destroy", - lambda: _bundle( - target_workspace, bundle_dir, env, "destroy", "--auto-approve", timeout=900 - ), - ) - return failures + +def _app_prefix(run_id: str) -> str: + return f"{APP_PREFIX}{project_prefix(run_id)}" -def _sweep_apps(workspace: Workspace, run_id: str) -> list[str]: +def _sweep_apps(workspace: Workspace, run: RunConfig) -> list[str]: """Delete Apps whose per-test cleanup failed or never ran, found by the run's name prefix.""" - prefix = f"{APP_PREFIX}{project_prefix(run_id)}" + prefix = _app_prefix(run.run_id) + recorded = recorded_app_identities(run.output) failures: list[str] = [] for app in list(workspace.client.apps.list()): if app.name and app.name.startswith(prefix): log(f"sweeping leaked App {app.name}") - failures += cleanup_app(workspace, app.name, runtime_store=True) + failures += cleanup_app( + workspace, app.name, runtime_store=True, identity=recorded.get(app.name) + ) + return failures + + +def _sweep_lakebase(workspace: Workspace, run: RunConfig) -> list[str]: + """Delete runtime databases and roles of this run's Apps, whether or not the App still exists.""" + recorded = list(recorded_app_identities(run.output).values()) + failures: list[str] = [] + for branch in sorted({item.branch for item in recorded if item.branch}): + failures += delete_lakebase_leftovers( + workspace, + 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 + ], + database_ids=[ + item.database_id for item in recorded if item.branch == branch and item.database_id + ], + ) + return failures + + +def _sweep_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"sweeping deployment source {path}") + try: + workspace.delete_workspace_path(path) + except Exception as exc: + failures.append(f"deployment source {path}: {exc}") + return failures + + +def _sweep_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"sweeping 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/run_context.py b/integrations/agentbricks/tests/run_context.py index aedefbbef..09c67da7e 100644 --- a/integrations/agentbricks/tests/run_context.py +++ b/integrations/agentbricks/tests/run_context.py @@ -16,7 +16,6 @@ import pytest from common import MatrixError, RunConfig, log, run_command -from shared_state import SharedResources from target_workspace import ( WORKSPACE_KINDS, TargetWorkspace, @@ -48,7 +47,7 @@ def pytest_addoption(parser: pytest.Parser) -> None: group.addoption( "--catalog", default=os.environ.get("AGENTBRICKS_E2E_CATALOG", "main"), - help="Existing catalog for scratch schemas (AGENTBRICKS_E2E_CATALOG, default main).", + help="Existing catalog for temporary test schemas (AGENTBRICKS_E2E_CATALOG, default main).", ) group.addoption("--bridge-sha", help="Immutable bridge commit for generated App dependencies.") group.addoption( @@ -63,7 +62,9 @@ def pytest_addoption(parser: pytest.Parser) -> None: help="Skip per-App USE CATALOG grants because catalog access is pre-provisioned.", ) group.addoption( - "--keep-resources", action="store_true", help="Skip teardown of provisioned resources." + "--keep-resources", + action="store_true", + help="Skip the end-of-run sweep of leaked Apps and temporary schemas.", ) @@ -126,27 +127,16 @@ def target_workspace(request: pytest.FixtureRequest) -> TargetWorkspace: @pytest.fixture(scope="session") -def shared_resources(request: pytest.FixtureRequest) -> SharedResources: - return SharedResources(base_run_config(request.config).output / "state") - - -@pytest.fixture(scope="session") -def run_config(request: pytest.FixtureRequest, shared_resources: SharedResources) -> RunConfig: +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 or run.bridge_sha: return run - - def build(_track: object) -> dict[str, Any]: - dist = run.output / "dist" - run_command( - ["uv", "build", "--wheel", "--out-dir", str(dist)], cwd=PACKAGE_DIR, timeout=600 - ) - wheels = sorted(dist.glob("*.whl")) - if not wheels: - raise MatrixError("uv build produced no wheel") - return {"path": str(wheels[-1])} - - built = shared_resources.get_or_create("wheel", build) - return dataclasses.replace(run, wheel=pathlib.Path(built["path"])) + # 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") + run_command(["uv", "build", "--wheel", "--out-dir", str(dist)], cwd=PACKAGE_DIR, timeout=600) + wheels = sorted(dist.glob("*.whl")) + if not wheels: + raise MatrixError("uv build produced no wheel") + return dataclasses.replace(run, wheel=wheels[-1]) diff --git a/integrations/agentbricks/tests/shared_state.py b/integrations/agentbricks/tests/shared_state.py deleted file mode 100644 index 952e529ab..000000000 --- a/integrations/agentbricks/tests/shared_state.py +++ /dev/null @@ -1,74 +0,0 @@ -"""Create-once, share-everywhere resources for a pytest run, including all xdist workers. - -Each resource is one JSON file in a directory every worker can reach, guarded by a per-resource -file lock (the pattern pytest-xdist documents). The first process to ask creates it while the others -wait; later processes read the cached value. Whatever a creator ``track``s lands in the same file, -which is how the controller finds what to destroy after every worker has finished. -""" - -from __future__ import annotations - -import contextlib -import fcntl -import json -import os -import pathlib -from collections.abc import Callable, Iterator -from typing import Any - -from common import MatrixError - -Track = Callable[..., None] - - -@contextlib.contextmanager -def _locked(path: pathlib.Path) -> Iterator[None]: - descriptor = os.open(path, os.O_CREAT | os.O_RDWR) - try: - fcntl.flock(descriptor, fcntl.LOCK_EX) - yield - finally: - os.close(descriptor) - - -def _read(path: pathlib.Path) -> dict[str, Any]: - return json.loads(path.read_text(encoding="utf-8")) if path.exists() else {} - - -class SharedResources: - def __init__(self, directory: pathlib.Path): - self._directory = directory - - def get_or_create(self, key: str, create: Callable[[Track], dict[str, Any]]) -> dict[str, Any]: - """The cached value for ``key``, calling ``create(track)`` in exactly one process. - - ``track(**fields)`` records what ``create`` has made so far, before it can fail, for - ``tracked``. A failure is cached too, so other workers fail fast instead of retrying a - deploy that takes minutes. - """ - self._directory.mkdir(parents=True, exist_ok=True) - path = self._directory / f"{key}.json" - with _locked(self._directory / f"{key}.lock"): - entry = _read(path) - if "error" in entry: - raise MatrixError(f"{key} setup already failed: {entry['error']}") - if "value" in entry: - return entry["value"] - - def track(**fields: Any) -> None: - entry.setdefault("cleanup", {}).update(fields) - path.write_text(json.dumps(entry), encoding="utf-8") - - try: - entry["value"] = create(track) - except Exception as exc: - entry["error"] = str(exc) - path.write_text(json.dumps(entry), encoding="utf-8") - raise - path.write_text(json.dumps(entry), encoding="utf-8") - return entry["value"] - - def tracked(self) -> dict[str, dict[str, Any]]: - """What each resource's creator recorded, for teardown.""" - entries = {path.stem: _read(path) for path in sorted(self._directory.glob("*.json"))} - return {key: entry["cleanup"] for key, entry in entries.items() if entry.get("cleanup")} diff --git a/integrations/agentbricks/tests/workspace_client.py b/integrations/agentbricks/tests/workspace_client.py index 096875506..095594250 100644 --- a/integrations/agentbricks/tests/workspace_client.py +++ b/integrations/agentbricks/tests/workspace_client.py @@ -5,23 +5,72 @@ from __future__ import annotations +import dataclasses import datetime as dt import io +import json import pathlib import re import subprocess import time -from collections.abc import Callable, Sequence +from collections.abc import Callable, Collection, Sequence from typing import Any, cast from common import TOOL_RESOURCE_PREFIX, MatrixError, log, now from databricks.sdk import WorkspaceClient from databricks.sdk.errors import DatabricksError, NotFound -from databricks.sdk.service.catalog import SecurableType +from databricks.sdk.service.catalog import SecurableType, VolumeType from databricks.sdk.service.sql import ExecuteStatementRequestOnWaitTimeout, State, StatementState GrantTuple = tuple[str, str, str, str] +_BRANCH = re.compile(r"projects/[^/]+/branches/[^/]+") +_RUNTIME_DATABASE_PREFIX = "runtime-" +_APPS_RECORD = "apps.jsonl" + + +@dataclasses.dataclass(frozen=True) +class AppIdentity: + """What cleanup needs to find an App's Lakebase leftovers once the App and its Runtime Store are gone.""" + + service_principal: str | None = None + branch: str | None = None + database_id: 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, + self.database_id or other.database_id, + ) + + +def record_app_identity(output: pathlib.Path, app: str, identity: AppIdentity) -> None: + """Append one line for the controller's end-of-run sweep, which cannot see worker memory.""" + line = json.dumps({"app": app, **dataclasses.asdict(identity)}) + "\n" + # 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"), data.get("database_id") + ) + identities[data["app"]] = identities.get(data["app"], AppIdentity()).merged(identity) + except (ValueError, KeyError, TypeError, AttributeError): + continue + return identities + class Workspace: def __init__( @@ -40,6 +89,7 @@ def __init__( self._app_auth_client: WorkspaceClient | None = None self._app_auth_checked = False self._warehouse_started = False + self._user_name: str | None = None # App auth: Databricks Apps /api routes need OAuth, so a PAT is rejected. @@ -137,19 +187,41 @@ def sql(self, statement: str, *, timeout: float = 600) -> None: if not response.status or response.status.state != StatementState.SUCCEEDED: raise MatrixError(f"SQL failed: {response.as_dict()}") - def drop(self, kind: str, full_name: str) -> None: - catalog, schema, name = full_name.split(".") - self.sql(f"DROP {kind} IF EXISTS `{catalog}`.`{schema}`.`{name}`") - def upload(self, path: str, data: bytes) -> None: self.client.files.upload(path, io.BytesIO(data), overwrite=True) log(f"# uploaded {path}") - def delete_file(self, path: str) -> None: + # 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 sweeper 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.files.delete(path) + self.client.schemas.delete(full_name, force=True) except NotFound: - pass + 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 @@ -329,43 +401,162 @@ def delete_runtime_store(self, app_name: str) -> None: except NotFound: pass - def app_role_target(self, app_name: str) -> str | None: - """The Lakebase role of the App's principal, only if ownership is verified end to end.""" + def app_identity(self, app_name: str) -> AppIdentity: + """The App's principal and Lakebase location as far as they still exist; never raises.""" + principal: str | None = None + branch: str | None = None + database_id: str | None = None try: principal = self.app(app_name).get("service_principal_client_id") - store = self.runtime_store(app_name) - owner = store.get("owner", {}).get("app", {}) - branch = store.get("storage_backend", {}).get("lakebase", {}).get("branch") + except NotFound: + pass + except Exception as exc: + log(f"cleanup warning | principal lookup for {app_name} | {exc}") + try: + lakebase = self.runtime_store(app_name).get("storage_backend", {}).get("lakebase", {}) + candidate = lakebase.get("branch") + if isinstance(candidate, str) and _BRANCH.fullmatch(candidate): + branch = candidate + database_id = lakebase.get("database_id") + except NotFound: + pass + except Exception as exc: + log(f"cleanup warning | Runtime Store lookup for {app_name} | {exc}") + return AppIdentity(principal, branch, database_id) + + def runtime_databases( + self, branch: str, principals: Collection[str], database_ids: Collection[str] + ) -> list[str]: + """Resource names of the branch's ``runtime-`` databases owned by these principals. + + Database ids embed a truncated App name plus a UUID, so ownership is read from the owner + role instead, which is named after the principal. ``database_ids`` are matched as given. + """ + wanted = [principal for principal in principals if principal] + names = [] + for database in self.client.postgres.list_databases(parent=branch): + database_id = (database.name or "").rsplit("/", 1)[-1] + owner = ( + (database.status.role if database.status else None) + or (database.spec.role if database.spec else None) + or "" + ) if ( - not principal - or store.get("name") != f"runtime-stores/{app_name}" - or owner.get("name") != app_name - or owner.get("service_principal_id") != principal - or not isinstance(branch, str) - or not re.fullmatch(r"projects/[^/]+/branches/[^/]+", branch) + database.name + and database_id.startswith(_RUNTIME_DATABASE_PREFIX) + and ( + database_id in database_ids + or any(owner.rsplit("/", 1)[-1].endswith(item) for item in wanted) + ) ): - log(f"cleanup warning | Lakebase role for {app_name} | ownership not verified") - return None - for role in self.client.postgres.list_roles(parent=branch): - data = role.as_dict() - status = data.get("status", {}) - name = data.get("name") - if ( - status.get("postgres_role") == principal - and status.get("identity_type") == "SERVICE_PRINCIPAL" - and isinstance(name, str) - and name.startswith(f"{branch}/roles/") - ): - return name + names.append(database.name) + return names + + 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_database(self, name: str) -> None: + try: + _wait(self.client.postgres.delete_database(name=name)) + except NotFound: + pass + + 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 MatrixError("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 delete_deployment_source(self, app_name: str) -> None: + """Delete the App's synced source folder; a missing folder counts as deleted.""" + self.delete_workspace_path(f"{self._deployments_dir()}/{app_name}") + + 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() + + +def _delete_each( + kind: str, find: Callable[[], list[str]], delete: Callable[[str], None] +) -> list[str]: + """Delete everything ``find`` returns, continuing past failures; returns one line per failure.""" + try: + names = find() + except Exception as exc: + return [f"{kind} lookup: {exc}"] + failures = [] + for name in names: + try: + delete(name) except Exception as exc: - log(f"cleanup warning | Lakebase role lookup for {app_name} | {exc}") - return None + failures.append(f"{kind} {name}: {exc}") + return failures - def delete_role(self, role_name: str) -> None: - operation = self.client.postgres.delete_role(name=role_name) - wait = getattr(operation, "wait", None) - if callable(wait): - wait() + +def delete_lakebase_leftovers( + workspace: Workspace, + branch: str, + *, + principals: Collection[str], + database_ids: Collection[str] = (), +) -> list[str]: + """Delete runtime databases and then the principals' roles on one branch. + + A role that owns a database cannot be deleted, so databases go first. Only ``runtime-`` + databases are touched, and both kinds are matched by principal UUID, never by name pattern. + """ + return _delete_each( + "runtime database", + lambda: workspace.runtime_databases(branch, principals, database_ids), + workspace.delete_database, + ) + _delete_each( + "Lakebase role", + lambda: workspace.roles_of(branch, principals), + workspace.delete_role, + ) def cleanup_app( @@ -377,46 +568,53 @@ def cleanup_app( session_store: str | None = None, has_app: bool = True, runtime_store: bool = False, + identity: AppIdentity | None = None, ) -> list[str]: - """Delete one project's stores and, if it has an App, its runtime store, App and Lakebase role. + """Delete one project's stores and, if it has an App, everything the App's deploy created. ``delete_store(kind, store)`` is the agentbricks CLI's store delete, which has no SDK surface. - Returns a description of each failure; an empty list means everything was deleted. + ``identity`` is what was recorded at deploy time; it fills in whatever the live lookup misses. + Every step runs regardless of earlier failures except those that need the App gone. Returns a + description of each failure; an empty list means everything was deleted. """ failures: list[str] = [] - # Looked up before anything is deleted, since it needs the App and its runtime store. - role_target = workspace.app_role_target(name) if has_app and runtime_store else None - for kind, store in (("memory", memory_store), ("sessions", session_store)): - if not (store and delete_store): - continue + + def attempt(label: str, action: Callable[..., object], *args: Any) -> bool: try: - result = delete_store(kind, store) - if result.returncode != 0: - failures.append(f"{kind} store {store}: {(result.stderr or result.stdout).strip()}") + action(*args) except Exception as exc: - failures.append(f"{kind} store {store}: {exc}") + failures.append(f"{label}: {exc}") + return False + return True + + def delete_cli_store(kind: str, store: str) -> None: + assert delete_store is not None + result = delete_store(kind, store) + if result.returncode != 0: + raise MatrixError((result.stderr or result.stdout).strip()) + + # Read before any delete: the principal and branch come from the App and its Runtime Store. + known = (identity or AppIdentity()).merged( + workspace.app_identity(name) if has_app else AppIdentity() + ) + for kind, store in (("memory", memory_store), ("sessions", session_store)): + if store and delete_store: + attempt(f"{kind} store {store}", delete_cli_store, kind, store) if not has_app: return failures - runtime_store_deleted = False if runtime_store: - # The Runtime Store owns a dedicated Lakebase database; it must be deleted before the - # role, or the role delete fails on database ownership. - try: - workspace.delete_runtime_store(name) - runtime_store_deleted = True - except Exception as exc: - failures.append(f"runtime store {name}: {exc}") - try: - workspace.delete_app(name) - except Exception as exc: - failures.append(f"app {name}: {exc}") - return failures - if role_target and not failures and runtime_store_deleted: - try: - workspace.delete_role(role_target) - except Exception as exc: - failures.append(f"Lakebase role {role_target}: {exc}") + attempt(f"runtime store {name}", workspace.delete_runtime_store, name) + app_deleted = attempt(f"app {name}", workspace.delete_app, name) + if app_deleted: + attempt(f"deployment source {name}", workspace.delete_deployment_source, name) + if known.branch: + failures += delete_lakebase_leftovers( + workspace, + known.branch, + principals=[known.service_principal or ""], + database_ids=[known.database_id or ""], + ) return failures From 4d3b3b1213b1f553d3f4d6a090e4ad5954d04e7f Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Mon, 12 Oct 2026 22:02:48 +0000 Subject: [PATCH 08/45] [agentbricks] rename e2e matrices to suites Co-authored-by: Isaac --- .../agentbricks/tests/agentbricks_cli.py | 34 ++++++------ integrations/agentbricks/tests/common.py | 8 +-- integrations/agentbricks/tests/conftest.py | 6 +-- .../agentbricks/tests/provisioning.py | 4 +- integrations/agentbricks/tests/run_context.py | 6 +-- .../memory_store/test_memory_store.py | 0 .../session_store/test_session_store.py | 0 .../{matrices => suites}/tools/conftest.py | 6 +-- .../{matrices => suites}/tools/test_deploy.py | 0 .../{matrices => suites}/tools/test_dev.py | 0 .../{matrices => suites}/tools/test_grants.py | 0 .../tools/test_selection.py | 0 .../tests/{matrices => suites}/tools/tools.py | 0 .../{matrices => suites}/tools/uc_objects.py | 6 +-- .../tracing/test_tracing.py | 0 .../agentbricks/tests/workspace_client.py | 54 +++++++++---------- 16 files changed, 59 insertions(+), 65 deletions(-) rename integrations/agentbricks/tests/{matrices => suites}/memory_store/test_memory_store.py (100%) rename integrations/agentbricks/tests/{matrices => suites}/session_store/test_session_store.py (100%) rename integrations/agentbricks/tests/{matrices => suites}/tools/conftest.py (93%) rename integrations/agentbricks/tests/{matrices => suites}/tools/test_deploy.py (100%) rename integrations/agentbricks/tests/{matrices => suites}/tools/test_dev.py (100%) rename integrations/agentbricks/tests/{matrices => suites}/tools/test_grants.py (100%) rename integrations/agentbricks/tests/{matrices => suites}/tools/test_selection.py (100%) rename integrations/agentbricks/tests/{matrices => suites}/tools/tools.py (100%) rename integrations/agentbricks/tests/{matrices => suites}/tools/uc_objects.py (93%) rename integrations/agentbricks/tests/{matrices => suites}/tracing/test_tracing.py (100%) diff --git a/integrations/agentbricks/tests/agentbricks_cli.py b/integrations/agentbricks/tests/agentbricks_cli.py index c1f01a723..109e4e027 100644 --- a/integrations/agentbricks/tests/agentbricks_cli.py +++ b/integrations/agentbricks/tests/agentbricks_cli.py @@ -2,7 +2,7 @@ Every call here is a real ``agentbricks`` subprocess, so a regression in the CLI fails the test that makes it. The only non-subprocess part is ``Agent.invoke``, which POSTs to a running agent. -Nothing in this module knows about any particular matrix; tests say what to bind. It records the +Nothing in this module knows about any particular suite; tests say what to bind. It records the projects it creates in ``projects`` but deletes nothing: the ``agentbricks_cli`` fixture does. """ @@ -34,7 +34,7 @@ APP_PREFIX, E2E_MODEL, FRAMEWORKS, - MatrixError, + E2EError, RunConfig, last_lines, last_nonempty_line, @@ -125,7 +125,7 @@ def invoke(self, prompt: str, label: str = "invoke") -> str: log(f"attempt {attempt}/3 | {label} | {exc}") if attempt < 3: time.sleep(15) - raise MatrixError(f"{label} failed after 3 attempts: {last}") + raise E2EError(f"{label} failed after 3 attempts: {last}") def warm_up(self, deadline_seconds: float = 180.0) -> None: """Invoke until the App answers once. A new App's session store can drop the first calls.""" @@ -134,7 +134,7 @@ def warm_up(self, deadline_seconds: float = 180.0) -> None: try: self.invoke("Reply with the single word OK.", label="warm-up") return - except MatrixError: + except E2EError: if time.monotonic() - started > deadline_seconds: raise @@ -178,7 +178,7 @@ def cli( if result.stderr.strip(): log(result.stderr) if check and result.returncode != 0: - raise MatrixError( + raise E2EError( f"Command failed ({result.returncode}): {shlex.join(argv)}\n" f"{result.stderr or result.stdout}" ) @@ -217,14 +217,14 @@ def _run_long(self, label: str, argv: Sequence[str], *, timeout: float) -> None: elapsed = time.monotonic() - started if elapsed >= timeout: terminate_process_group(process, grace=0) - raise MatrixError(f"{label} timed out after {timeout:.0f}s; log: {log_path}") + raise E2EError(f"{label} timed out after {timeout:.0f}s; log: {log_path}") if elapsed >= next_tick: log(f"tick {now():%H:%M} | {label} | running | {last_nonempty_line(log_path)}") next_tick += 60.0 time.sleep(2) log(log_path.read_text(encoding="utf-8", errors="replace")) if process.returncode != 0: - raise MatrixError(f"{label} failed ({process.returncode}); log: {log_path}") + raise E2EError(f"{label} failed ({process.returncode}); log: {log_path}") log(f"tick {now():%H:%M} | {label} | success") def _next(self) -> int: @@ -424,7 +424,7 @@ def deploy(self, project: Project, app: str | None = None) -> Agent: """`agentbricks deploy` the project; deploying again updates the same App.""" workspace = self.target_workspace.client if workspace is None: - raise MatrixError( + raise E2EError( f"deploy needs a live workspace; got --workspace {self.target_workspace.name}" ) workspace.check_app_auth() @@ -447,7 +447,7 @@ def deploy(self, project: Project, app: str | None = None) -> Agent: self._record_identity(project, name, workspace) url = str(deployed.get("url") or "").rstrip("/") if not url: - raise MatrixError(f"App {name} has no URL: {deployed}") + raise E2EError(f"App {name} has no URL: {deployed}") agent = Agent( url, lambda: workspace.app_headers, @@ -473,9 +473,7 @@ def _wait_for_local( next_tick = 60.0 while True: if process.poll() is not None: - raise MatrixError( - f"{label} exited {process.returncode}: {last_lines(log_path, 30)}" - ) + raise E2EError(f"{label} exited {process.returncode}: {last_lines(log_path, 30)}") try: with urllib.request.urlopen(f"http://127.0.0.1:{port}/", timeout=5): return @@ -486,7 +484,7 @@ def _wait_for_local( pass elapsed = time.monotonic() - started if elapsed > 1200: - raise MatrixError(f"{label} did not become reachable: {last_lines(log_path, 30)}") + raise E2EError(f"{label} did not become reachable: {last_lines(log_path, 30)}") if elapsed >= next_tick: log(f"tick {now():%H:%M} | {label} | starting | {last_nonempty_line(log_path)}") next_tick += 60 @@ -512,7 +510,7 @@ def _agentbricks_bin() -> str: sibling = pathlib.Path(sys.executable).parent / "agentbricks" found = str(sibling) if sibling.exists() else shutil.which("agentbricks") if found is None: - raise MatrixError("The agentbricks CLI is not installed in the test environment.") + raise E2EError("The agentbricks CLI is not installed in the test environment.") return found @@ -553,7 +551,7 @@ def _monitored( elapsed = time.monotonic() - started log(f"tick {now():%H:%M} | {label} | running | {elapsed:.0f}s") if elapsed >= timeout: - raise MatrixError(f"{label} timed out after {timeout:.0f}s") from None + raise E2EError(f"{label} timed out after {timeout:.0f}s") from None def _http_json(url: str, body: dict[str, Any], headers: dict[str, str]) -> dict[str, Any]: @@ -568,11 +566,11 @@ def _http_json(url: str, body: dict[str, Any], headers: dict[str, str]) -> dict[ payload = response.read().decode() except urllib.error.HTTPError as exc: detail = exc.read().decode(errors="replace") - raise MatrixError(f"HTTP {exc.code} from {url}: {detail}") from exc + raise E2EError(f"HTTP {exc.code} from {url}: {detail}") from exc try: value = json.loads(payload) except json.JSONDecodeError as exc: - raise MatrixError(f"Invalid JSON from {url}: {payload[:2000]}") from exc + raise E2EError(f"Invalid JSON from {url}: {payload[:2000]}") from exc if not isinstance(value, dict): - raise MatrixError(f"Expected object response from {url}, got {type(value).__name__}") + raise E2EError(f"Expected object response from {url}, got {type(value).__name__}") return value diff --git a/integrations/agentbricks/tests/common.py b/integrations/agentbricks/tests/common.py index 71d8f1ef4..19e2d36e4 100644 --- a/integrations/agentbricks/tests/common.py +++ b/integrations/agentbricks/tests/common.py @@ -1,4 +1,4 @@ -"""Constants, ``RunConfig`` and logging/subprocess helpers shared by the matrix tests. +"""Constants, ``RunConfig`` and logging/subprocess helpers shared by the E2E suites. ``run_context.py`` builds one ``RunConfig`` per pytest run; fixtures and tests read it through the ``run_config`` fixture. @@ -21,7 +21,7 @@ APP_PREFIX = "agent-bricks-" -class MatrixError(RuntimeError): +class E2EError(RuntimeError): """A reproducible setup or execution failure.""" @@ -69,7 +69,7 @@ def run_command( env: Mapping[str, str] | None = None, timeout: float = 300, ) -> str: - """Run a command, echo its output, and return stdout; a non-zero exit is a ``MatrixError``.""" + """Run a command, echo its output, and return stdout; a non-zero exit is a ``E2EError``.""" log_command(argv, cwd) result = subprocess.run( list(argv), @@ -84,7 +84,7 @@ def run_command( if text.strip(): log(text) if result.returncode != 0: - raise MatrixError( + raise E2EError( f"Command failed ({result.returncode}): {shlex.join(argv)}\n" f"{result.stderr or result.stdout}" ) diff --git a/integrations/agentbricks/tests/conftest.py b/integrations/agentbricks/tests/conftest.py index 4a355babb..91f16a398 100644 --- a/integrations/agentbricks/tests/conftest.py +++ b/integrations/agentbricks/tests/conftest.py @@ -1,9 +1,9 @@ -"""Fixtures shared by every matrix under ``matrices/``. +"""Fixtures shared by every suite under ``suites/``. ``run_context`` (options, target workspace, run identity) and ``provisioning`` (the catalog -requirement and the end-of-run cleanup) are plugins. One command runs a matrix, from ``integrations/agentbricks``: +requirement and the end-of-run cleanup) are plugins. One command runs a suite, from ``integrations/agentbricks``: - uv run pytest tests/matrices/tools -n 6 --dist loadgroup \ + uv run pytest tests/suites/tools -n 6 --dist loadgroup \ --databricks-profile ``--workspace fake`` runs against no workspace; whatever needs one is skipped. diff --git a/integrations/agentbricks/tests/provisioning.py b/integrations/agentbricks/tests/provisioning.py index dcbd5bb67..26c63fcdb 100644 --- a/integrations/agentbricks/tests/provisioning.py +++ b/integrations/agentbricks/tests/provisioning.py @@ -1,6 +1,6 @@ -"""pytest plugin: the workspace requirements shared by every matrix, and the end-of-run cleanup. +"""pytest plugin: the workspace requirements shared by every suite, and the end-of-run cleanup. -``catalog`` is the one fixture here; a matrix that needs more (the tools matrix's UC schema) builds +``catalog`` is the one fixture here; a suite that needs more (the tools suite's UC schema) builds it from ``require_workspace`` in its own conftest, and tears it down in the test's own teardown. Cleanup of leftovers happens once, in ``pytest_sessionfinish`` on the controller (the only process diff --git a/integrations/agentbricks/tests/run_context.py b/integrations/agentbricks/tests/run_context.py index 09c67da7e..bd4368e53 100644 --- a/integrations/agentbricks/tests/run_context.py +++ b/integrations/agentbricks/tests/run_context.py @@ -15,7 +15,7 @@ from typing import Any import pytest -from common import MatrixError, RunConfig, log, run_command +from common import E2EError, RunConfig, log, run_command from target_workspace import ( WORKSPACE_KINDS, TargetWorkspace, @@ -29,7 +29,7 @@ def pytest_addoption(parser: pytest.Parser) -> None: - group = parser.getgroup("agentbricks", "Agent Bricks E2E matrices") + group = parser.getgroup("agentbricks", "Agent Bricks E2E suites") group.addoption( "--workspace", choices=WORKSPACE_KINDS, @@ -138,5 +138,5 @@ def run_config(request: pytest.FixtureRequest) -> RunConfig: run_command(["uv", "build", "--wheel", "--out-dir", str(dist)], cwd=PACKAGE_DIR, timeout=600) wheels = sorted(dist.glob("*.whl")) if not wheels: - raise MatrixError("uv build produced no wheel") + raise E2EError("uv build produced no wheel") return dataclasses.replace(run, wheel=wheels[-1]) diff --git a/integrations/agentbricks/tests/matrices/memory_store/test_memory_store.py b/integrations/agentbricks/tests/suites/memory_store/test_memory_store.py similarity index 100% rename from integrations/agentbricks/tests/matrices/memory_store/test_memory_store.py rename to integrations/agentbricks/tests/suites/memory_store/test_memory_store.py diff --git a/integrations/agentbricks/tests/matrices/session_store/test_session_store.py b/integrations/agentbricks/tests/suites/session_store/test_session_store.py similarity index 100% rename from integrations/agentbricks/tests/matrices/session_store/test_session_store.py rename to integrations/agentbricks/tests/suites/session_store/test_session_store.py diff --git a/integrations/agentbricks/tests/matrices/tools/conftest.py b/integrations/agentbricks/tests/suites/tools/conftest.py similarity index 93% rename from integrations/agentbricks/tests/matrices/tools/conftest.py rename to integrations/agentbricks/tests/suites/tools/conftest.py index e4f33e4af..135eb2639 100644 --- a/integrations/agentbricks/tests/matrices/tools/conftest.py +++ b/integrations/agentbricks/tests/suites/tools/conftest.py @@ -1,4 +1,4 @@ -"""Tools matrix: every test runs once per authoring path, over the tools the target workspace can support. +"""Tools suite: every test runs once per authoring path, over the tools the target workspace can support. Each test gets its own temporary UC schema (``ab_e2e__``) holding the function and volume its tools use, and deletes it, grants included, in its own teardown. The controller sweeps @@ -16,7 +16,7 @@ import pytest from common import ( AUTHORING_PATHS, - MatrixError, + E2EError, RunConfig, log, now, @@ -64,7 +64,7 @@ def genie_space(target_workspace: TargetWorkspace) -> str: if not space_id: pytest.skip("no Genie space; set AGENTBRICKS_E2E_GENIE_SPACE_ID") if re.fullmatch(r"[0-9a-f]{32}", space_id) is None: - raise MatrixError(f"Genie space id {space_id!r} is not 32 lowercase hex characters.") + 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 diff --git a/integrations/agentbricks/tests/matrices/tools/test_deploy.py b/integrations/agentbricks/tests/suites/tools/test_deploy.py similarity index 100% rename from integrations/agentbricks/tests/matrices/tools/test_deploy.py rename to integrations/agentbricks/tests/suites/tools/test_deploy.py diff --git a/integrations/agentbricks/tests/matrices/tools/test_dev.py b/integrations/agentbricks/tests/suites/tools/test_dev.py similarity index 100% rename from integrations/agentbricks/tests/matrices/tools/test_dev.py rename to integrations/agentbricks/tests/suites/tools/test_dev.py diff --git a/integrations/agentbricks/tests/matrices/tools/test_grants.py b/integrations/agentbricks/tests/suites/tools/test_grants.py similarity index 100% rename from integrations/agentbricks/tests/matrices/tools/test_grants.py rename to integrations/agentbricks/tests/suites/tools/test_grants.py diff --git a/integrations/agentbricks/tests/matrices/tools/test_selection.py b/integrations/agentbricks/tests/suites/tools/test_selection.py similarity index 100% rename from integrations/agentbricks/tests/matrices/tools/test_selection.py rename to integrations/agentbricks/tests/suites/tools/test_selection.py diff --git a/integrations/agentbricks/tests/matrices/tools/tools.py b/integrations/agentbricks/tests/suites/tools/tools.py similarity index 100% rename from integrations/agentbricks/tests/matrices/tools/tools.py rename to integrations/agentbricks/tests/suites/tools/tools.py diff --git a/integrations/agentbricks/tests/matrices/tools/uc_objects.py b/integrations/agentbricks/tests/suites/tools/uc_objects.py similarity index 93% rename from integrations/agentbricks/tests/matrices/tools/uc_objects.py rename to integrations/agentbricks/tests/suites/tools/uc_objects.py index bfd9234c9..9f781f96a 100644 --- a/integrations/agentbricks/tests/matrices/tools/uc_objects.py +++ b/integrations/agentbricks/tests/suites/tools/uc_objects.py @@ -1,11 +1,11 @@ -"""The UC objects the tools matrix uses, created fresh in each test's temporary schema.""" +"""The UC objects the tools suite uses, created fresh in each test's temporary schema.""" from __future__ import annotations import dataclasses import uuid -from common import MatrixError +from common import E2EError from workspace_client import Workspace NESTED_FUNCTION = "ab_nested_marker" @@ -32,7 +32,7 @@ def create_uc_function(workspace: Workspace, scratch_schema: str) -> UcFunction: catalog, _, schema = scratch_schema.partition(".") exposed_tool_name = f"{scratch_schema}.{FUNCTION}".replace(".", "__") if len(exposed_tool_name) > MAX_TOOL_NAME_LENGTH: - raise MatrixError( + 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." ) diff --git a/integrations/agentbricks/tests/matrices/tracing/test_tracing.py b/integrations/agentbricks/tests/suites/tracing/test_tracing.py similarity index 100% rename from integrations/agentbricks/tests/matrices/tracing/test_tracing.py rename to integrations/agentbricks/tests/suites/tracing/test_tracing.py diff --git a/integrations/agentbricks/tests/workspace_client.py b/integrations/agentbricks/tests/workspace_client.py index 095594250..507a607f4 100644 --- a/integrations/agentbricks/tests/workspace_client.py +++ b/integrations/agentbricks/tests/workspace_client.py @@ -1,4 +1,4 @@ -"""Databricks SDK helpers for the matrix tests and fixtures: the workspace, which is not under test. +"""Databricks SDK helpers for the E2E suites and fixtures: the workspace, which is not under test. Everything here goes through the SDK except ``app_logs``, since the SDK has no Apps log API. """ @@ -16,7 +16,7 @@ from collections.abc import Callable, Collection, Sequence from typing import Any, cast -from common import TOOL_RESOURCE_PREFIX, MatrixError, log, now +from common import TOOL_RESOURCE_PREFIX, E2EError, log, now from databricks.sdk import WorkspaceClient from databricks.sdk.errors import DatabricksError, NotFound from databricks.sdk.service.catalog import SecurableType, VolumeType @@ -109,7 +109,7 @@ def check_app_auth(self) -> None: if self._app_auth_checked: return if self._app_client().config.auth_type == "pat": - raise MatrixError( + raise E2EError( f"{self._app_auth_label()} uses a PAT. " "Databricks Apps /api routes require OAuth; run `databricks auth login` " "for a profile on the same workspace." @@ -121,7 +121,7 @@ def app_headers(self) -> dict[str, str]: # Resolved per call so OAuth tokens refresh during a long run. authorization = self._app_client().config.authenticate().get("Authorization") if not authorization: - raise MatrixError(f"Could not resolve credentials from {self._app_auth_label()}.") + raise E2EError(f"Could not resolve credentials from {self._app_auth_label()}.") return {"Authorization": authorization} # Existing resources the suite depends on @@ -130,13 +130,13 @@ def require_catalog(self, name: str) -> None: try: self.client.catalogs.get(name) except DatabricksError as exc: - raise MatrixError(f"Catalog {name!r} is not accessible: {exc}") from 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 MatrixError(f"Genie space {space_id!r} is not accessible: {exc}") from exc + raise E2EError(f"Genie space {space_id!r} is not accessible: {exc}") from exc # SQL, warehouse, files @@ -146,7 +146,7 @@ def start_warehouse(self, override: str | None = None) -> str: else: warehouses = list(self.client.warehouses.list()) if not warehouses: - raise MatrixError("The workspace has no SQL warehouse available for UC setup.") + 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] ) @@ -178,14 +178,14 @@ def sql(self, statement: str, *, timeout: float = 600) -> None: StatementState.RUNNING, }: if not statement_id: - raise MatrixError(f"SQL response has no statement_id: {response}") + raise E2EError(f"SQL response has no statement_id: {response}") if timeout <= 0: - raise MatrixError(f"SQL statement timed out: {statement_id}") + 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 MatrixError(f"SQL failed: {response.as_dict()}") + 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) @@ -234,8 +234,8 @@ def assert_app_absent(self, name: str) -> None: except NotFound: return except DatabricksError as exc: - raise MatrixError(f"Could not verify App {name} is absent: {exc}") from exc - raise MatrixError(f"App {name} already exists; refusing to deploy over it.") + raise E2EError(f"Could not verify App {name} is absent: {exc}") from exc + raise E2EError(f"App {name} already exists; refusing to deploy over it.") def wait_for_app(self, name: str) -> dict[str, Any]: """The App once it is ACTIVE with a URL.""" @@ -252,7 +252,7 @@ def wait_for_app(self, name: str) -> dict[str, Any]: log(f"tick {now():%H:%M} | app-{name} | {state or 'UNKNOWN'}") next_tick += 60 time.sleep(15) - raise MatrixError(f"App {name} did not become ACTIVE.") + raise E2EError(f"App {name} did not become ACTIVE.") 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.""" @@ -285,7 +285,7 @@ def granted_tuples(self, app_name: str) -> set[GrantTuple]: def experiment_id(self, name: str) -> str: experiment = self.client.experiments.get_by_name(name).experiment if experiment is None or not experiment.experiment_id: - raise MatrixError(f"Experiment {name!r} does not exist.") + raise E2EError(f"Experiment {name!r} does not exist.") return str(experiment.experiment_id) def non_tool_resources(self, app_name: str) -> list[dict[str, Any]]: @@ -297,12 +297,12 @@ def transitive_state(self, app_name: str, function: str) -> dict[str, Any]: principal = self._principal(app) direct = direct_privileges(self.client, SecurableType.FUNCTION, function, principal) if "EXECUTE" in direct: - raise MatrixError( + raise E2EError( f"Agent Bricks granted the transitive function directly: {function} -> {direct}" ) effective = effective_privileges(self.client, SecurableType.FUNCTION, function, principal) if "EXECUTE" in effective: - raise MatrixError( + raise E2EError( "The transitive function was already effective before the manual grant: " f"{function} -> {effective}" ) @@ -310,9 +310,7 @@ def transitive_state(self, app_name: str, function: str) -> dict[str, Any]: resource.get("uc_securable", {}).get("securable_full_name") == function for resource in tool_resources(app) ): - raise MatrixError( - "The transitive function appeared in Agent Bricks-owned Apps resources." - ) + raise E2EError("The transitive function appeared in Agent Bricks-owned Apps resources.") return { "service_principal_client_id": principal, "direct_privileges": direct, @@ -339,12 +337,12 @@ def grant_transitive(self, app_name: str, function: str) -> dict[str, Any]: direct = direct_privileges(self.client, SecurableType.FUNCTION, function, principal) effective = effective_privileges(self.client, SecurableType.FUNCTION, function, principal) if "EXECUTE" not in effective: - raise MatrixError( + raise E2EError( "Manual EXECUTE grant did not become effective for the transitive function: " f"{function} -> {effective}" ) if "EXECUTE" not in direct: - raise MatrixError( + raise E2EError( "Manual EXECUTE grant was not persisted directly for the transitive function: " f"{function} -> {direct}" ) @@ -358,7 +356,7 @@ def grant_transitive(self, app_name: str, function: str) -> dict[str, Any]: def _principal(app: dict[str, Any]) -> str: principal = app.get("service_principal_client_id") if not principal: - raise MatrixError(f"App response has no service_principal_client_id: {app}") + raise E2EError(f"App response has no service_principal_client_id: {app}") return str(principal) # Cleanup @@ -377,12 +375,10 @@ def _wait_for_app_deleted(self, name: str, timeout: float = 1200) -> None: log(f"tick {now():%H:%M} | delete-{name} | absent") return except DatabricksError as exc: - raise MatrixError(f"Could not verify deletion of App {name!r}: {exc}") from exc + raise E2EError(f"Could not verify deletion of App {name!r}: {exc}") from exc elapsed = time.monotonic() - started if elapsed >= timeout: - raise MatrixError( - f"App {name!r} still existed {timeout:.0f}s after delete returned." - ) + 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 @@ -484,7 +480,7 @@ def user_name(self) -> str: if self._user_name is None: name = self.client.current_user.me().user_name if not name: - raise MatrixError("The current user has no user_name.") + raise E2EError("The current user has no user_name.") self._user_name = name return self._user_name @@ -591,7 +587,7 @@ def delete_cli_store(kind: str, store: str) -> None: assert delete_store is not None result = delete_store(kind, store) if result.returncode != 0: - raise MatrixError((result.stderr or result.stdout).strip()) + raise E2EError((result.stderr or result.stdout).strip()) # Read before any delete: the principal and branch come from the App and its Runtime Store. known = (identity or AppIdentity()).merged( @@ -624,7 +620,7 @@ def delete_cli_store(kind: str, store: str) -> None: def _split_resources(app: dict[str, Any]) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: resources = app.get("resources") or [] if not isinstance(resources, list): - raise MatrixError(f"App resources are not a list: {resources}") + raise E2EError(f"App resources are not a list: {resources}") def name(resource: dict[str, Any]) -> str: return str(resource.get("name", "")) From 607fc2b51d22b2f9d819f86c070fdcab9e2446ca Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Mon, 12 Oct 2026 22:55:23 +0000 Subject: [PATCH 09/45] [agentbricks] revoke e2e catalog grants and add a standalone leak sweep Co-authored-by: Isaac --- integrations/agentbricks/tests/common.py | 11 ++ .../agentbricks/tests/provisioning.py | 95 +----------- integrations/agentbricks/tests/run_context.py | 16 +- integrations/agentbricks/tests/sweep.py | 146 ++++++++++++++++++ .../agentbricks/tests/workspace_client.py | 20 ++- 5 files changed, 194 insertions(+), 94 deletions(-) create mode 100644 integrations/agentbricks/tests/sweep.py diff --git a/integrations/agentbricks/tests/common.py b/integrations/agentbricks/tests/common.py index 19e2d36e4..926a692b2 100644 --- a/integrations/agentbricks/tests/common.py +++ b/integrations/agentbricks/tests/common.py @@ -9,9 +9,11 @@ import dataclasses import datetime as dt import pathlib +import re import shlex import subprocess import sys +import tempfile from collections.abc import Mapping, Sequence FRAMEWORKS = ("langgraph",) @@ -34,6 +36,15 @@ def project_prefix(run_id: str) -> str: 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``; the sweep 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 the controller can sweep ones a test left behind.""" return f"ab_e2e_{run_id}_" diff --git a/integrations/agentbricks/tests/provisioning.py b/integrations/agentbricks/tests/provisioning.py index 26c63fcdb..61088fd54 100644 --- a/integrations/agentbricks/tests/provisioning.py +++ b/integrations/agentbricks/tests/provisioning.py @@ -4,23 +4,18 @@ it from ``require_workspace`` in its own conftest, and tears it down in the test's own teardown. Cleanup of leftovers happens once, in ``pytest_sessionfinish`` on the controller (the only process -that sees every worker finish): it sweeps what a crashed test or a failed per-test cleanup leaked, -found by the run's name prefix (Apps, source folders, schemas) and by the App identities the workers -recorded in ``apps.jsonl`` (Lakebase databases and roles, which outlive their App). +that sees every worker finish), by ``sweep.sweep_run``. A killed run never gets here, so CI also +runs ``tests/sweep.py`` with the same ``--run-id`` afterwards. """ from __future__ import annotations import pytest -from common import APP_PREFIX, RunConfig, log, project_prefix, scratch_schema_prefix +from common import log from run_context import build_target_workspace, started_run_config +from sweep import sweep_run from target_workspace import TargetWorkspace, TargetWorkspaceUnavailable -from workspace_client import ( - Workspace, - cleanup_app, - delete_lakebase_leftovers, - recorded_app_identities, -) +from workspace_client import Workspace def require_workspace(target_workspace: TargetWorkspace, what: str) -> Workspace: @@ -58,86 +53,8 @@ def pytest_sessionfinish(session: pytest.Session) -> None: workspace = target_workspace.client if workspace is None: return - failures: list[str] = [] - sweeps = ( - ("leaked Apps", lambda: _sweep_apps(workspace, run)), - ("leaked Lakebase databases and roles", lambda: _sweep_lakebase(workspace, run)), - ("leaked source folders", lambda: _sweep_deployment_sources(workspace, run.run_id)), - ( - "leaked schemas", - lambda: _sweep_schemas(workspace, config.getoption("catalog"), run.run_id), - ), - ) - for label, action in sweeps: - try: - failures.extend(action()) - except Exception as exc: - failures.append(f"{label}: {exc}") + failures = sweep_run(workspace, run.run_id, run.output, config.getoption("catalog")) for failure in failures: log(f"cleanup warning | {failure}") if failures and session.exitstatus == 0: session.exitstatus = pytest.ExitCode.TESTS_FAILED - - -def _app_prefix(run_id: str) -> str: - return f"{APP_PREFIX}{project_prefix(run_id)}" - - -def _sweep_apps(workspace: Workspace, run: RunConfig) -> list[str]: - """Delete Apps whose per-test cleanup failed or never ran, found by the run's name prefix.""" - prefix = _app_prefix(run.run_id) - recorded = recorded_app_identities(run.output) - failures: list[str] = [] - for app in list(workspace.client.apps.list()): - if app.name and app.name.startswith(prefix): - log(f"sweeping leaked App {app.name}") - failures += cleanup_app( - workspace, app.name, runtime_store=True, identity=recorded.get(app.name) - ) - return failures - - -def _sweep_lakebase(workspace: Workspace, run: RunConfig) -> list[str]: - """Delete runtime databases and roles of this run's Apps, whether or not the App still exists.""" - recorded = list(recorded_app_identities(run.output).values()) - failures: list[str] = [] - for branch in sorted({item.branch for item in recorded if item.branch}): - failures += delete_lakebase_leftovers( - workspace, - 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 - ], - database_ids=[ - item.database_id for item in recorded if item.branch == branch and item.database_id - ], - ) - return failures - - -def _sweep_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"sweeping deployment source {path}") - try: - workspace.delete_workspace_path(path) - except Exception as exc: - failures.append(f"deployment source {path}: {exc}") - return failures - - -def _sweep_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"sweeping 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/run_context.py b/integrations/agentbricks/tests/run_context.py index bd4368e53..c859bc2de 100644 --- a/integrations/agentbricks/tests/run_context.py +++ b/integrations/agentbricks/tests/run_context.py @@ -10,12 +10,11 @@ import os import pathlib import re -import tempfile import uuid from typing import Any import pytest -from common import E2EError, RunConfig, log, run_command +from common import RUN_ID_PATTERN, E2EError, RunConfig, log, run_command, run_output_dir from target_workspace import ( WORKSPACE_KINDS, TargetWorkspace, @@ -61,6 +60,12 @@ def pytest_addoption(parser: pytest.Parser) -> None: action="store_true", help="Skip per-App USE CATALOG grants because catalog access is pre-provisioned.", ) + group.addoption( + "--run-id", + default=os.environ.get("AGENTBRICKS_E2E_RUN_ID"), + help="Names every resource of the run, so `tests/sweep.py --run-id` can clean up after a " + "killed run (AGENTBRICKS_E2E_RUN_ID; default: random). 1-8 lowercase letters or digits.", + ) group.addoption( "--keep-resources", action="store_true", @@ -72,6 +77,9 @@ def pytest_configure(config: pytest.Config) -> None: sha = config.getoption("bridge_sha") if sha is not None and re.fullmatch(r"[0-9a-f]{40}", sha) is None: raise pytest.UsageError("--bridge-sha must be a 40-character lowercase Git commit SHA.") + 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.") @pytest.hookimpl(optionalhook=True) @@ -89,8 +97,8 @@ def base_run_config(config: pytest.Config) -> RunConfig: if handoff: run_id, output = handoff["run_id"], pathlib.Path(handoff["output"]) else: - run_id = uuid.uuid4().hex[:6] - output = pathlib.Path(tempfile.gettempdir()) / f"agentbricks-e2e-{run_id}" + 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, diff --git a/integrations/agentbricks/tests/sweep.py b/integrations/agentbricks/tests/sweep.py new file mode 100644 index 000000000..327cb056f --- /dev/null +++ b/integrations/agentbricks/tests/sweep.py @@ -0,0 +1,146 @@ +"""Delete whatever one E2E run leaked, inside pytest or as a separate step after it. + +Per-test teardown deletes what each test created; this catches what a crashed test or a failed +teardown left behind. pytest runs it at the end of a run, but a job timeout or cancellation kills +pytest before that, which is exactly when leaks happen. So CI picks the run id up front and runs +this again in an ``always()`` step, from ``integrations/agentbricks``: + + uv run pytest tests/suites/tools --run-id "$RUN_ID" ... + uv run python tests/sweep.py --run-id "$RUN_ID" --databricks-profile + +Everything is found by the run's name prefix (Apps, source folders, schemas) or by the App +identities workers recorded in ``/apps.jsonl`` (Lakebase databases and roles, which outlive +their App), so it touches nothing outside the run and is safe to run twice. +""" + +from __future__ import annotations + +import argparse +import os +import pathlib +import sys + +from common import ( + APP_PREFIX, + RUN_ID_PATTERN, + log, + project_prefix, + run_output_dir, + scratch_schema_prefix, +) +from workspace_client import ( + Workspace, + cleanup_app, + delete_lakebase_leftovers, + recorded_app_identities, +) + +DEFAULT_CATALOG = os.environ.get("AGENTBRICKS_E2E_CATALOG", "main") + + +def sweep_run(workspace: Workspace, run_id: str, output: pathlib.Path, catalog: str) -> list[str]: + """Delete everything the run left; returns a description of each failure.""" + failures: list[str] = [] + sweeps = ( + ("leaked Apps", lambda: _sweep_apps(workspace, run_id, output)), + ("leaked Lakebase databases and roles", lambda: _sweep_lakebase(workspace, output)), + ("leaked source folders", lambda: _sweep_deployment_sources(workspace, run_id)), + ("leaked schemas", lambda: _sweep_schemas(workspace, catalog, run_id)), + ) + for label, action in sweeps: + 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 _sweep_apps(workspace: Workspace, run_id: str, output: pathlib.Path) -> list[str]: + """Delete Apps whose per-test cleanup failed or never ran, found by the run's name prefix.""" + prefix = _app_prefix(run_id) + recorded = recorded_app_identities(output) + failures: list[str] = [] + for app in list(workspace.client.apps.list()): + if app.name and app.name.startswith(prefix): + log(f"sweeping leaked App {app.name}") + failures += cleanup_app( + workspace, app.name, runtime_store=True, identity=recorded.get(app.name) + ) + return failures + + +def _sweep_lakebase(workspace: Workspace, output: pathlib.Path) -> list[str]: + """Delete runtime databases and roles of this run's Apps, whether or not the App still exists.""" + recorded = list(recorded_app_identities(output).values()) + failures: list[str] = [] + for branch in sorted({item.branch for item in recorded if item.branch}): + failures += delete_lakebase_leftovers( + workspace, + 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 + ], + database_ids=[ + item.database_id for item in recorded if item.branch == branch and item.database_id + ], + ) + return failures + + +def _sweep_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"sweeping deployment source {path}") + try: + workspace.delete_workspace_path(path) + except Exception as exc: + failures.append(f"deployment source {path}: {exc}") + return failures + + +def _sweep_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"sweeping leaked schema {schema}") + try: + workspace.delete_schema(schema) + except Exception as exc: + failures.append(f"schema {schema}: {exc}") + return failures + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description="Delete whatever one Agent Bricks E2E run leaked.") + parser.add_argument("--run-id", required=True, help="The run id the pytest run was given.") + parser.add_argument("--databricks-profile", help="Databricks CLI profile; default is ambient.") + parser.add_argument("--catalog", default=DEFAULT_CATALOG, help="Catalog of the run's schemas.") + parser.add_argument( + "--output", + type=pathlib.Path, + help="The run's output dir, for apps.jsonl (default: the one pytest derives from the id).", + ) + args = parser.parse_args(argv) + if RUN_ID_PATTERN.fullmatch(args.run_id) is None: + parser.error("--run-id must be 1-8 lowercase letters or digits") + output = args.output or run_output_dir(args.run_id) + if not (output / "apps.jsonl").exists(): + # Without the record, Lakebase leftovers of Apps that are already gone can't be found. + log(f"no apps.jsonl in {output}; sweeping only what the run prefix finds") + failures = sweep_run(Workspace(args.databricks_profile), args.run_id, output, args.catalog) + for failure in failures: + log(f"cleanup warning | {failure}") + return 1 if failures else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/integrations/agentbricks/tests/workspace_client.py b/integrations/agentbricks/tests/workspace_client.py index 507a607f4..bc448a853 100644 --- a/integrations/agentbricks/tests/workspace_client.py +++ b/integrations/agentbricks/tests/workspace_client.py @@ -19,7 +19,7 @@ from common import TOOL_RESOURCE_PREFIX, E2EError, log, now from databricks.sdk import WorkspaceClient from databricks.sdk.errors import DatabricksError, NotFound -from databricks.sdk.service.catalog import SecurableType, VolumeType +from databricks.sdk.service.catalog import PermissionsChange, Privilege, SecurableType, VolumeType from databricks.sdk.service.sql import ExecuteStatementRequestOnWaitTimeout, State, StatementState GrantTuple = tuple[str, str, str, str] @@ -90,6 +90,8 @@ def __init__( self._app_auth_checked = False self._warehouse_started = False self._user_name: str | None = None + # App name -> (catalog, principal) of each USE CATALOG this suite granted, for cleanup. + self._catalog_grants: dict[str, tuple[str, str]] = {} # App auth: Databricks Apps /api routes need OAuth, so a PAT is rejected. @@ -317,6 +319,18 @@ def transitive_state(self, app_name: str, function: str) -> dict[str, Any]: "effective_privileges": effective, } + def revoke_catalog_grant(self, app_name: str) -> None: + """Revoke the USE CATALOG ``grant_transitive`` gave the App's principal, if it gave one.""" + if app_name not in self._catalog_grants: + return + catalog, principal = self._catalog_grants[app_name] + self.client.grants.update( + SecurableType.CATALOG.value, + catalog, + changes=[PermissionsChange(principal=principal, remove=[Privilege.USE_CATALOG])], + ) + del self._catalog_grants[app_name] + def grant_transitive(self, app_name: str, function: str) -> dict[str, Any]: """Grant the App principal EXECUTE on a nested function and confirm it is direct and effective.""" principal = self._principal(self.app(app_name)) @@ -324,6 +338,9 @@ def grant_transitive(self, app_name: str, function: str) -> dict[str, Any]: quoted_principal = f"`{principal.replace('`', '``')}`" statements = [] if not self.preprovisioned_app_catalog_access: + # Recorded first so a partial failure is still revoked. The schema-level grants go + # with the test's schema, but the catalog is shared and outlives the run. + self._catalog_grants[app_name] = (catalog, principal) statements.append(f"GRANT USE CATALOG ON CATALOG `{catalog}` TO {quoted_principal}") statements.extend( ( @@ -599,6 +616,7 @@ def delete_cli_store(kind: str, store: str) -> None: if not has_app: return failures + attempt(f"catalog grant {name}", workspace.revoke_catalog_grant, name) if runtime_store: attempt(f"runtime store {name}", workspace.delete_runtime_store, name) app_deleted = attempt(f"app {name}", workspace.delete_app, name) From ef830d764a01cd24f0209a4cb73c71440f55e762 Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Tue, 13 Oct 2026 00:06:54 +0000 Subject: [PATCH 10/45] [agentbricks] move e2e suite harness under tests/suites and run leftover cleanup via --cleanup-only Co-authored-by: Isaac --- integrations/agentbricks/tests/common.py | 116 ------- integrations/agentbricks/tests/conftest.py | 77 ----- .../agentbricks/tests/provisioning.py | 60 ---- integrations/agentbricks/tests/run_context.py | 150 --------- .../agentbricks/tests/suites/README.md | 63 ++++ .../agentbricks/tests/suites/conftest.py | 286 ++++++++++++++++++ .../tests/suites/harness/__init__.py | 1 + .../tests/suites/harness/cleanup.py | 245 +++++++++++++++ .../harness/cli.py} | 23 +- .../tests/suites/harness/process.py | 63 ++++ .../agentbricks/tests/suites/harness/run.py | 60 ++++ .../harness/target.py} | 14 +- .../harness/workspace.py} | 140 +-------- .../suites/memory_store/test_memory_store.py | 2 +- .../session_store/test_session_store.py | 2 +- .../tests/suites/tools/conftest.py | 21 +- .../tests/suites/tools/test_deploy.py | 6 +- .../tests/suites/tools/test_dev.py | 4 +- .../tests/suites/tools/test_grants.py | 8 +- .../tests/suites/tools/test_selection.py | 6 +- .../tools/{tools.py => tool_definitions.py} | 4 +- .../tests/suites/tools/uc_objects.py | 4 +- .../tests/suites/tracing/test_tracing.py | 6 +- integrations/agentbricks/tests/sweep.py | 146 --------- 24 files changed, 773 insertions(+), 734 deletions(-) delete mode 100644 integrations/agentbricks/tests/common.py delete mode 100644 integrations/agentbricks/tests/conftest.py delete mode 100644 integrations/agentbricks/tests/provisioning.py delete mode 100644 integrations/agentbricks/tests/run_context.py create mode 100644 integrations/agentbricks/tests/suites/README.md create mode 100644 integrations/agentbricks/tests/suites/conftest.py create mode 100644 integrations/agentbricks/tests/suites/harness/__init__.py create mode 100644 integrations/agentbricks/tests/suites/harness/cleanup.py rename integrations/agentbricks/tests/{agentbricks_cli.py => suites/harness/cli.py} (98%) create mode 100644 integrations/agentbricks/tests/suites/harness/process.py create mode 100644 integrations/agentbricks/tests/suites/harness/run.py rename integrations/agentbricks/tests/{target_workspace.py => suites/harness/target.py} (85%) rename integrations/agentbricks/tests/{workspace_client.py => suites/harness/workspace.py} (82%) rename integrations/agentbricks/tests/suites/tools/{tools.py => tool_definitions.py} (98%) delete mode 100644 integrations/agentbricks/tests/sweep.py diff --git a/integrations/agentbricks/tests/common.py b/integrations/agentbricks/tests/common.py deleted file mode 100644 index 926a692b2..000000000 --- a/integrations/agentbricks/tests/common.py +++ /dev/null @@ -1,116 +0,0 @@ -"""Constants, ``RunConfig`` and logging/subprocess helpers shared by the E2E suites. - -``run_context.py`` builds one ``RunConfig`` per pytest run; fixtures and tests read it through the -``run_config`` fixture. -""" - -from __future__ import annotations - -import dataclasses -import datetime as dt -import pathlib -import re -import shlex -import subprocess -import sys -import tempfile -from collections.abc import Mapping, Sequence - -FRAMEWORKS = ("langgraph",) -AUTHORING_PATHS = ("cli", "direct") -E2E_MODEL = "system.ai.gpt-5-2" -TOOL_RESOURCE_PREFIX = "agentbricks-tool-" -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 teardown can sweep the Apps a crashed test leaked.""" - 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``; the sweep 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 the controller can sweep ones a test left behind.""" - return f"ab_e2e_{run_id}_" - - -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" - - -def run_command( - argv: Sequence[str], - *, - cwd: pathlib.Path | None = None, - env: Mapping[str, str] | None = None, - timeout: float = 300, -) -> str: - """Run a command, echo its output, and return stdout; a non-zero exit is a ``E2EError``.""" - log_command(argv, cwd) - result = subprocess.run( - list(argv), - cwd=cwd, - env=None if env is None else dict(env), - text=True, - capture_output=True, - timeout=timeout, - check=False, - ) - for text in (result.stdout, result.stderr): - if text.strip(): - log(text) - if result.returncode != 0: - raise E2EError( - f"Command failed ({result.returncode}): {shlex.join(argv)}\n" - f"{result.stderr or result.stdout}" - ) - return result.stdout - - -@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 - app_auth_profile: str | None - # Operator override; None lets the workspace pick the first running warehouse. - warehouse_id: str | None - preprovisioned_app_catalog_access: bool - bridge_sha: str | None - wheel: pathlib.Path | None diff --git a/integrations/agentbricks/tests/conftest.py b/integrations/agentbricks/tests/conftest.py deleted file mode 100644 index 91f16a398..000000000 --- a/integrations/agentbricks/tests/conftest.py +++ /dev/null @@ -1,77 +0,0 @@ -"""Fixtures shared by every suite under ``suites/``. - -``run_context`` (options, target workspace, run identity) and ``provisioning`` (the catalog -requirement and the end-of-run cleanup) are plugins. One command runs a suite, from ``integrations/agentbricks``: - - uv run pytest tests/suites/tools -n 6 --dist loadgroup \ - --databricks-profile - -``--workspace fake`` runs against no workspace; whatever needs one is skipped. - -The CLI wrapper is function-scoped: each test authors, runs and deploys its own projects, and the -fixture deletes every App and store they created when the test ends. -""" - -from __future__ import annotations - -from collections.abc import Iterator - -import pytest -from agentbricks_cli import AgentbricksCli -from common import RunConfig, log -from target_workspace import TargetWorkspace -from workspace_client import Workspace, cleanup_app - -pytest_plugins = ["run_context", "provisioning"] - -_failed_tests: set[str] = set() - - -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(scope="session") -def workspace_client(target_workspace: TargetWorkspace) -> Workspace: - workspace = target_workspace.client - if workspace is None: - pytest.skip(f"--workspace {target_workspace.name} has no Databricks SDK access") - return workspace - - -@pytest.fixture -def agentbricks_cli( - request: pytest.FixtureRequest, target_workspace: TargetWorkspace, run_config: RunConfig -) -> Iterator[AgentbricksCli]: - cli = AgentbricksCli(target_workspace, run_config) - yield cli - if cli.projects and target_workspace.client is not None: - _delete_projects(cli, target_workspace.client, failed=request.node.nodeid in _failed_tests) - - -def _delete_projects(cli: AgentbricksCli, workspace: Workspace, *, failed: bool) -> None: - """Delete every App, store, source folder and Lakebase leftover the test's projects created.""" - for project in cli.projects: - if failed and project.app_registered: - # The App is deleted next, so a failing test must capture its logs while it exists. - workspace.app_logs( - project.app_name, cli.logs_dir / f"deploy-runtime-{project.app_name}.log" - ) - try: - failures = cleanup_app( - workspace, - project.app_name, - delete_store=cli.delete_store, - memory_store=project.memory_store_name, - session_store=project.session_store_name, - has_app=project.app_registered, - # Every managed deploy creates a Runtime Store, with or without memory/session stores. - runtime_store=project.app_registered, - identity=project.identity, - ) - except Exception as exc: - failures = [str(exc)] - for failure in failures: - log(f"cleanup warning | {project.app_name} | {failure}") diff --git a/integrations/agentbricks/tests/provisioning.py b/integrations/agentbricks/tests/provisioning.py deleted file mode 100644 index 61088fd54..000000000 --- a/integrations/agentbricks/tests/provisioning.py +++ /dev/null @@ -1,60 +0,0 @@ -"""pytest plugin: the workspace requirements shared by every suite, and the end-of-run cleanup. - -``catalog`` is the one fixture here; a suite that needs more (the tools suite's UC schema) builds -it from ``require_workspace`` in its own conftest, and tears it down in the test's own teardown. - -Cleanup of leftovers happens once, in ``pytest_sessionfinish`` on the controller (the only process -that sees every worker finish), by ``sweep.sweep_run``. A killed run never gets here, so CI also -runs ``tests/sweep.py`` with the same ``--run-id`` afterwards. -""" - -from __future__ import annotations - -import pytest -from common import log -from run_context import build_target_workspace, started_run_config -from sweep import sweep_run -from target_workspace import TargetWorkspace, TargetWorkspaceUnavailable -from workspace_client import Workspace - - -def require_workspace(target_workspace: TargetWorkspace, what: str) -> Workspace: - workspace = target_workspace.client - if workspace is None: - pytest.skip(f"{what} needs a live workspace; got --workspace {target_workspace.name}") - return workspace - - -@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 - - -@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 - run = started_run_config(config) - # The output dir exists once any test has run, so its absence means nothing was provisioned. - if run is None or not run.output.exists(): - return - if config.getoption("keep_resources"): - log("Resources retained (--keep-resources); skipping cleanup.") - return - try: - target_workspace = build_target_workspace(config) - except TargetWorkspaceUnavailable: - return - workspace = target_workspace.client - if workspace is None: - return - failures = sweep_run(workspace, run.run_id, run.output, config.getoption("catalog")) - for failure in failures: - log(f"cleanup warning | {failure}") - if failures and session.exitstatus == 0: - session.exitstatus = pytest.ExitCode.TESTS_FAILED diff --git a/integrations/agentbricks/tests/run_context.py b/integrations/agentbricks/tests/run_context.py deleted file mode 100644 index c859bc2de..000000000 --- a/integrations/agentbricks/tests/run_context.py +++ /dev/null @@ -1,150 +0,0 @@ -"""pytest plugin: command-line options, the run's identity, and the workspace the suite runs against. - -Under xdist the controller picks the run id and output directory and hands them to each worker in -``workerinput``, so every process names, finds and later tears down the same resources. -""" - -from __future__ import annotations - -import dataclasses -import os -import pathlib -import re -import uuid -from typing import Any - -import pytest -from common import RUN_ID_PATTERN, E2EError, RunConfig, log, run_command, run_output_dir -from target_workspace import ( - WORKSPACE_KINDS, - TargetWorkspace, - TargetWorkspaceUnavailable, - make_target_workspace, -) - -PACKAGE_DIR = pathlib.Path(__file__).resolve().parent.parent -_RUN_KEY = pytest.StashKey[RunConfig]() -_HANDOFF = "agentbricks_run" - - -def pytest_addoption(parser: pytest.Parser) -> None: - group = parser.getgroup("agentbricks", "Agent Bricks E2E suites") - group.addoption( - "--workspace", - choices=WORKSPACE_KINDS, - default="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( - "--app-auth-profile", - help="OAuth profile for deployed App /api calls; defaults to --databricks-profile.", - ) - group.addoption( - "--catalog", - default=os.environ.get("AGENTBRICKS_E2E_CATALOG", "main"), - help="Existing catalog for temporary test schemas (AGENTBRICKS_E2E_CATALOG, default main).", - ) - group.addoption("--bridge-sha", help="Immutable bridge commit for generated App dependencies.") - 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( - "--preprovisioned-app-catalog-access", - action="store_true", - help="Skip per-App USE CATALOG grants because catalog access is pre-provisioned.", - ) - group.addoption( - "--run-id", - default=os.environ.get("AGENTBRICKS_E2E_RUN_ID"), - help="Names every resource of the run, so `tests/sweep.py --run-id` can clean up after a " - "killed run (AGENTBRICKS_E2E_RUN_ID; default: random). 1-8 lowercase letters or digits.", - ) - group.addoption( - "--keep-resources", - action="store_true", - help="Skip the end-of-run sweep of leaked Apps and temporary schemas.", - ) - - -def pytest_configure(config: pytest.Config) -> None: - sha = config.getoption("bridge_sha") - if sha is not None and re.fullmatch(r"[0-9a-f]{40}", sha) is None: - raise pytest.UsageError("--bridge-sha must be a 40-character lowercase Git commit SHA.") - 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.") - - -@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, - app_auth_profile=option("app_auth_profile") or option("databricks_profile"), - warehouse_id=option("warehouse_id"), - preprovisioned_app_catalog_access=bool(option("preprovisioned_app_catalog_access")), - bridge_sha=option("bridge_sha"), - wheel=option("wheel"), - ) - config.stash[_RUN_KEY] = run - return run - - -def started_run_config(config: pytest.Config) -> RunConfig | None: - """The run's config if this process ever created one; None means no test needed it.""" - return config.stash.get(_RUN_KEY, None) - - -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 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 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 or run.bridge_sha: - 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") - run_command(["uv", "build", "--wheel", "--out-dir", str(dist)], cwd=PACKAGE_DIR, timeout=600) - wheels = sorted(dist.glob("*.whl")) - if not wheels: - raise E2EError("uv build produced no wheel") - return dataclasses.replace(run, wheel=wheels[-1]) diff --git a/integrations/agentbricks/tests/suites/README.md b/integrations/agentbricks/tests/suites/README.md new file mode 100644 index 000000000..7db656a5b --- /dev/null +++ b/integrations/agentbricks/tests/suites/README.md @@ -0,0 +1,63 @@ +# Agent Bricks E2E suites + +Black-box tests of the `agentbricks` CLI: each test runs the real CLI (`init`, `tools add`, `dev`, +`deploy`, `endpoint invoke`) and checks the agent it produces. + +## Running + +From `integrations/agentbricks`: + +```bash +uv run pytest tests/suites -n 6 --dist loadgroup \ + --databricks-profile --catalog +``` + +- Narrow it to a suite (`tests/suites/tools`), a file, or a test (`"...test_deploy.py::test_deploy[cli]"`). +- `-n 6 --dist loadgroup` runs tests in parallel; `loadgroup` keeps the dev tests on one worker, + because local dev servers share a proxy port. +- `AGENTBRICKS_E2E_GENIE_SPACE_ID=` adds the Genie tool; without it that tool is skipped. +- `--workspace fake` runs against no workspace; everything that needs one is skipped. +- `--keep-resources` keeps leftovers for debugging; each test still deletes its own resources. + +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**, by the fixture that made it: `agentbricks_cli` deletes the test's Apps and + everything their deploys created (`harness.cleanup.delete_app_resources`); `scratch_schema` deletes + the test's schema. +2. **When the run ends**, `pytest_sessionfinish` deletes whatever a crashed test or a failed teardown + left (`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` (or `AGENTBRICKS_E2E_RUN_ID`) and runs this in an `always()` step. + +## Layout + +| Path | What it is | +|---|---| +| `conftest.py` | All pytest wiring: options, hooks, and the fixtures every suite shares | +| `harness/run.py` | `RunConfig`, the run id, and the name prefixes every resource carries | +| `harness/target.py` | The `--workspace live\|fake` seam the CLI and fixtures talk to | +| `harness/cli.py` | The `agentbricks` CLI under test, and the agents it produces | +| `harness/workspace.py` | Databricks SDK helpers for what is not under test: SQL, UC, Apps, grants | +| `harness/cleanup.py` | Deleting what the suites create, per App and per run | +| `harness/process.py` | Logging and subprocess helpers | +| `/` | One folder per feature: `tools`, `memory_store`, `session_store`, `tracing` | + +`conftest.py` files are the only code pytest calls by name; everything in `harness/` is plain Python +that you can follow through its imports. + +## Adding things + +- **A tool:** add one `ToolDefinition(name, requires, build)` in `tools/tool_definitions.py`. + `requires` names the fixtures it needs; tools whose fixtures skip are left out and logged. +- **A suite:** add a folder with `test_*.py` files. Add a `conftest.py` only for resources the CLI + under test does not create itself; create them per test and delete them 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..e4edd643d --- /dev/null +++ b/integrations/agentbricks/tests/suites/conftest.py @@ -0,0 +1,286 @@ +"""All pytest wiring for the E2E suites: options, hooks, and the fixtures every suite shares. + +pytest calls everything in this file by name; ``harness/`` is plain code, so following its imports +answers what calls what. One command runs a suite, from ``integrations/agentbricks``: + + uv run pytest tests/suites/tools -n 6 --dist loadgroup \ + --databricks-profile + +``--workspace fake`` runs against no workspace; whatever needs one is skipped. + +Each test deletes what it created when it ends (the ``agentbricks_cli`` fixture here, and suite +fixtures such as the tools suite's ``scratch_schema``). After every live run, the controller also +deletes whatever a crashed test or a failed teardown 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 os +import pathlib +import re +import uuid +from collections.abc import Iterator +from typing import Any + +import pytest +from harness.cleanup import delete_app_resources, delete_run_leftovers +from harness.cli import AgentbricksCli +from harness.process import log, run_command +from harness.run import RUN_ID_PATTERN, E2EError, RunConfig, run_output_dir +from harness.target import ( + WORKSPACE_KINDS, + TargetWorkspace, + TargetWorkspaceUnavailable, + make_target_workspace, + require_workspace, +) +from harness.workspace import Workspace + +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", + choices=WORKSPACE_KINDS, + default="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( + "--app-auth-profile", + help="OAuth profile for deployed App /api calls; defaults to --databricks-profile.", + ) + group.addoption( + "--catalog", + default=os.environ.get("AGENTBRICKS_E2E_CATALOG", "main"), + help="Existing catalog for temporary test schemas (AGENTBRICKS_E2E_CATALOG, default main).", + ) + group.addoption("--bridge-sha", help="Immutable bridge commit for generated App dependencies.") + 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( + "--preprovisioned-app-catalog-access", + action="store_true", + help="Skip per-App USE CATALOG grants because catalog access is pre-provisioned.", + ) + group.addoption( + "--run-id", + default=os.environ.get("AGENTBRICKS_E2E_RUN_ID"), + help="Names every resource of the run, so --cleanup-only can delete what a killed run left " + "(AGENTBRICKS_E2E_RUN_ID; 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: + sha = config.getoption("bridge_sha") + if sha is not None and re.fullmatch(r"[0-9a-f]{40}", sha) is None: + raise pytest.UsageError("--bridge-sha must be a 40-character lowercase Git commit SHA.") + 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, + app_auth_profile=option("app_auth_profile") or option("databricks_profile"), + warehouse_id=option("warehouse_id"), + preprovisioned_app_catalog_access=bool(option("preprovisioned_app_catalog_access")), + bridge_sha=option("bridge_sha"), + 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 or run.bridge_sha: + 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") + run_command(["uv", "build", "--wheel", "--out-dir", str(dist)], cwd=PACKAGE_DIR, timeout=600) + wheels = sorted(dist.glob("*.whl")) + if not wheels: + raise E2EError("uv build produced no wheel") + return dataclasses.replace(run, wheel=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 + + +# 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 +) -> Iterator[AgentbricksCli]: + """The CLI wrapper for one test; deletes every App and store the test created when it ends.""" + cli = AgentbricksCli(target_workspace, run_config) + yield cli + if cli.projects and target_workspace.client is not None: + _delete_projects(cli, target_workspace.client, failed=request.node.nodeid in _failed_tests) + + +def _delete_projects(cli: AgentbricksCli, workspace: Workspace, *, failed: bool) -> None: + for project in cli.projects: + if failed and project.app_registered: + # The App is deleted next, so a failing test must capture its logs while it exists. + workspace.app_logs( + project.app_name, cli.logs_dir / f"deploy-runtime-{project.app_name}.log" + ) + try: + failures = delete_app_resources( + workspace, + project.app_name, + delete_store=cli.delete_store, + memory_store=project.memory_store_name, + session_store=project.session_store_name, + has_app=project.app_registered, + # Every managed deploy creates a Runtime Store, with or without memory/session stores. + runtime_store=project.app_registered, + identity=project.identity, + ) + except Exception as exc: + failures = [str(exc)] + for failure in failures: + log(f"cleanup warning | {project.app_name} | {failure}") + + +# 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: + workspace = _build_target_workspace(config).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") + failures = delete_run_leftovers(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/harness/__init__.py b/integrations/agentbricks/tests/suites/harness/__init__.py new file mode 100644 index 000000000..47679ee61 --- /dev/null +++ b/integrations/agentbricks/tests/suites/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/suites/harness/cleanup.py b/integrations/agentbricks/tests/suites/harness/cleanup.py new file mode 100644 index 000000000..cf8507e27 --- /dev/null +++ b/integrations/agentbricks/tests/suites/harness/cleanup.py @@ -0,0 +1,245 @@ +"""Deleting what the E2E suites create: per App when a test ends, and per run afterwards. + +Two entry points, both safe to repeat (a missing resource counts as deleted): + +- ``delete_app_resources``: everything one App's deploy created. The ``agentbricks_cli`` fixture + calls it for each project when its test ends. +- ``delete_run_leftovers``: whatever a crashed test or a failed teardown left. It runs once after + every live run, and on its own with ``pytest tests/suites --cleanup-only --run-id `` after a + run that was killed before it could. It finds resources by the run's name prefix (Apps, source + folders, schemas) and by the App identities recorded at deploy time in ``/apps.jsonl`` + (Lakebase databases and roles, which outlive their App), so it touches nothing outside the run. +""" + +from __future__ import annotations + +import dataclasses +import json +import pathlib +import subprocess +from collections.abc import Callable, Collection +from typing import Any + +from harness.process import log +from harness.run import APP_PREFIX, E2EError, project_prefix, scratch_schema_prefix +from harness.workspace import AppIdentity, Workspace + +_APPS_RECORD = "apps.jsonl" + + +# Identities recorded at deploy time + + +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" + # 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"), data.get("database_id") + ) + identities[data["app"]] = identities.get(data["app"], AppIdentity()).merged(identity) + except (ValueError, KeyError, TypeError, AttributeError): + continue + return identities + + +# One App + + +def _delete_each( + kind: str, find: Callable[[], list[str]], delete: Callable[[str], None] +) -> list[str]: + """Delete everything ``find`` returns, continuing past failures; returns one line per failure.""" + try: + names = find() + except Exception as exc: + return [f"{kind} lookup: {exc}"] + failures = [] + for name in names: + try: + delete(name) + except Exception as exc: + failures.append(f"{kind} {name}: {exc}") + return failures + + +def delete_lakebase_leftovers( + workspace: Workspace, + branch: str, + *, + principals: Collection[str], + database_ids: Collection[str] = (), +) -> list[str]: + """Delete runtime databases and then the principals' roles on one branch. + + A role that owns a database cannot be deleted, so databases go first. Only ``runtime-`` + databases are touched, and both kinds are matched by principal UUID, never by name pattern. + """ + return _delete_each( + "runtime database", + lambda: workspace.runtime_databases(branch, principals, database_ids), + workspace.delete_database, + ) + _delete_each( + "Lakebase role", + lambda: workspace.roles_of(branch, principals), + workspace.delete_role, + ) + + +def delete_app_resources( + workspace: Workspace, + name: str, + *, + delete_store: Callable[[str, str], subprocess.CompletedProcess[str]] | None = None, + memory_store: str | None = None, + session_store: str | None = None, + has_app: bool = True, + runtime_store: bool = False, + identity: AppIdentity | None = None, +) -> list[str]: + """Delete one project's stores and, if it has an App, everything the App's deploy created. + + ``delete_store(kind, store)`` is the agentbricks CLI's store delete, which has no SDK surface. + ``identity`` is what was recorded at deploy time; it fills in whatever the live lookup misses. + Every step runs regardless of earlier failures except those that need the App gone. Returns a + description of each failure; an empty list means everything was deleted. + """ + failures: list[str] = [] + + def attempt(label: str, action: Callable[..., object], *args: Any) -> bool: + try: + action(*args) + except Exception as exc: + failures.append(f"{label}: {exc}") + return False + return True + + def delete_cli_store(kind: str, store: str) -> None: + assert delete_store is not None + result = delete_store(kind, store) + if result.returncode != 0: + raise E2EError((result.stderr or result.stdout).strip()) + + # Read before any delete: the principal and branch come from the App and its Runtime Store. + known = (identity or AppIdentity()).merged( + workspace.app_identity(name) if has_app else AppIdentity() + ) + for kind, store in (("memory", memory_store), ("sessions", session_store)): + if store and delete_store: + attempt(f"{kind} store {store}", delete_cli_store, kind, store) + if not has_app: + return failures + + attempt(f"catalog grant {name}", workspace.revoke_catalog_grant, name) + if runtime_store: + attempt(f"runtime store {name}", workspace.delete_runtime_store, name) + app_deleted = attempt(f"app {name}", workspace.delete_app, name) + if app_deleted: + attempt(f"deployment source {name}", workspace.delete_deployment_source, name) + if known.branch: + failures += delete_lakebase_leftovers( + workspace, + known.branch, + principals=[known.service_principal or ""], + database_ids=[known.database_id or ""], + ) + return failures + + +# One run + + +def delete_run_leftovers( + 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 Apps", lambda: _delete_leaked_apps(workspace, run_id, output)), + ("leaked Lakebase databases and roles", lambda: _delete_leaked_lakebase(workspace, output)), + ("leaked source folders", lambda: _delete_leaked_deployment_sources(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_apps(workspace: Workspace, run_id: str, output: pathlib.Path) -> list[str]: + """Delete Apps whose per-test cleanup failed or never ran, found by the run's name prefix.""" + prefix = _app_prefix(run_id) + recorded = recorded_app_identities(output) + failures: list[str] = [] + for app in list(workspace.client.apps.list()): + if app.name and app.name.startswith(prefix): + log(f"deleting leaked App {app.name}") + failures += delete_app_resources( + workspace, app.name, runtime_store=True, identity=recorded.get(app.name) + ) + return failures + + +def _delete_leaked_lakebase(workspace: Workspace, output: pathlib.Path) -> list[str]: + """Delete runtime databases and roles of this run's Apps, whether or not the App still exists.""" + recorded = list(recorded_app_identities(output).values()) + failures: list[str] = [] + for branch in sorted({item.branch for item in recorded if item.branch}): + failures += delete_lakebase_leftovers( + workspace, + 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 + ], + database_ids=[ + item.database_id for item in recorded if item.branch == branch and item.database_id + ], + ) + 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_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/agentbricks_cli.py b/integrations/agentbricks/tests/suites/harness/cli.py similarity index 98% rename from integrations/agentbricks/tests/agentbricks_cli.py rename to integrations/agentbricks/tests/suites/harness/cli.py index 109e4e027..7ad11b916 100644 --- a/integrations/agentbricks/tests/agentbricks_cli.py +++ b/integrations/agentbricks/tests/suites/harness/cli.py @@ -30,21 +30,12 @@ import tomli import tomlkit -from common import ( - APP_PREFIX, - E2E_MODEL, - FRAMEWORKS, - E2EError, - RunConfig, - last_lines, - last_nonempty_line, - log, - log_command, - now, - project_prefix, -) -from target_workspace import TargetWorkspace -from workspace_client import AppIdentity, Workspace, record_app_identity + +from harness.cleanup import record_app_identity +from harness.process import last_lines, last_nonempty_line, log, log_command +from harness.run import APP_PREFIX, E2E_MODEL, FRAMEWORKS, E2EError, RunConfig, now, project_prefix +from harness.target import TargetWorkspace +from harness.workspace import AppIdentity, Workspace _STORE_KINDS = {"memory": ("memory", "stores"), "sessions": ("sessions", "stores")} _DIRECT_HEADER = ( @@ -460,7 +451,7 @@ def deploy(self, project: Project, app: str | None = None) -> Agent: return agent def _record_identity(self, project: Project, app: str, workspace: Workspace) -> None: - """Keep the App's principal and Lakebase branch for cleanup, and for the controller's sweep.""" + """Keep the App's principal and Lakebase branch for cleanup, including leftover cleanup after the run.""" identity = project.identity.merged(workspace.app_identity(app)) if identity != project.identity: project.identity = identity diff --git a/integrations/agentbricks/tests/suites/harness/process.py b/integrations/agentbricks/tests/suites/harness/process.py new file mode 100644 index 000000000..9251ab5d3 --- /dev/null +++ b/integrations/agentbricks/tests/suites/harness/process.py @@ -0,0 +1,63 @@ +"""Logging and subprocess helpers: every command the suites run is echoed, so a log replays the run.""" + +from __future__ import annotations + +import pathlib +import shlex +import subprocess +import sys +from collections.abc import Mapping, Sequence + +from harness.run import E2EError + + +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" + + +def run_command( + argv: Sequence[str], + *, + cwd: pathlib.Path | None = None, + env: Mapping[str, str] | None = None, + timeout: float = 300, +) -> str: + """Run a command, echo its output, and return stdout; a non-zero exit is a ``E2EError``.""" + log_command(argv, cwd) + result = subprocess.run( + list(argv), + cwd=cwd, + env=None if env is None else dict(env), + text=True, + capture_output=True, + timeout=timeout, + check=False, + ) + for text in (result.stdout, result.stderr): + if text.strip(): + log(text) + if result.returncode != 0: + raise E2EError( + f"Command failed ({result.returncode}): {shlex.join(argv)}\n" + f"{result.stderr or result.stdout}" + ) + return result.stdout diff --git a/integrations/agentbricks/tests/suites/harness/run.py b/integrations/agentbricks/tests/suites/harness/run.py new file mode 100644 index 000000000..c73b5aa00 --- /dev/null +++ b/integrations/agentbricks/tests/suites/harness/run.py @@ -0,0 +1,60 @@ +"""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 + +FRAMEWORKS = ("langgraph",) +AUTHORING_PATHS = ("cli", "direct") +E2E_MODEL = "system.ai.gpt-5-2" +TOOL_RESOURCE_PREFIX = "agentbricks-tool-" +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 + app_auth_profile: str | None + # Operator override; None lets the workspace pick the first running warehouse. + warehouse_id: str | None + preprovisioned_app_catalog_access: bool + bridge_sha: str | None + wheel: pathlib.Path | None diff --git a/integrations/agentbricks/tests/target_workspace.py b/integrations/agentbricks/tests/suites/harness/target.py similarity index 85% rename from integrations/agentbricks/tests/target_workspace.py rename to integrations/agentbricks/tests/suites/harness/target.py index 079a1b636..6590cfd2f 100644 --- a/integrations/agentbricks/tests/target_workspace.py +++ b/integrations/agentbricks/tests/suites/harness/target.py @@ -9,8 +9,10 @@ from collections.abc import Mapping, Sequence from typing import Protocol -from common import RunConfig -from workspace_client import Workspace +import pytest + +from harness.run import RunConfig +from harness.workspace import Workspace WORKSPACE_KINDS = ("live", "fake") @@ -79,6 +81,14 @@ 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(name: str, profile: str | None, run_config: RunConfig) -> TargetWorkspace: if name == "fake": return FakeWorkspace(run_config) diff --git a/integrations/agentbricks/tests/workspace_client.py b/integrations/agentbricks/tests/suites/harness/workspace.py similarity index 82% rename from integrations/agentbricks/tests/workspace_client.py rename to integrations/agentbricks/tests/suites/harness/workspace.py index bc448a853..bd60c088e 100644 --- a/integrations/agentbricks/tests/workspace_client.py +++ b/integrations/agentbricks/tests/suites/harness/workspace.py @@ -1,6 +1,8 @@ """Databricks SDK helpers for the E2E suites and fixtures: the workspace, which is not under test. -Everything here goes through the SDK except ``app_logs``, since the SDK has no Apps log API. +``Workspace`` holds one operation per method (create, read, delete one thing); deciding what to +delete and in which order is ``harness.cleanup``. Everything goes through the SDK except +``app_logs``, since the SDK has no Apps log API. """ from __future__ import annotations @@ -8,25 +10,25 @@ import dataclasses import datetime as dt import io -import json import pathlib import re import subprocess import time -from collections.abc import Callable, Collection, Sequence +from collections.abc import Collection, Sequence from typing import Any, cast -from common import TOOL_RESOURCE_PREFIX, E2EError, log, now from databricks.sdk import WorkspaceClient from databricks.sdk.errors import DatabricksError, NotFound from databricks.sdk.service.catalog import PermissionsChange, Privilege, SecurableType, VolumeType from databricks.sdk.service.sql import ExecuteStatementRequestOnWaitTimeout, State, StatementState +from harness.process import log +from harness.run import TOOL_RESOURCE_PREFIX, E2EError, now + GrantTuple = tuple[str, str, str, str] _BRANCH = re.compile(r"projects/[^/]+/branches/[^/]+") _RUNTIME_DATABASE_PREFIX = "runtime-" -_APPS_RECORD = "apps.jsonl" @dataclasses.dataclass(frozen=True) @@ -46,32 +48,6 @@ def merged(self, other: AppIdentity) -> AppIdentity: ) -def record_app_identity(output: pathlib.Path, app: str, identity: AppIdentity) -> None: - """Append one line for the controller's end-of-run sweep, which cannot see worker memory.""" - line = json.dumps({"app": app, **dataclasses.asdict(identity)}) + "\n" - # 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"), data.get("database_id") - ) - identities[data["app"]] = identities.get(data["app"], AppIdentity()).merged(identity) - except (ValueError, KeyError, TypeError, AttributeError): - continue - return identities - - class Workspace: def __init__( self, @@ -196,7 +172,7 @@ def upload(self, path: str, data: bytes) -> None: # 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 sweeper delete it later.""" + """Create a schema and return its full name; ``RemoveAfter`` lets a workspace janitor delete it later.""" info = self.client.schemas.create( name, catalog, @@ -532,106 +508,6 @@ def _wait(operation: Any) -> None: wait() -def _delete_each( - kind: str, find: Callable[[], list[str]], delete: Callable[[str], None] -) -> list[str]: - """Delete everything ``find`` returns, continuing past failures; returns one line per failure.""" - try: - names = find() - except Exception as exc: - return [f"{kind} lookup: {exc}"] - failures = [] - for name in names: - try: - delete(name) - except Exception as exc: - failures.append(f"{kind} {name}: {exc}") - return failures - - -def delete_lakebase_leftovers( - workspace: Workspace, - branch: str, - *, - principals: Collection[str], - database_ids: Collection[str] = (), -) -> list[str]: - """Delete runtime databases and then the principals' roles on one branch. - - A role that owns a database cannot be deleted, so databases go first. Only ``runtime-`` - databases are touched, and both kinds are matched by principal UUID, never by name pattern. - """ - return _delete_each( - "runtime database", - lambda: workspace.runtime_databases(branch, principals, database_ids), - workspace.delete_database, - ) + _delete_each( - "Lakebase role", - lambda: workspace.roles_of(branch, principals), - workspace.delete_role, - ) - - -def cleanup_app( - workspace: Workspace, - name: str, - *, - delete_store: Callable[[str, str], subprocess.CompletedProcess[str]] | None = None, - memory_store: str | None = None, - session_store: str | None = None, - has_app: bool = True, - runtime_store: bool = False, - identity: AppIdentity | None = None, -) -> list[str]: - """Delete one project's stores and, if it has an App, everything the App's deploy created. - - ``delete_store(kind, store)`` is the agentbricks CLI's store delete, which has no SDK surface. - ``identity`` is what was recorded at deploy time; it fills in whatever the live lookup misses. - Every step runs regardless of earlier failures except those that need the App gone. Returns a - description of each failure; an empty list means everything was deleted. - """ - failures: list[str] = [] - - def attempt(label: str, action: Callable[..., object], *args: Any) -> bool: - try: - action(*args) - except Exception as exc: - failures.append(f"{label}: {exc}") - return False - return True - - def delete_cli_store(kind: str, store: str) -> None: - assert delete_store is not None - result = delete_store(kind, store) - if result.returncode != 0: - raise E2EError((result.stderr or result.stdout).strip()) - - # Read before any delete: the principal and branch come from the App and its Runtime Store. - known = (identity or AppIdentity()).merged( - workspace.app_identity(name) if has_app else AppIdentity() - ) - for kind, store in (("memory", memory_store), ("sessions", session_store)): - if store and delete_store: - attempt(f"{kind} store {store}", delete_cli_store, kind, store) - if not has_app: - return failures - - attempt(f"catalog grant {name}", workspace.revoke_catalog_grant, name) - if runtime_store: - attempt(f"runtime store {name}", workspace.delete_runtime_store, name) - app_deleted = attempt(f"app {name}", workspace.delete_app, name) - if app_deleted: - attempt(f"deployment source {name}", workspace.delete_deployment_source, name) - if known.branch: - failures += delete_lakebase_leftovers( - workspace, - known.branch, - principals=[known.service_principal or ""], - database_ids=[known.database_id or ""], - ) - return failures - - # Grant helpers diff --git a/integrations/agentbricks/tests/suites/memory_store/test_memory_store.py b/integrations/agentbricks/tests/suites/memory_store/test_memory_store.py index 1a674ec8a..f93e5cdb0 100644 --- a/integrations/agentbricks/tests/suites/memory_store/test_memory_store.py +++ b/integrations/agentbricks/tests/suites/memory_store/test_memory_store.py @@ -6,7 +6,7 @@ from __future__ import annotations import pytest -from agentbricks_cli import AgentbricksCli +from harness.cli import AgentbricksCli @pytest.mark.usefixtures("workspace_client") diff --git a/integrations/agentbricks/tests/suites/session_store/test_session_store.py b/integrations/agentbricks/tests/suites/session_store/test_session_store.py index 2bd4ffa45..a15f6c713 100644 --- a/integrations/agentbricks/tests/suites/session_store/test_session_store.py +++ b/integrations/agentbricks/tests/suites/session_store/test_session_store.py @@ -7,7 +7,7 @@ from __future__ import annotations import pytest -from agentbricks_cli import AgentbricksCli +from harness.cli import AgentbricksCli @pytest.mark.usefixtures("workspace_client") diff --git a/integrations/agentbricks/tests/suites/tools/conftest.py b/integrations/agentbricks/tests/suites/tools/conftest.py index 135eb2639..e99cb73bb 100644 --- a/integrations/agentbricks/tests/suites/tools/conftest.py +++ b/integrations/agentbricks/tests/suites/tools/conftest.py @@ -1,8 +1,8 @@ """Tools suite: every test runs once per authoring path, over the tools the target workspace can support. Each test gets its own temporary UC schema (``ab_e2e__``) holding the function and -volume its tools use, and deletes it, grants included, in its own teardown. The controller sweeps -any a crashed test left behind by the run's prefix (see ``provisioning``). +volume its tools use, and deletes it, grants included, in its own teardown. Leftover cleanup deletes +any a crashed test left behind, by the run's prefix (see ``harness.cleanup``). """ from __future__ import annotations @@ -14,17 +14,10 @@ from collections.abc import Iterator import pytest -from common import ( - AUTHORING_PATHS, - E2EError, - RunConfig, - log, - now, - scratch_schema_prefix, -) -from provisioning import require_workspace -from target_workspace import TargetWorkspace -from tools import TOOL_DEFINITIONS, Tool +from harness.process import log +from harness.run import AUTHORING_PATHS, E2EError, RunConfig, now, scratch_schema_prefix +from harness.target import TargetWorkspace, require_workspace +from tool_definitions import TOOL_DEFINITIONS, Tool from uc_objects import UcFunction, create_uc_function @@ -40,7 +33,7 @@ def scratch_schema( """``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 separate sweeper delete schemas of runs that crashed before any cleanup. + # 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 diff --git a/integrations/agentbricks/tests/suites/tools/test_deploy.py b/integrations/agentbricks/tests/suites/tools/test_deploy.py index 0b03d41cc..16277909a 100644 --- a/integrations/agentbricks/tests/suites/tools/test_deploy.py +++ b/integrations/agentbricks/tests/suites/tools/test_deploy.py @@ -8,9 +8,9 @@ from __future__ import annotations -from agentbricks_cli import AgentbricksCli -from tools import Tool, assert_endpoint_invoke, bind, invoke_and_check -from workspace_client import Workspace +from harness.cli import AgentbricksCli +from harness.workspace import Workspace +from tool_definitions import Tool, assert_endpoint_invoke, bind, invoke_and_check def test_deploy( diff --git a/integrations/agentbricks/tests/suites/tools/test_dev.py b/integrations/agentbricks/tests/suites/tools/test_dev.py index 7dd2ec3ae..078105934 100644 --- a/integrations/agentbricks/tests/suites/tools/test_dev.py +++ b/integrations/agentbricks/tests/suites/tools/test_dev.py @@ -8,8 +8,8 @@ from __future__ import annotations import pytest -from agentbricks_cli import AgentbricksCli -from tools import ( +from harness.cli import AgentbricksCli +from tool_definitions import ( Tool, assert_endpoint_invoke, bind, diff --git a/integrations/agentbricks/tests/suites/tools/test_grants.py b/integrations/agentbricks/tests/suites/tools/test_grants.py index eb4af2c3d..34b84166c 100644 --- a/integrations/agentbricks/tests/suites/tools/test_grants.py +++ b/integrations/agentbricks/tests/suites/tools/test_grants.py @@ -2,11 +2,11 @@ from __future__ import annotations -from agentbricks_cli import AgentbricksCli -from common import TOOL_RESOURCE_PREFIX -from tools import bind, uc_function_tool +from harness.cli import AgentbricksCli +from harness.run import TOOL_RESOURCE_PREFIX +from harness.workspace import Workspace +from tool_definitions import bind, uc_function_tool from uc_objects import UcFunction -from workspace_client import Workspace def test_grants( diff --git a/integrations/agentbricks/tests/suites/tools/test_selection.py b/integrations/agentbricks/tests/suites/tools/test_selection.py index 1dd4f49b1..5b4f5f567 100644 --- a/integrations/agentbricks/tests/suites/tools/test_selection.py +++ b/integrations/agentbricks/tests/suites/tools/test_selection.py @@ -6,10 +6,10 @@ from __future__ import annotations -from agentbricks_cli import AgentbricksCli -from tools import PYTHON_MARKER, UC_MARKER, Tool, bind +from harness.cli import AgentbricksCli +from harness.workspace import Workspace +from tool_definitions import PYTHON_MARKER, UC_MARKER, Tool, bind from uc_objects import UcFunction -from workspace_client import Workspace # Names the Unity Catalog function without its tool id, and the Python helper only to exclude it. PROMPT = ( diff --git a/integrations/agentbricks/tests/suites/tools/tools.py b/integrations/agentbricks/tests/suites/tools/tool_definitions.py similarity index 98% rename from integrations/agentbricks/tests/suites/tools/tools.py rename to integrations/agentbricks/tests/suites/tools/tool_definitions.py index 264e7b67e..767d823eb 100644 --- a/integrations/agentbricks/tests/suites/tools/tools.py +++ b/integrations/agentbricks/tests/suites/tools/tool_definitions.py @@ -6,9 +6,9 @@ import json from collections.abc import Callable, Mapping -from agentbricks_cli import Agent, AgentbricksCli, Project +from harness.cli import Agent, AgentbricksCli, Project +from harness.workspace import GrantTuple from uc_objects import UcFunction -from workspace_client import GrantTuple PYTHON_MARKER_FILE = "agent/tools/matrix_marker.py" PYTHON_MARKER_SOURCE = ( diff --git a/integrations/agentbricks/tests/suites/tools/uc_objects.py b/integrations/agentbricks/tests/suites/tools/uc_objects.py index 9f781f96a..b0b50ac69 100644 --- a/integrations/agentbricks/tests/suites/tools/uc_objects.py +++ b/integrations/agentbricks/tests/suites/tools/uc_objects.py @@ -5,8 +5,8 @@ import dataclasses import uuid -from common import E2EError -from workspace_client import Workspace +from harness.run import E2EError +from harness.workspace import Workspace NESTED_FUNCTION = "ab_nested_marker" FUNCTION = "ab_uc_marker" diff --git a/integrations/agentbricks/tests/suites/tracing/test_tracing.py b/integrations/agentbricks/tests/suites/tracing/test_tracing.py index 54a1b3663..ac1859a8f 100644 --- a/integrations/agentbricks/tests/suites/tracing/test_tracing.py +++ b/integrations/agentbricks/tests/suites/tracing/test_tracing.py @@ -6,9 +6,9 @@ from __future__ import annotations import pytest -from agentbricks_cli import AgentbricksCli -from common import TOOL_RESOURCE_PREFIX -from workspace_client import Workspace +from harness.cli import AgentbricksCli +from harness.run import TOOL_RESOURCE_PREFIX +from harness.workspace import Workspace def test_tracing_experiment_bound( diff --git a/integrations/agentbricks/tests/sweep.py b/integrations/agentbricks/tests/sweep.py deleted file mode 100644 index 327cb056f..000000000 --- a/integrations/agentbricks/tests/sweep.py +++ /dev/null @@ -1,146 +0,0 @@ -"""Delete whatever one E2E run leaked, inside pytest or as a separate step after it. - -Per-test teardown deletes what each test created; this catches what a crashed test or a failed -teardown left behind. pytest runs it at the end of a run, but a job timeout or cancellation kills -pytest before that, which is exactly when leaks happen. So CI picks the run id up front and runs -this again in an ``always()`` step, from ``integrations/agentbricks``: - - uv run pytest tests/suites/tools --run-id "$RUN_ID" ... - uv run python tests/sweep.py --run-id "$RUN_ID" --databricks-profile - -Everything is found by the run's name prefix (Apps, source folders, schemas) or by the App -identities workers recorded in ``/apps.jsonl`` (Lakebase databases and roles, which outlive -their App), so it touches nothing outside the run and is safe to run twice. -""" - -from __future__ import annotations - -import argparse -import os -import pathlib -import sys - -from common import ( - APP_PREFIX, - RUN_ID_PATTERN, - log, - project_prefix, - run_output_dir, - scratch_schema_prefix, -) -from workspace_client import ( - Workspace, - cleanup_app, - delete_lakebase_leftovers, - recorded_app_identities, -) - -DEFAULT_CATALOG = os.environ.get("AGENTBRICKS_E2E_CATALOG", "main") - - -def sweep_run(workspace: Workspace, run_id: str, output: pathlib.Path, catalog: str) -> list[str]: - """Delete everything the run left; returns a description of each failure.""" - failures: list[str] = [] - sweeps = ( - ("leaked Apps", lambda: _sweep_apps(workspace, run_id, output)), - ("leaked Lakebase databases and roles", lambda: _sweep_lakebase(workspace, output)), - ("leaked source folders", lambda: _sweep_deployment_sources(workspace, run_id)), - ("leaked schemas", lambda: _sweep_schemas(workspace, catalog, run_id)), - ) - for label, action in sweeps: - 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 _sweep_apps(workspace: Workspace, run_id: str, output: pathlib.Path) -> list[str]: - """Delete Apps whose per-test cleanup failed or never ran, found by the run's name prefix.""" - prefix = _app_prefix(run_id) - recorded = recorded_app_identities(output) - failures: list[str] = [] - for app in list(workspace.client.apps.list()): - if app.name and app.name.startswith(prefix): - log(f"sweeping leaked App {app.name}") - failures += cleanup_app( - workspace, app.name, runtime_store=True, identity=recorded.get(app.name) - ) - return failures - - -def _sweep_lakebase(workspace: Workspace, output: pathlib.Path) -> list[str]: - """Delete runtime databases and roles of this run's Apps, whether or not the App still exists.""" - recorded = list(recorded_app_identities(output).values()) - failures: list[str] = [] - for branch in sorted({item.branch for item in recorded if item.branch}): - failures += delete_lakebase_leftovers( - workspace, - 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 - ], - database_ids=[ - item.database_id for item in recorded if item.branch == branch and item.database_id - ], - ) - return failures - - -def _sweep_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"sweeping deployment source {path}") - try: - workspace.delete_workspace_path(path) - except Exception as exc: - failures.append(f"deployment source {path}: {exc}") - return failures - - -def _sweep_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"sweeping leaked schema {schema}") - try: - workspace.delete_schema(schema) - except Exception as exc: - failures.append(f"schema {schema}: {exc}") - return failures - - -def main(argv: list[str] | None = None) -> int: - parser = argparse.ArgumentParser(description="Delete whatever one Agent Bricks E2E run leaked.") - parser.add_argument("--run-id", required=True, help="The run id the pytest run was given.") - parser.add_argument("--databricks-profile", help="Databricks CLI profile; default is ambient.") - parser.add_argument("--catalog", default=DEFAULT_CATALOG, help="Catalog of the run's schemas.") - parser.add_argument( - "--output", - type=pathlib.Path, - help="The run's output dir, for apps.jsonl (default: the one pytest derives from the id).", - ) - args = parser.parse_args(argv) - if RUN_ID_PATTERN.fullmatch(args.run_id) is None: - parser.error("--run-id must be 1-8 lowercase letters or digits") - output = args.output or run_output_dir(args.run_id) - if not (output / "apps.jsonl").exists(): - # Without the record, Lakebase leftovers of Apps that are already gone can't be found. - log(f"no apps.jsonl in {output}; sweeping only what the run prefix finds") - failures = sweep_run(Workspace(args.databricks_profile), args.run_id, output, args.catalog) - for failure in failures: - log(f"cleanup warning | {failure}") - return 1 if failures else 0 - - -if __name__ == "__main__": - sys.exit(main()) From d6b490194d5964f1317659777fc8c3aeeb962555 Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Tue, 13 Oct 2026 00:23:55 +0000 Subject: [PATCH 11/45] [agentbricks] rename e2e harness process.py to logs.py and inline the wheel build Co-authored-by: Isaac --- .../agentbricks/tests/suites/README.md | 2 +- .../agentbricks/tests/suites/conftest.py | 18 +++++- .../tests/suites/harness/cleanup.py | 2 +- .../agentbricks/tests/suites/harness/cli.py | 2 +- .../agentbricks/tests/suites/harness/logs.py | 31 +++++++++ .../tests/suites/harness/process.py | 63 ------------------- .../tests/suites/harness/workspace.py | 2 +- .../tests/suites/tools/conftest.py | 2 +- 8 files changed, 51 insertions(+), 71 deletions(-) create mode 100644 integrations/agentbricks/tests/suites/harness/logs.py delete mode 100644 integrations/agentbricks/tests/suites/harness/process.py diff --git a/integrations/agentbricks/tests/suites/README.md b/integrations/agentbricks/tests/suites/README.md index 7db656a5b..ce1d732a9 100644 --- a/integrations/agentbricks/tests/suites/README.md +++ b/integrations/agentbricks/tests/suites/README.md @@ -49,7 +49,7 @@ Every resource a run creates carries its run id, and is deleted at three points: | `harness/cli.py` | The `agentbricks` CLI under test, and the agents it produces | | `harness/workspace.py` | Databricks SDK helpers for what is not under test: SQL, UC, Apps, grants | | `harness/cleanup.py` | Deleting what the suites create, per App and per run | -| `harness/process.py` | Logging and subprocess helpers | +| `harness/logs.py` | Test-log output; every command is echoed so a run's log replays it | | `/` | One folder per feature: `tools`, `memory_store`, `session_store`, `tracing` | `conftest.py` files are the only code pytest calls by name; everything in `harness/` is plain Python diff --git a/integrations/agentbricks/tests/suites/conftest.py b/integrations/agentbricks/tests/suites/conftest.py index e4edd643d..e1a255570 100644 --- a/integrations/agentbricks/tests/suites/conftest.py +++ b/integrations/agentbricks/tests/suites/conftest.py @@ -20,6 +20,7 @@ import os import pathlib import re +import subprocess import uuid from collections.abc import Iterator from typing import Any @@ -27,7 +28,7 @@ import pytest from harness.cleanup import delete_app_resources, delete_run_leftovers from harness.cli import AgentbricksCli -from harness.process import log, run_command +from harness.logs import log, log_command from harness.run import RUN_ID_PATTERN, E2EError, RunConfig, run_output_dir from harness.target import ( WORKSPACE_KINDS, @@ -169,11 +170,22 @@ def run_config(request: pytest.FixtureRequest) -> RunConfig: 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") - run_command(["uv", "build", "--wheel", "--out-dir", str(dist)], cwd=PACKAGE_DIR, timeout=600) + 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 dataclasses.replace(run, wheel=wheels[-1]) + return wheels[-1] # The workspace diff --git a/integrations/agentbricks/tests/suites/harness/cleanup.py b/integrations/agentbricks/tests/suites/harness/cleanup.py index cf8507e27..0a35f3080 100644 --- a/integrations/agentbricks/tests/suites/harness/cleanup.py +++ b/integrations/agentbricks/tests/suites/harness/cleanup.py @@ -20,7 +20,7 @@ from collections.abc import Callable, Collection from typing import Any -from harness.process import log +from harness.logs import log from harness.run import APP_PREFIX, E2EError, project_prefix, scratch_schema_prefix from harness.workspace import AppIdentity, Workspace diff --git a/integrations/agentbricks/tests/suites/harness/cli.py b/integrations/agentbricks/tests/suites/harness/cli.py index 7ad11b916..e7a66f7ba 100644 --- a/integrations/agentbricks/tests/suites/harness/cli.py +++ b/integrations/agentbricks/tests/suites/harness/cli.py @@ -32,7 +32,7 @@ import tomlkit from harness.cleanup import record_app_identity -from harness.process import last_lines, last_nonempty_line, log, log_command +from harness.logs import last_lines, last_nonempty_line, log, log_command from harness.run import APP_PREFIX, E2E_MODEL, FRAMEWORKS, E2EError, RunConfig, now, project_prefix from harness.target import TargetWorkspace from harness.workspace import AppIdentity, Workspace diff --git a/integrations/agentbricks/tests/suites/harness/logs.py b/integrations/agentbricks/tests/suites/harness/logs.py new file mode 100644 index 000000000..d196f41bf --- /dev/null +++ b/integrations/agentbricks/tests/suites/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/suites/harness/process.py b/integrations/agentbricks/tests/suites/harness/process.py deleted file mode 100644 index 9251ab5d3..000000000 --- a/integrations/agentbricks/tests/suites/harness/process.py +++ /dev/null @@ -1,63 +0,0 @@ -"""Logging and subprocess helpers: every command the suites run is echoed, so a log replays the run.""" - -from __future__ import annotations - -import pathlib -import shlex -import subprocess -import sys -from collections.abc import Mapping, Sequence - -from harness.run import E2EError - - -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" - - -def run_command( - argv: Sequence[str], - *, - cwd: pathlib.Path | None = None, - env: Mapping[str, str] | None = None, - timeout: float = 300, -) -> str: - """Run a command, echo its output, and return stdout; a non-zero exit is a ``E2EError``.""" - log_command(argv, cwd) - result = subprocess.run( - list(argv), - cwd=cwd, - env=None if env is None else dict(env), - text=True, - capture_output=True, - timeout=timeout, - check=False, - ) - for text in (result.stdout, result.stderr): - if text.strip(): - log(text) - if result.returncode != 0: - raise E2EError( - f"Command failed ({result.returncode}): {shlex.join(argv)}\n" - f"{result.stderr or result.stdout}" - ) - return result.stdout diff --git a/integrations/agentbricks/tests/suites/harness/workspace.py b/integrations/agentbricks/tests/suites/harness/workspace.py index bd60c088e..86fd86d8b 100644 --- a/integrations/agentbricks/tests/suites/harness/workspace.py +++ b/integrations/agentbricks/tests/suites/harness/workspace.py @@ -22,7 +22,7 @@ from databricks.sdk.service.catalog import PermissionsChange, Privilege, SecurableType, VolumeType from databricks.sdk.service.sql import ExecuteStatementRequestOnWaitTimeout, State, StatementState -from harness.process import log +from harness.logs import log from harness.run import TOOL_RESOURCE_PREFIX, E2EError, now GrantTuple = tuple[str, str, str, str] diff --git a/integrations/agentbricks/tests/suites/tools/conftest.py b/integrations/agentbricks/tests/suites/tools/conftest.py index e99cb73bb..2c358076d 100644 --- a/integrations/agentbricks/tests/suites/tools/conftest.py +++ b/integrations/agentbricks/tests/suites/tools/conftest.py @@ -14,7 +14,7 @@ from collections.abc import Iterator import pytest -from harness.process import log +from harness.logs import log from harness.run import AUTHORING_PATHS, E2EError, RunConfig, now, scratch_schema_prefix from harness.target import TargetWorkspace, require_workspace from tool_definitions import TOOL_DEFINITIONS, Tool From fe2628ccf2b7f01c3029978eb443c9511498e6fe Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Tue, 13 Oct 2026 17:14:19 +0000 Subject: [PATCH 12/45] [agentbricks] disable the coding-agent hint in e2e CLI subprocesses Co-authored-by: Isaac --- integrations/agentbricks/tests/suites/harness/cli.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/integrations/agentbricks/tests/suites/harness/cli.py b/integrations/agentbricks/tests/suites/harness/cli.py index e7a66f7ba..cc47964a5 100644 --- a/integrations/agentbricks/tests/suites/harness/cli.py +++ b/integrations/agentbricks/tests/suites/harness/cli.py @@ -138,7 +138,9 @@ def __init__(self, target_workspace: TargetWorkspace, run_config: RunConfig): self.run_config = run_config self.logs_dir = run_config.output / "logs" self.projects: list[Project] = [] - self._env = {**os.environ, **target_workspace.env} + # 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. + self._env = {**os.environ, **target_workspace.env, "AGENTBRICKS_DISABLE_AGENT_HINT": "1"} self._sequence = 0 # Process plumbing From a2ade45244f8f250d230a6873d7c06b0d93274d6 Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Tue, 13 Oct 2026 17:43:34 +0000 Subject: [PATCH 13/45] [agentbricks] shorten the e2e UC marker function name to fit the MCP tool-name limit Co-authored-by: Isaac --- integrations/agentbricks/tests/suites/tools/uc_objects.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/integrations/agentbricks/tests/suites/tools/uc_objects.py b/integrations/agentbricks/tests/suites/tools/uc_objects.py index b0b50ac69..2d732ef7d 100644 --- a/integrations/agentbricks/tests/suites/tools/uc_objects.py +++ b/integrations/agentbricks/tests/suites/tools/uc_objects.py @@ -9,7 +9,7 @@ from harness.workspace import Workspace NESTED_FUNCTION = "ab_nested_marker" -FUNCTION = "ab_uc_marker" +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 From 516934187ac2a69e68134fdac28f81c30aa91b6b Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Tue, 13 Oct 2026 18:32:22 +0000 Subject: [PATCH 14/45] [agentbricks] retry e2e endpoint invoke on gateway errors, add --genie-space-id, drop per-option env vars Co-authored-by: Isaac --- .../agentbricks/tests/suites/README.md | 7 +++-- .../agentbricks/tests/suites/conftest.py | 16 ++++++++---- .../agentbricks/tests/suites/harness/cli.py | 26 ++++++++++++++++--- .../tests/suites/tools/conftest.py | 7 +++-- .../tests/suites/tools/test_selection.py | 6 ++--- .../tests/suites/tools/tool_definitions.py | 14 +++++++--- 6 files changed, 56 insertions(+), 20 deletions(-) diff --git a/integrations/agentbricks/tests/suites/README.md b/integrations/agentbricks/tests/suites/README.md index ce1d732a9..c2d13da0b 100644 --- a/integrations/agentbricks/tests/suites/README.md +++ b/integrations/agentbricks/tests/suites/README.md @@ -15,7 +15,10 @@ uv run pytest tests/suites -n 6 --dist loadgroup \ - Narrow it to a suite (`tests/suites/tools`), a file, or a test (`"...test_deploy.py::test_deploy[cli]"`). - `-n 6 --dist loadgroup` runs tests in parallel; `loadgroup` keeps the dev tests on one worker, because local dev servers share a proxy port. -- `AGENTBRICKS_E2E_GENIE_SPACE_ID=` adds the Genie tool; without it that tool is skipped. +- 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 "`. - `--workspace fake` runs against no workspace; everything that needs one is skipped. - `--keep-resources` keeps leftovers for debugging; each test still deletes its own resources. @@ -37,7 +40,7 @@ Every resource a run creates carries its run id, and is deleted at three points: uv run pytest tests/suites --cleanup-only --run-id --databricks-profile --catalog ``` - CI passes its own `--run-id` (or `AGENTBRICKS_E2E_RUN_ID`) and runs this in an `always()` step. + CI passes its own `--run-id` and runs this in an `always()` step. ## Layout diff --git a/integrations/agentbricks/tests/suites/conftest.py b/integrations/agentbricks/tests/suites/conftest.py index e1a255570..2d76185fe 100644 --- a/integrations/agentbricks/tests/suites/conftest.py +++ b/integrations/agentbricks/tests/suites/conftest.py @@ -66,8 +66,8 @@ def pytest_addoption(parser: pytest.Parser) -> None: ) group.addoption( "--catalog", - default=os.environ.get("AGENTBRICKS_E2E_CATALOG", "main"), - help="Existing catalog for temporary test schemas (AGENTBRICKS_E2E_CATALOG, default main).", + default="main", + help="Existing catalog for temporary test schemas (default: main).", ) group.addoption("--bridge-sha", help="Immutable bridge commit for generated App dependencies.") group.addoption( @@ -75,7 +75,14 @@ def pytest_addoption(parser: pytest.Parser) -> None: 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( + "--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( "--preprovisioned-app-catalog-access", action="store_true", @@ -83,9 +90,8 @@ def pytest_addoption(parser: pytest.Parser) -> None: ) group.addoption( "--run-id", - default=os.environ.get("AGENTBRICKS_E2E_RUN_ID"), help="Names every resource of the run, so --cleanup-only can delete what a killed run left " - "(AGENTBRICKS_E2E_RUN_ID; default: random). 1-8 lowercase letters or digits.", + "(default: random). 1-8 lowercase letters or digits.", ) group.addoption( "--keep-resources", diff --git a/integrations/agentbricks/tests/suites/harness/cli.py b/integrations/agentbricks/tests/suites/harness/cli.py index cc47964a5..62f51f30d 100644 --- a/integrations/agentbricks/tests/suites/harness/cli.py +++ b/integrations/agentbricks/tests/suites/harness/cli.py @@ -130,6 +130,9 @@ def warm_up(self, deadline_seconds: float = 180.0) -> None: raise +_GATEWAY_UNAVAILABLE = re.compile(r"HTTP 50[234]\b") + + class AgentbricksCli: """Authors, runs and deploys projects with the CLI against a target workspace.""" @@ -186,10 +189,27 @@ def endpoint_invoke(self, agent: Agent, prompt: str) -> str: "input": [{"role": "user", "content": prompt}], } target = [agent.app_name] if agent.app_name else ["--url", agent.base_url, "--no-auth"] - result = self.cli( - "endpoint", "invoke", *target, "--path", "/api/invocations", "--json", json.dumps(body) + argv = ( + "endpoint", + "invoke", + *target, + "--path", + "/api/invocations", + "--json", + json.dumps(body), ) - return result.stdout + for attempt in range(1, 4): + result = self.cli(*argv, check=False) + if result.returncode == 0: + return result.stdout + output = result.stderr or result.stdout + # Only the Apps gateway's "App not up" answers are retried (a staging App can restart + # under the test); any other failure is the CLI's and fails the test. + if attempt == 3 or _GATEWAY_UNAVAILABLE.search(output) is None: + raise E2EError(f"endpoint invoke failed ({result.returncode}): {output}") + log(f"attempt {attempt}/3 | endpoint invoke | App unavailable; retrying") + time.sleep(15) + raise AssertionError("unreachable") def _run_long(self, label: str, argv: Sequence[str], *, timeout: float) -> None: log_command(argv) diff --git a/integrations/agentbricks/tests/suites/tools/conftest.py b/integrations/agentbricks/tests/suites/tools/conftest.py index 2c358076d..518c91f05 100644 --- a/integrations/agentbricks/tests/suites/tools/conftest.py +++ b/integrations/agentbricks/tests/suites/tools/conftest.py @@ -8,7 +8,6 @@ from __future__ import annotations import datetime as dt -import os import re import uuid from collections.abc import Iterator @@ -51,11 +50,11 @@ def uc_function(scratch_schema: str, target_workspace: TargetWorkspace) -> UcFun @pytest.fixture(scope="session") -def genie_space(target_workspace: TargetWorkspace) -> str: +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 = os.environ.get("AGENTBRICKS_E2E_GENIE_SPACE_ID") + space_id = request.config.getoption("genie_space_id") if not space_id: - pytest.skip("no Genie space; set AGENTBRICKS_E2E_GENIE_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) diff --git a/integrations/agentbricks/tests/suites/tools/test_selection.py b/integrations/agentbricks/tests/suites/tools/test_selection.py index 5b4f5f567..bb600d7ca 100644 --- a/integrations/agentbricks/tests/suites/tools/test_selection.py +++ b/integrations/agentbricks/tests/suites/tools/test_selection.py @@ -8,7 +8,7 @@ from harness.cli import AgentbricksCli from harness.workspace import Workspace -from tool_definitions import PYTHON_MARKER, UC_MARKER, Tool, bind +from tool_definitions import PYTHON_MARKER, UC_MARKER, Tool, bind, excerpt from uc_objects import UcFunction # Names the Unity Catalog function without its tool id, and the Python helper only to exclude it. @@ -37,7 +37,7 @@ def test_selects_uc_function_not_python_marker( workspace_client.grant_transitive(app, uc_function.nested_function) serialized = agent.invoke(PROMPT, "deploy-selection") - assert UC_MARKER in serialized, f"The UC function tool was not called: {serialized[:2000]}" + assert UC_MARKER in serialized, f"The UC function tool was not called: {excerpt(serialized)}" assert PYTHON_MARKER not in serialized, ( - f"The Python marker tool was called too: {serialized[:2000]}" + f"The Python marker tool was called too: {excerpt(serialized)}" ) diff --git a/integrations/agentbricks/tests/suites/tools/tool_definitions.py b/integrations/agentbricks/tests/suites/tools/tool_definitions.py index 767d823eb..e2a7b4513 100644 --- a/integrations/agentbricks/tests/suites/tools/tool_definitions.py +++ b/integrations/agentbricks/tests/suites/tools/tool_definitions.py @@ -55,9 +55,17 @@ class ToolDefinition: build: Callable[..., 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 contains(marker: str) -> Check: def check(serialized: str) -> None: - assert marker in serialized, f"Missing exact marker {marker!r}: {serialized[:2000]}" + assert marker in serialized, f"Missing exact marker {marker!r}: {excerpt(serialized)}" return check @@ -66,7 +74,7 @@ def check_web_search(serialized: str) -> None: lowered = serialized.lower() tool_evidence = any(v in lowered for v in ("web_search", "web search", "search")) assert tool_evidence and "https" in lowered and len(serialized) >= 80, ( - f"Missing web-search execution/result evidence: {serialized[:2000]}" + f"Missing web-search execution/result evidence: {excerpt(serialized)}" ) @@ -75,7 +83,7 @@ def check_genie(serialized: str) -> None: assert ( any(v in lowered for v in ("genie_ask", "genie", "conversation_id")) and len(serialized) >= 80 - ), f"Missing Genie execution/result evidence: {serialized[:2000]}" + ), f"Missing Genie execution/result evidence: {excerpt(serialized)}" def sandbox_tool(uc_function: UcFunction) -> Tool: From cb2ad6d4752e0f93900f7d5f34cfb281977b0173 Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Tue, 13 Oct 2026 20:59:16 +0000 Subject: [PATCH 15/45] [agentbricks] wait for e2e tool grants to take effect and run the sandbox with default user auth Co-authored-by: Isaac --- .../tests/suites/harness/workspace.py | 23 +++++++++++++++++++ .../tests/suites/tools/test_deploy.py | 8 +++++++ .../tests/suites/tools/tool_definitions.py | 12 ++++++---- 3 files changed, 38 insertions(+), 5 deletions(-) diff --git a/integrations/agentbricks/tests/suites/harness/workspace.py b/integrations/agentbricks/tests/suites/harness/workspace.py index 86fd86d8b..6a846127e 100644 --- a/integrations/agentbricks/tests/suites/harness/workspace.py +++ b/integrations/agentbricks/tests/suites/harness/workspace.py @@ -257,6 +257,29 @@ def granted(self, app_name: str, grant: GrantTuple) -> bool: """Whether Agent Bricks granted the App ``grant``, as (kind, name, type, permission).""" return grant in app_resource_tuples(tool_resources(self.app(app_name))) + def wait_until_effective(self, app_name: str, grant: GrantTuple, timeout: float = 120) -> float: + """Wait until a UC ``grant`` deploy added is effective for the App; returns the seconds waited. + + Deploy reconciles the App's UC resources before returning, but the privilege can become + effective for its principal a little later; a tool called in between is denied. + """ + kind, full_name, securable_type, privilege = grant + if kind != "uc_securable": + return 0.0 + principal = self._principal(self.app(app_name)) + started = time.monotonic() + while True: + held = effective_privileges( + self.client, SecurableType(securable_type), full_name, principal + ) + if privilege in held or "ALL_PRIVILEGES" in held: + return time.monotonic() - started + if time.monotonic() - started > timeout: + raise E2EError( + f"{privilege} on {full_name} not effective for {app_name} after {timeout:.0f}s: {held}" + ) + time.sleep(5) + def granted_tuples(self, app_name: str) -> set[GrantTuple]: return app_resource_tuples(tool_resources(self.app(app_name))) diff --git a/integrations/agentbricks/tests/suites/tools/test_deploy.py b/integrations/agentbricks/tests/suites/tools/test_deploy.py index 16277909a..9cd4b0b6c 100644 --- a/integrations/agentbricks/tests/suites/tools/test_deploy.py +++ b/integrations/agentbricks/tests/suites/tools/test_deploy.py @@ -9,6 +9,7 @@ from __future__ import annotations from harness.cli import AgentbricksCli +from harness.logs import log from harness.workspace import Workspace from tool_definitions import Tool, assert_endpoint_invoke, bind, invoke_and_check @@ -30,6 +31,11 @@ def assert_granted(tool: Tool) -> None: f"{sorted(workspace_client.granted_tuples(app))}" ) + def wait_for_grant(tool: Tool) -> None: + if tool.grant is not None: + waited = workspace_client.wait_until_effective(app, tool.grant) + log(f"{tool.name}: {tool.grant[3]} effective for {app} after {waited:.0f}s") + def grant_transitive_if_needed(tool: Tool) -> None: # The nested function is deliberately outside Agent Bricks' grants; the tool only works # once this manual grant is in place. @@ -41,6 +47,7 @@ def grant_transitive_if_needed(tool: Tool) -> None: with subtests.test(tool=tool.name): bind(agentbricks_cli, project, tool) agent = agentbricks_cli.deploy(project, app) + wait_for_grant(tool) grant_transitive_if_needed(tool) invoke_and_check(agent, tool) for added in tools[: index + 1]: @@ -53,6 +60,7 @@ def grant_transitive_if_needed(tool: Tool) -> None: agent = agentbricks_cli.deploy(project, app) for tool in tools: with subtests.test(tool=tool.name): + wait_for_grant(tool) grant_transitive_if_needed(tool) invoke_and_check(agent, tool) assert_granted(tool) diff --git a/integrations/agentbricks/tests/suites/tools/tool_definitions.py b/integrations/agentbricks/tests/suites/tools/tool_definitions.py index e2a7b4513..96a2857c2 100644 --- a/integrations/agentbricks/tests/suites/tools/tool_definitions.py +++ b/integrations/agentbricks/tests/suites/tools/tool_definitions.py @@ -95,9 +95,11 @@ def sandbox_tool(uc_function: UcFunction) -> Tool: "do not fabricate them." ), check=contains(uc_function.volume_marker), - add_args=("sandbox", "--scope", f"volume:{uc_function.volume}", "--auth", "app"), + # The CLI's default auth: the sandbox reads the volume as the user invoking the agent, so + # deploy adds no App grant. A manifest entry without `auth` would run as the App instead. + add_args=("sandbox", "--scope", f"volume:{uc_function.volume}"), toml=( - '[[tools]]\nid = "sandbox"\n' + '[[tools]]\nid = "sandbox"\nauth = "user"\n' 'source = { kind = "sandbox", service = "system.ai.sandbox" }\n' "policy = { downscope = [\n" f' {{ resource = "volume:{uc_function.volume}", permission = "read_only" }},\n' @@ -105,8 +107,8 @@ def sandbox_tool(uc_function: UcFunction) -> Tool: ), files={}, tool_id="sandbox", - app_auth=True, - grant=("uc_securable", uc_function.volume, "VOLUME", "READ_VOLUME"), + app_auth=False, + grant=None, ) @@ -121,7 +123,7 @@ def web_search_tool(**_resources: object) -> Tool: check=check_web_search, add_args=("mcp", "system.ai.web_search", "--auth", "app"), toml=( - '[[tools]]\nid = "web_search"\n' + '[[tools]]\nid = "web_search"\nauth = "app"\n' 'source = { kind = "mcp", service = "system.ai.web_search" }\n' ), files={}, From a28ea5a206c13fc3ace001304340b138dafe69e1 Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Wed, 14 Oct 2026 18:17:07 +0000 Subject: [PATCH 16/45] [agentbricks] E2E suites: run CLI steps as command strings; dev/deploy off the CLI wrapper Co-authored-by: Isaac --- .../agentbricks/tests/suites/README.md | 2 +- .../agentbricks/tests/suites/conftest.py | 34 +- .../tests/suites/harness/cleanup.py | 6 +- .../agentbricks/tests/suites/harness/cli.py | 487 +++++++----------- .../tests/suites/harness/workspace.py | 6 +- .../suites/memory_store/test_memory_store.py | 9 +- .../session_store/test_session_store.py | 11 +- .../tests/suites/tools/test_deploy.py | 18 +- .../tests/suites/tools/test_dev.py | 32 +- .../tests/suites/tools/test_grants.py | 7 +- .../tests/suites/tools/test_selection.py | 9 +- .../tests/suites/tools/tool_definitions.py | 62 ++- .../tests/suites/tracing/test_tracing.py | 5 +- 13 files changed, 322 insertions(+), 366 deletions(-) diff --git a/integrations/agentbricks/tests/suites/README.md b/integrations/agentbricks/tests/suites/README.md index c2d13da0b..b0cfc037d 100644 --- a/integrations/agentbricks/tests/suites/README.md +++ b/integrations/agentbricks/tests/suites/README.md @@ -49,7 +49,7 @@ Every resource a run creates carries its run id, and is deleted at three points: | `conftest.py` | All pytest wiring: options, hooks, and the fixtures every suite shares | | `harness/run.py` | `RunConfig`, the run id, and the name prefixes every resource carries | | `harness/target.py` | The `--workspace live\|fake` seam the CLI and fixtures talk to | -| `harness/cli.py` | The `agentbricks` CLI under test, and the agents it produces | +| `harness/cli.py` | Runs `agentbricks` commands (`cli.run("tools add ...")`), and the endpoints of the agents they start | | `harness/workspace.py` | Databricks SDK helpers for what is not under test: SQL, UC, Apps, grants | | `harness/cleanup.py` | Deleting what the suites create, per App and per run | | `harness/logs.py` | Test-log output; every command is echoed so a run's log replays it | diff --git a/integrations/agentbricks/tests/suites/conftest.py b/integrations/agentbricks/tests/suites/conftest.py index 2d76185fe..7504bda90 100644 --- a/integrations/agentbricks/tests/suites/conftest.py +++ b/integrations/agentbricks/tests/suites/conftest.py @@ -26,7 +26,7 @@ from typing import Any import pytest -from harness.cleanup import delete_app_resources, delete_run_leftovers +from harness.cleanup import delete_app_resources, delete_run_leftovers, record_app_identity from harness.cli import AgentbricksCli from harness.logs import log, log_command from harness.run import RUN_ID_PATTERN, E2EError, RunConfig, run_output_dir @@ -37,7 +37,7 @@ make_target_workspace, require_workspace, ) -from harness.workspace import Workspace +from harness.workspace import AppIdentity, Workspace PACKAGE_DIR = pathlib.Path(__file__).resolve().parents[2] _RUN_KEY = pytest.StashKey[RunConfig]() @@ -232,7 +232,7 @@ def agentbricks_cli( request: pytest.FixtureRequest, target_workspace: TargetWorkspace, run_config: RunConfig ) -> Iterator[AgentbricksCli]: """The CLI wrapper for one test; deletes every App and store the test created when it ends.""" - cli = AgentbricksCli(target_workspace, run_config) + cli = AgentbricksCli(target_workspace, run_config, request.node.name) yield cli if cli.projects and target_workspace.client is not None: _delete_projects(cli, target_workspace.client, failed=request.node.nodeid in _failed_tests) @@ -240,27 +240,35 @@ def agentbricks_cli( def _delete_projects(cli: AgentbricksCli, workspace: Workspace, *, failed: bool) -> None: for project in cli.projects: - if failed and project.app_registered: + name = project.app_name + # Whatever deploy got as far as creating; read before anything is deleted, because the + # principal and branch come from the App and its Runtime Store. + identity = workspace.app_identity(name) + deployed = identity != AppIdentity() + if deployed: + # Leftover cleanup retries the Lakebase deletes from this record if they fail here. + record_app_identity(cli.run_config.output, name, identity) + if failed and identity.service_principal: # The App is deleted next, so a failing test must capture its logs while it exists. - workspace.app_logs( - project.app_name, cli.logs_dir / f"deploy-runtime-{project.app_name}.log" - ) + workspace.app_logs(name, cli.logs_dir / f"deploy-runtime-{name}.log") try: failures = delete_app_resources( workspace, - project.app_name, - delete_store=cli.delete_store, + name, + delete_store=lambda kind, store: cli.run( + f"{kind} stores delete {store} --yes", timeout=600, check=False + ), memory_store=project.memory_store_name, session_store=project.session_store_name, - has_app=project.app_registered, + has_app=deployed, # Every managed deploy creates a Runtime Store, with or without memory/session stores. - runtime_store=project.app_registered, - identity=project.identity, + runtime_store=deployed, + identity=identity, ) except Exception as exc: failures = [str(exc)] for failure in failures: - log(f"cleanup warning | {project.app_name} | {failure}") + log(f"cleanup warning | {name} | {failure}") # Deleting what the whole run left behind diff --git a/integrations/agentbricks/tests/suites/harness/cleanup.py b/integrations/agentbricks/tests/suites/harness/cleanup.py index 0a35f3080..a27affb7e 100644 --- a/integrations/agentbricks/tests/suites/harness/cleanup.py +++ b/integrations/agentbricks/tests/suites/harness/cleanup.py @@ -7,7 +7,7 @@ - ``delete_run_leftovers``: whatever a crashed test or a failed teardown left. It runs once after every live run, and on its own with ``pytest tests/suites --cleanup-only --run-id `` after a run that was killed before it could. It finds resources by the run's name prefix (Apps, source - folders, schemas) and by the App identities recorded at deploy time in ``/apps.jsonl`` + folders, schemas) and by the App identities per-test cleanup recorded in ``/apps.jsonl`` (Lakebase databases and roles, which outlive their App), so it touches nothing outside the run. """ @@ -27,7 +27,7 @@ _APPS_RECORD = "apps.jsonl" -# Identities recorded at deploy time +# Identities recorded by per-test cleanup def record_app_identity(output: pathlib.Path, app: str, identity: AppIdentity) -> None: @@ -113,7 +113,7 @@ def delete_app_resources( """Delete one project's stores and, if it has an App, everything the App's deploy created. ``delete_store(kind, store)`` is the agentbricks CLI's store delete, which has no SDK surface. - ``identity`` is what was recorded at deploy time; it fills in whatever the live lookup misses. + ``identity`` is what was recorded earlier; it fills in whatever the live lookup misses. Every step runs regardless of earlier failures except those that need the App gone. Returns a description of each failure; an empty list means everything was deleted. """ diff --git a/integrations/agentbricks/tests/suites/harness/cli.py b/integrations/agentbricks/tests/suites/harness/cli.py index 62f51f30d..8cf314b76 100644 --- a/integrations/agentbricks/tests/suites/harness/cli.py +++ b/integrations/agentbricks/tests/suites/harness/cli.py @@ -1,9 +1,10 @@ """The ``agentbricks`` CLI and the agents it produces: the system under test. -Every call here is a real ``agentbricks`` subprocess, so a regression in the CLI fails the test -that makes it. The only non-subprocess part is ``Agent.invoke``, which POSTs to a running agent. -Nothing in this module knows about any particular suite; tests say what to bind. It records the -projects it creates in ``projects`` but deletes nothing: the ``agentbricks_cli`` fixture does. +Every CLI 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. The only +non-subprocess part is ``AgentEndpoint.invoke``, which POSTs to a running agent. Nothing in this +module knows about any particular suite. It records the projects it creates in ``projects`` but +deletes nothing: the ``agentbricks_cli`` fixture does. """ from __future__ import annotations @@ -31,16 +32,16 @@ import tomli import tomlkit -from harness.cleanup import record_app_identity from harness.logs import last_lines, last_nonempty_line, log, log_command from harness.run import APP_PREFIX, E2E_MODEL, FRAMEWORKS, E2EError, RunConfig, now, project_prefix from harness.target import TargetWorkspace -from harness.workspace import AppIdentity, Workspace +from harness.workspace import Workspace -_STORE_KINDS = {"memory": ("memory", "stores"), "sessions": ("sessions", "stores")} _DIRECT_HEADER = ( 'schema_version = 1\n\n[agent]\nframework = "{framework}"\nserver = "agentbricks"\n' ) +# The Apps gateway's answers while an App is not up, e.g. restarting after a redeploy. +_GATEWAY_UNAVAILABLE = re.compile(r"HTTP 50[234]\b") class ToolSpec(Protocol): @@ -62,31 +63,104 @@ class Project: memory_store_name: str | None = None session_store_name: str | None = None experiment_name: str | None = None - # Set only once this project registered the App, so cleanup never deletes one it did not create. - app_registered: bool = False - environment_prepared: bool = False - # Recorded at deploy time: cleanup cannot look this up once the App or its store is gone. - identity: AppIdentity = dataclasses.field(default_factory=AppIdentity) + + def manifest(self) -> dict[str, Any]: + return tomli.loads((self.path / "agent.toml").read_text(encoding="utf-8")) + + 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_manifest(self, tool_specs: Sequence[ToolSpec]) -> None: + """Direct authoring: write agent.toml with exactly these tools plus a tracing binding.""" + sections: list[str] = [_DIRECT_HEADER.format(framework=FRAMEWORKS[0])] + sections.extend(spec.toml for spec in tool_specs if spec.toml) + # Follows default_experiment_name's shape so direct authoring carries the same tracing + # binding `agentbricks init` gives CLI authoring. + slug = re.sub(r"[^a-z0-9-]+", "-", self.path.name.lower()).strip("-") or "agent" + self.experiment_name = f"/Shared/agentbricks_traces/{slug}" + sections.append(f'[tracing]\nexperiment_name = "{self.experiment_name}"\n') + target = self.path / "agent.toml" + log(f"# write {target}: direct authoring; no agentbricks tools command") + target.write_text("\n".join(sections), encoding="utf-8") + for spec in tool_specs: + for relative_path, content in spec.files.items(): + self.write_file(relative_path, content) + + +@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 -class Agent: - """A running agent: a dev server or a deployed App.""" +class AgentEndpoint: + """Where a running agent answers, a dev server or a deployed App, and how to call it.""" def __init__( self, base_url: str, headers: Callable[[], dict[str, str]], - log_path: pathlib.Path, *, app_name: str | None = None, app: dict[str, Any] | None = None, ): self.base_url = base_url - self.log_path = log_path self.app_name = app_name self.app = app self._headers = headers + @classmethod + def local(cls, server: BackgroundCommand, port: int) -> AgentEndpoint: + """The `agentbricks dev` server on ``port``, once it answers; fails if the server exits.""" + started = time.monotonic() + next_tick = 60.0 + while True: + if server.process.poll() is not None: + raise E2EError( + f"{server.label} exited {server.process.returncode}: " + f"{last_lines(server.log_path, 30)}" + ) + try: + with urllib.request.urlopen(f"http://127.0.0.1:{port}/", timeout=5): + break + except urllib.error.HTTPError as exc: + if exc.code < 500: + break + except (urllib.error.URLError, TimeoutError): + pass + elapsed = time.monotonic() - started + if elapsed > 1200: + raise E2EError( + f"{server.label} did not become reachable: {last_lines(server.log_path, 30)}" + ) + if elapsed >= next_tick: + log( + f"tick {now():%H:%M} | {server.label} | starting | " + f"{last_nonempty_line(server.log_path)}" + ) + next_tick += 60 + time.sleep(5) + return cls(f"http://127.0.0.1:{port}", lambda: {}) + + @classmethod + def deployed(cls, workspace: Workspace, app_name: str) -> AgentEndpoint: + """The App once it is up and has answered once.""" + workspace.check_app_auth() + app = workspace.wait_for_app(app_name) + url = str(app.get("url") or "").rstrip("/") + if not url: + raise E2EError(f"App {app_name} has no URL: {app}") + endpoint = cls(url, lambda: workspace.app_headers, app_name=app_name, app=app) + endpoint.warm_up() + return endpoint + @property def runtime(self) -> str: return "deploy" if self.app_name else "dev" @@ -119,7 +193,8 @@ def invoke(self, prompt: str, label: str = "invoke") -> str: raise E2EError(f"{label} failed after 3 attempts: {last}") def warm_up(self, deadline_seconds: float = 180.0) -> None: - """Invoke until the App answers once. A new App's session store can drop the first calls.""" + """Invoke until the agent answers once. A new App's session store can drop the first calls, + and a redeployed one restarts.""" started = time.monotonic() while True: try: @@ -130,119 +205,130 @@ def warm_up(self, deadline_seconds: float = 180.0) -> None: raise -_GATEWAY_UNAVAILABLE = re.compile(r"HTTP 50[234]\b") - - class AgentbricksCli: - """Authors, runs and deploys projects with the CLI against a target workspace.""" + """Runs ``agentbricks`` commands against a target workspace and authors projects with them.""" - def __init__(self, target_workspace: TargetWorkspace, run_config: RunConfig): + def __init__(self, target_workspace: TargetWorkspace, run_config: RunConfig, test_name: str): self.target_workspace = target_workspace self.run_config = run_config - self.logs_dir = run_config.output / "logs" + # 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("-") + ) self.projects: list[Project] = [] # 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. self._env = {**os.environ, **target_workspace.env, "AGENTBRICKS_DISABLE_AGENT_HINT": "1"} self._sequence = 0 - # Process plumbing - - def argv(self, *args: str) -> list[str]: - return [_agentbricks_bin(), *self.target_workspace.cli_args, *args] - - def cli( + def run( self, - *args: str, - cwd: pathlib.Path | None = None, + command: str, + *, timeout: float = 300, check: bool = True, + retry_unavailable: bool = False, ) -> subprocess.CompletedProcess[str]: - argv = self.argv(*args) - log_command(argv, cwd) - result = subprocess.run( - argv, - cwd=cwd, - env=self._env, - text=True, - capture_output=True, - timeout=timeout, - check=False, - ) - if result.stdout.strip(): - log(result.stdout) - if result.stderr.strip(): - log(result.stderr) - if check and result.returncode != 0: - raise E2EError( - f"Command failed ({result.returncode}): {shlex.join(argv)}\n" - f"{result.stderr or result.stdout}" - ) - return result + """Run ``agentbricks `` to completion. - def endpoint_invoke(self, agent: Agent, prompt: str) -> str: - """`agentbricks endpoint invoke`, the call the CLI tells users to run after dev/deploy.""" - invocation_id = str(uuid.uuid4()) - body = { - "id": invocation_id, - "session_id": invocation_id, - "input": [{"role": "user", "content": prompt}], - } - target = [agent.app_name] if agent.app_name else ["--url", agent.base_url, "--no-auth"] - argv = ( - "endpoint", - "invoke", - *target, - "--path", - "/api/invocations", - "--json", - json.dumps(body), - ) + ``retry_unavailable`` retries only the Apps gateway's "App not up" answers (a staging App + can restart under the test); any other failure is the CLI's. + """ + argv = self._argv(command) for attempt in range(1, 4): - result = self.cli(*argv, check=False) - if result.returncode == 0: - return result.stdout + result = self._run_once(argv, self._label(command), timeout) output = result.stderr or result.stdout - # Only the Apps gateway's "App not up" answers are retried (a staging App can restart - # under the test); any other failure is the CLI's and fails the test. - if attempt == 3 or _GATEWAY_UNAVAILABLE.search(output) is None: - raise E2EError(f"endpoint invoke failed ({result.returncode}): {output}") - log(f"attempt {attempt}/3 | endpoint invoke | App unavailable; retrying") + if ( + result.returncode == 0 + or not retry_unavailable + or attempt == 3 + or _GATEWAY_UNAVAILABLE.search(output) is None + ): + break + log(f"attempt {attempt}/3 | {command.split()[0]} | App unavailable; retrying") time.sleep(15) - raise AssertionError("unreachable") + if check and result.returncode != 0: + raise E2EError(f"Command failed ({result.returncode}): {shlex.join(argv)}\n{output}") + return result - def _run_long(self, label: str, argv: Sequence[str], *, timeout: float) -> None: - log_command(argv) + @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) - started = time.monotonic() + log_command(argv) with log_path.open("w", encoding="utf-8") as log_file: process = subprocess.Popen( - list(argv), - env=self._env, + 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, grace=0) - raise E2EError(f"{label} timed out after {timeout:.0f}s; log: {log_path}") + raise E2EError(f"{label} timed out after {timeout:.0f}s; log: {stderr_path}") if elapsed >= next_tick: - log(f"tick {now():%H:%M} | {label} | running | {last_nonempty_line(log_path)}") + 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(2) - log(log_path.read_text(encoding="utf-8", errors="replace")) - if process.returncode != 0: - raise E2EError(f"{label} failed ({process.returncode}); log: {log_path}") - log(f"tick {now():%H:%M} | {label} | success") - - def _next(self) -> int: - self._sequence += 1 - return self._sequence + 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) # Authoring @@ -250,103 +336,44 @@ def new_project(self, authoring: str) -> Project: """`agentbricks init` a uniquely named project and pin the runtime under test. `init` binds a session store, a memory store and a tracing experiment by default, and CLI - authoring keeps them. Direct authoring replaces the whole manifest in ``write_manifest``. + authoring keeps them, so their stores are created here. Direct authoring replaces the + whole manifest with ``Project.write_manifest``. """ name = f"{project_prefix(self.run_config.run_id)}{uuid.uuid4().hex[:6]}" path = self.run_config.output / "projects" / name path.parent.mkdir(parents=True, exist_ok=True) project = Project(authoring, path, f"{APP_PREFIX}{name}") + workspace = self.target_workspace.client + if workspace is not None: + # Teardown deletes whatever App has this name, so it must not belong to anyone else. + workspace.assert_app_absent(project.app_name) self.projects.append(project) - self._run_long( - f"init-{path.name}", - self.argv("init", "--framework", FRAMEWORKS[0], str(path)), - timeout=600, - ) + self.run(f"init --framework {FRAMEWORKS[0]} {path}", timeout=600) if self.run_config.bridge_sha: self._pin_bridge_sources(path) else: self._pin_project_wheel(path) if authoring == "cli": - self._sync_manifest(project) + self._create_declared_stores(project) return project - def tools_add( - self, project: Project, *args: str, check: bool = True, json_output: bool = False - ) -> subprocess.CompletedProcess[str]: - return self._tools("add", project, args, check, json_output) - - def tools_remove( - self, project: Project, *args: str, check: bool = True, json_output: bool = False - ) -> subprocess.CompletedProcess[str]: - return self._tools("remove", project, args, check, json_output) - - def _tools( - self, - verb: str, - project: Project, - args: Sequence[str], - check: bool, - json_output: bool, - ) -> subprocess.CompletedProcess[str]: - prefix = ("--output", "json") if json_output else () - return self.cli(*prefix, "tools", verb, *args, "--source", str(project.path), check=check) - - def manifest(self, project: Project) -> dict[str, Any]: - return tomli.loads((project.path / "agent.toml").read_text(encoding="utf-8")) - - def write_file(self, project: Project, relative_path: str, content: str) -> None: - """Write a user-owned file in the project, as a developer would.""" - target = project.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_manifest(self, project: Project, tool_specs: Sequence[ToolSpec]) -> None: - """Direct authoring: write agent.toml with exactly these tools plus a tracing binding.""" - sections: list[str] = [_DIRECT_HEADER.format(framework=FRAMEWORKS[0])] - sections.extend(spec.toml for spec in tool_specs if spec.toml) - # Follows default_experiment_name's shape so direct authoring carries the same tracing - # binding `agentbricks init` gives CLI authoring. - slug = re.sub(r"[^a-z0-9-]+", "-", project.path.name.lower()).strip("-") or "agent" - experiment = f"/Shared/agentbricks_traces/{slug}" - sections.append(f'[tracing]\nexperiment_name = "{experiment}"\n') - target = project.path / "agent.toml" - log(f"# write {target}: direct authoring; no agentbricks tools command") - target.write_text("\n".join(sections), encoding="utf-8") - for spec in tool_specs: - for relative_path, content in spec.files.items(): - self.write_file(project, relative_path, content) - self._sync_manifest(project) - - def _sync_manifest(self, project: Project) -> None: - """Read the authored agent.toml and create the stores it declares.""" - manifest = self.manifest(project) + def _create_declared_stores(self, project: Project) -> None: + """Create the stores `init` declared in agent.toml, which deploy expects to exist.""" + manifest = project.manifest() project.memory_store_config = manifest.get("memory_store", {}).get("name") project.session_store_config = manifest.get("session_store", {}).get("name") project.experiment_name = manifest.get("tracing", {}).get("experiment_name") - if project.memory_store_config and not project.memory_store_name: - result = self.cli( - "--output", - "json", - "memory", - "stores", - "create", - "--name", - project.memory_store_config, + if project.memory_store_config: + result = self.run( + f"--output json memory stores create --name {project.memory_store_config}" ) project.memory_resource = json.loads(result.stdout) name = project.memory_resource.get("name") if isinstance(name, str) and name: project.memory_store_name = name - if project.session_store_config and not project.session_store_name: - result = self.cli( - "--output", - "json", - "sessions", - "stores", - "create", - "--name", - project.session_store_config, + if project.session_store_config: + result = self.run( + f"--output json sessions stores create --name {project.session_store_config}" ) project.session_resource = json.loads(result.stdout) project.session_store_name = project.session_store_config @@ -401,122 +428,6 @@ def _pin_bridge_sources(self, project: pathlib.Path) -> None: log(f"# write {pyproject}: pin Agent Bricks and LangChain to {bridge_sha}") pyproject.write_text(tomlkit.dumps(document), encoding="utf-8") - # Running - - @contextlib.contextmanager - def dev(self, project: Project) -> Iterator[Agent]: - """`agentbricks dev` the project locally for the duration of the block.""" - label = f"dev-{project.path.name}-{self._next()}" - log_path = self.logs_dir / f"{label}.log" - port = _free_port() - # Only the first start prepares the venv; run-local refuses to recreate an existing one. - prepare = [] if project.environment_prepared else ["--prepare-environment"] - argv = self.argv("dev", "--source", str(project.path), "--app-port", str(port), *prepare) - 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, - # The template server binds $PORT, so an inherited PORT would override --app-port. - env={**self._env, "PORT": str(port)}, - ) - try: - self._wait_for_local(process, port, label, log_path) - project.environment_prepared = True - yield Agent(f"http://127.0.0.1:{port}", lambda: {}, log_path) - finally: - if process.poll() is None: - terminate_process_group(process) - log(f"{label} | stopped") - - def deploy(self, project: Project, app: str | None = None) -> Agent: - """`agentbricks deploy` the project; deploying again updates the same App.""" - workspace = self.target_workspace.client - if workspace is None: - raise E2EError( - f"deploy needs a live workspace; got --workspace {self.target_workspace.name}" - ) - workspace.check_app_auth() - name = app or project.app_name - creating = not project.app_registered - if creating: - workspace.assert_app_absent(name) - # Deploy can create the App and then fail while waiting for it, so mark it first. - project.app_registered = True - label = f"deploy-{name}-{self._next()}" - try: - self._run_long( - label, - self.argv("deploy", name, "--source", str(project.path)), - timeout=2400, - ) - deployed = workspace.wait_for_app(name) - finally: - # Also after a failed deploy, which can leave an App, store and role behind. - self._record_identity(project, name, workspace) - url = str(deployed.get("url") or "").rstrip("/") - if not url: - raise E2EError(f"App {name} has no URL: {deployed}") - agent = Agent( - url, - lambda: workspace.app_headers, - self.logs_dir / f"{label}.log", - app_name=name, - app=deployed, - ) - if creating: - agent.warm_up() - return agent - - def _record_identity(self, project: Project, app: str, workspace: Workspace) -> None: - """Keep the App's principal and Lakebase branch for cleanup, including leftover cleanup after the run.""" - identity = project.identity.merged(workspace.app_identity(app)) - if identity != project.identity: - project.identity = identity - record_app_identity(self.run_config.output, app, identity) - - def _wait_for_local( - self, process: subprocess.Popen[str], port: int, label: str, log_path: pathlib.Path - ) -> None: - started = time.monotonic() - next_tick = 60.0 - while True: - if process.poll() is not None: - raise E2EError(f"{label} exited {process.returncode}: {last_lines(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 > 1200: - raise E2EError(f"{label} did not become reachable: {last_lines(log_path, 30)}") - if elapsed >= next_tick: - log(f"tick {now():%H:%M} | {label} | starting | {last_nonempty_line(log_path)}") - next_tick += 60 - time.sleep(5) - - # Stores - - def delete_store(self, kind: str, name: str) -> subprocess.CompletedProcess[str]: - """`agentbricks stores delete`.""" - argv = [*self.argv(*_STORE_KINDS[kind]), "delete", name, "--yes"] - log_command(argv) - result = subprocess.run( - argv, env=self._env, text=True, capture_output=True, timeout=600, check=False - ) - for text in (result.stdout, result.stderr): - if text.strip(): - log(text) - return result - def _agentbricks_bin() -> str: """The CLI of the interpreter running the tests, i.e. the source checkout's own install.""" @@ -543,7 +454,7 @@ def terminate_process_group(process: subprocess.Popen[str], *, grace: float = 20 pass -def _free_port() -> int: +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/suites/harness/workspace.py b/integrations/agentbricks/tests/suites/harness/workspace.py index 6a846127e..bc72cf1b1 100644 --- a/integrations/agentbricks/tests/suites/harness/workspace.py +++ b/integrations/agentbricks/tests/suites/harness/workspace.py @@ -378,7 +378,11 @@ def _principal(app: dict[str, Any]) -> str: # Cleanup def delete_app(self, name: str) -> None: - self.client.apps.delete(name) + """Delete the App; a missing App counts as deleted.""" + try: + self.client.apps.delete(name) + except NotFound: + return self._wait_for_app_deleted(name) def _wait_for_app_deleted(self, name: str, timeout: float = 1200) -> None: diff --git a/integrations/agentbricks/tests/suites/memory_store/test_memory_store.py b/integrations/agentbricks/tests/suites/memory_store/test_memory_store.py index f93e5cdb0..682be0d80 100644 --- a/integrations/agentbricks/tests/suites/memory_store/test_memory_store.py +++ b/integrations/agentbricks/tests/suites/memory_store/test_memory_store.py @@ -6,13 +6,14 @@ from __future__ import annotations import pytest -from harness.cli import AgentbricksCli +from harness.cli import AgentbricksCli, AgentEndpoint +from harness.workspace import Workspace -@pytest.mark.usefixtures("workspace_client") -def test_memory_store_created(agentbricks_cli: AgentbricksCli) -> None: +def test_memory_store_created(agentbricks_cli: AgentbricksCli, workspace_client: Workspace) -> None: project = agentbricks_cli.new_project("cli") - agentbricks_cli.deploy(project) + agentbricks_cli.run(f"deploy {project.app_name} --source {project.path}", timeout=2400) + AgentEndpoint.deployed(workspace_client, project.app_name) resource = project.memory_resource or {} assert resource.get("display_name") == project.memory_store_config, resource diff --git a/integrations/agentbricks/tests/suites/session_store/test_session_store.py b/integrations/agentbricks/tests/suites/session_store/test_session_store.py index a15f6c713..56a7d9709 100644 --- a/integrations/agentbricks/tests/suites/session_store/test_session_store.py +++ b/integrations/agentbricks/tests/suites/session_store/test_session_store.py @@ -7,13 +7,16 @@ from __future__ import annotations import pytest -from harness.cli import AgentbricksCli +from harness.cli import AgentbricksCli, AgentEndpoint +from harness.workspace import Workspace -@pytest.mark.usefixtures("workspace_client") -def test_session_store_created(agentbricks_cli: AgentbricksCli) -> None: +def test_session_store_created( + agentbricks_cli: AgentbricksCli, workspace_client: Workspace +) -> None: project = agentbricks_cli.new_project("cli") - agentbricks_cli.deploy(project) + agentbricks_cli.run(f"deploy {project.app_name} --source {project.path}", timeout=2400) + AgentEndpoint.deployed(workspace_client, project.app_name) resource = project.session_resource or {} assert resource.get("session_store_name") == project.session_store_config, resource diff --git a/integrations/agentbricks/tests/suites/tools/test_deploy.py b/integrations/agentbricks/tests/suites/tools/test_deploy.py index 9cd4b0b6c..0f268fe02 100644 --- a/integrations/agentbricks/tests/suites/tools/test_deploy.py +++ b/integrations/agentbricks/tests/suites/tools/test_deploy.py @@ -8,7 +8,7 @@ from __future__ import annotations -from harness.cli import AgentbricksCli +from harness.cli import AgentbricksCli, AgentEndpoint from harness.logs import log from harness.workspace import Workspace from tool_definitions import Tool, assert_endpoint_invoke, bind, invoke_and_check @@ -46,23 +46,25 @@ def grant_transitive_if_needed(tool: Tool) -> None: for index, tool in enumerate(tools): with subtests.test(tool=tool.name): bind(agentbricks_cli, project, tool) - agent = agentbricks_cli.deploy(project, app) + agentbricks_cli.run(f"deploy {app} --source {project.path}", timeout=2400) + endpoint = AgentEndpoint.deployed(workspace_client, app) wait_for_grant(tool) grant_transitive_if_needed(tool) - invoke_and_check(agent, tool) + invoke_and_check(endpoint, tool) for added in tools[: index + 1]: assert_granted(added) with subtests.test(step="endpoint-invoke"): - assert_endpoint_invoke(agentbricks_cli, agent) + assert_endpoint_invoke(agentbricks_cli, endpoint) return - agentbricks_cli.write_manifest(project, tools) - agent = agentbricks_cli.deploy(project, app) + project.write_manifest(tools) + agentbricks_cli.run(f"deploy {app} --source {project.path}", timeout=2400) + endpoint = AgentEndpoint.deployed(workspace_client, app) for tool in tools: with subtests.test(tool=tool.name): wait_for_grant(tool) grant_transitive_if_needed(tool) - invoke_and_check(agent, tool) + invoke_and_check(endpoint, tool) assert_granted(tool) with subtests.test(step="endpoint-invoke"): - assert_endpoint_invoke(agentbricks_cli, agent) + assert_endpoint_invoke(agentbricks_cli, endpoint) diff --git a/integrations/agentbricks/tests/suites/tools/test_dev.py b/integrations/agentbricks/tests/suites/tools/test_dev.py index 078105934..6f99b7b9f 100644 --- a/integrations/agentbricks/tests/suites/tools/test_dev.py +++ b/integrations/agentbricks/tests/suites/tools/test_dev.py @@ -7,8 +7,11 @@ from __future__ import annotations +import contextlib +from collections.abc import Iterator + import pytest -from harness.cli import AgentbricksCli +from harness.cli import AgentbricksCli, AgentEndpoint, Project, free_port from tool_definitions import ( Tool, assert_endpoint_invoke, @@ -21,6 +24,17 @@ pytestmark = pytest.mark.xdist_group("dev") +@contextlib.contextmanager +def dev(cli: AgentbricksCli, project: Project) -> Iterator[AgentEndpoint]: + """`agentbricks dev` the project for the duration of the block.""" + port = free_port() + # The template server binds $PORT, so an inherited PORT would override --app-port. + with cli.start( + f"dev --source {project.path} --app-port {port}", env={"PORT": str(port)} + ) as server: + yield AgentEndpoint.local(server, port) + + def test_dev( agentbricks_cli: AgentbricksCli, tools: tuple[Tool, ...], authoring: str, subtests ) -> None: @@ -32,16 +46,16 @@ def test_dev( for tool in tools: with subtests.test(tool=tool.name): bind(agentbricks_cli, project, tool) - with agentbricks_cli.dev(project) as agent: - invoke_and_check(agent, tool) - with subtests.test(step="endpoint-invoke"), agentbricks_cli.dev(project) as agent: - assert_endpoint_invoke(agentbricks_cli, agent) + with dev(agentbricks_cli, project) as endpoint: + invoke_and_check(endpoint, tool) + with subtests.test(step="endpoint-invoke"), dev(agentbricks_cli, project) as endpoint: + assert_endpoint_invoke(agentbricks_cli, endpoint) return - agentbricks_cli.write_manifest(project, tools) - with agentbricks_cli.dev(project) as agent: + project.write_manifest(tools) + with dev(agentbricks_cli, project) as endpoint: for tool in tools: with subtests.test(tool=tool.name): - invoke_and_check(agent, tool) + invoke_and_check(endpoint, tool) with subtests.test(step="endpoint-invoke"): - assert_endpoint_invoke(agentbricks_cli, agent) + assert_endpoint_invoke(agentbricks_cli, endpoint) diff --git a/integrations/agentbricks/tests/suites/tools/test_grants.py b/integrations/agentbricks/tests/suites/tools/test_grants.py index 34b84166c..92df4cc70 100644 --- a/integrations/agentbricks/tests/suites/tools/test_grants.py +++ b/integrations/agentbricks/tests/suites/tools/test_grants.py @@ -2,7 +2,7 @@ from __future__ import annotations -from harness.cli import AgentbricksCli +from harness.cli import AgentbricksCli, AgentEndpoint from harness.run import TOOL_RESOURCE_PREFIX from harness.workspace import Workspace from tool_definitions import bind, uc_function_tool @@ -23,8 +23,9 @@ def test_grants( if authoring == "cli": bind(agentbricks_cli, project, tool) else: - agentbricks_cli.write_manifest(project, [tool]) - agentbricks_cli.deploy(project, app) + project.write_manifest([tool]) + agentbricks_cli.run(f"deploy {app} --source {project.path}", timeout=2400) + AgentEndpoint.deployed(workspace_client, app) with subtests.test(check="transitive-function-not-auto-granted"): # Raises if the hidden nested function is already granted or listed as an App resource. diff --git a/integrations/agentbricks/tests/suites/tools/test_selection.py b/integrations/agentbricks/tests/suites/tools/test_selection.py index bb600d7ca..24b9fde80 100644 --- a/integrations/agentbricks/tests/suites/tools/test_selection.py +++ b/integrations/agentbricks/tests/suites/tools/test_selection.py @@ -6,7 +6,7 @@ from __future__ import annotations -from harness.cli import AgentbricksCli +from harness.cli import AgentbricksCli, AgentEndpoint from harness.workspace import Workspace from tool_definitions import PYTHON_MARKER, UC_MARKER, Tool, bind, excerpt from uc_objects import UcFunction @@ -31,11 +31,12 @@ def test_selects_uc_function_not_python_marker( for tool in tools: bind(agentbricks_cli, project, tool) else: - agentbricks_cli.write_manifest(project, tools) - agent = agentbricks_cli.deploy(project, app) + project.write_manifest(tools) + agentbricks_cli.run(f"deploy {app} --source {project.path}", timeout=2400) + endpoint = AgentEndpoint.deployed(workspace_client, app) # The nested function is deliberately outside Agent Bricks' grants. workspace_client.grant_transitive(app, uc_function.nested_function) - serialized = agent.invoke(PROMPT, "deploy-selection") + serialized = endpoint.invoke(PROMPT, "deploy-selection") assert UC_MARKER in serialized, f"The UC function tool was not called: {excerpt(serialized)}" assert PYTHON_MARKER not in serialized, ( diff --git a/integrations/agentbricks/tests/suites/tools/tool_definitions.py b/integrations/agentbricks/tests/suites/tools/tool_definitions.py index 96a2857c2..255d9aabd 100644 --- a/integrations/agentbricks/tests/suites/tools/tool_definitions.py +++ b/integrations/agentbricks/tests/suites/tools/tool_definitions.py @@ -4,9 +4,11 @@ import dataclasses import json +import shlex +import uuid from collections.abc import Callable, Mapping -from harness.cli import Agent, AgentbricksCli, Project +from harness.cli import AgentbricksCli, AgentEndpoint, Project from harness.workspace import GrantTuple from uc_objects import UcFunction @@ -31,8 +33,8 @@ class Tool: prompt: str # Asserts the serialized response proves the tool ran. check: Check - # `agentbricks tools add` arguments; None when the tool is a user-owned file instead. - add_args: tuple[str, ...] | None + # What follows `agentbricks tools add`; None when the tool is a user-owned file instead. + add_command: str | None # The [[tools]] entry for direct authoring, and any user-owned files the tool needs. toml: str files: Mapping[str, str] @@ -97,7 +99,7 @@ def sandbox_tool(uc_function: UcFunction) -> Tool: check=contains(uc_function.volume_marker), # The CLI's default auth: the sandbox reads the volume as the user invoking the agent, so # deploy adds no App grant. A manifest entry without `auth` would run as the App instead. - add_args=("sandbox", "--scope", f"volume:{uc_function.volume}"), + add_command=f"sandbox --scope volume:{uc_function.volume}", toml=( '[[tools]]\nid = "sandbox"\nauth = "user"\n' 'source = { kind = "sandbox", service = "system.ai.sandbox" }\n' @@ -121,7 +123,7 @@ def web_search_tool(**_resources: object) -> Tool: "return the title and https URL of one result. Do not answer from memory." ), check=check_web_search, - add_args=("mcp", "system.ai.web_search", "--auth", "app"), + add_command="mcp system.ai.web_search --auth app", toml=( '[[tools]]\nid = "web_search"\nauth = "app"\n' 'source = { kind = "mcp", service = "system.ai.web_search" }\n' @@ -140,7 +142,7 @@ def python_tool() -> Tool: "Return its exact result." ), check=contains(PYTHON_MARKER), - add_args=None, + add_command=None, toml="", files={PYTHON_MARKER_FILE: PYTHON_MARKER_SOURCE}, tool_id=None, @@ -156,7 +158,7 @@ def uc_function_tool(uc_function: UcFunction) -> Tool: "'matrix'. Return the called tool's exact result." ), check=contains(UC_MARKER), - add_args=("uc-function", uc_function.function, "--name", "agentbricks_uc_marker"), + add_command=f"uc-function {uc_function.function} --name agentbricks_uc_marker", toml=( '[[tools]]\nid = "agentbricks_uc_marker"\n' f'source = {{ kind = "uc_function", function = "{uc_function.function}" }}\n' @@ -178,7 +180,7 @@ def genie_tool(genie_space: str) -> Tool: "response." ), check=check_genie, - add_args=("genie-agent", genie_space, "--name", "genie", "--auth", "app"), + add_command=f"genie-agent {genie_space} --name genie --auth app", toml=( '[[tools]]\nid = "genie"\nauth = "app"\n' f'source = {{ kind = "genie_agent", space_id = "{genie_space}" }}\n' @@ -202,13 +204,13 @@ def genie_tool(genie_space: str) -> Tool: def bind(cli: AgentbricksCli, project: Project, tool: Tool) -> None: """CLI authoring: add one tool and confirm agent.toml declares it as intended.""" - if tool.add_args is not None: - cli.tools_add(project, *tool.add_args) + if tool.add_command is not None: + cli.run(f"tools add {tool.add_command} --source {project.path}") for relative_path, content in tool.files.items(): - cli.write_file(project, relative_path, content) + project.write_file(relative_path, content) if tool.tool_id is None: return - declared = {t["id"]: t for t in cli.manifest(project).get("tools", [])} + declared = {t["id"]: t for t in project.manifest().get("tools", [])} assert tool.tool_id in declared, ( f"tools add did not declare {tool.tool_id!r}: {sorted(declared)}" ) @@ -222,14 +224,10 @@ def reject_unavailable_mcp(cli: AgentbricksCli, project: Project) -> None: """`tools add` must refuse an MCP service that does not exist and leave agent.toml untouched.""" manifest = project.path / "agent.toml" before = manifest.read_bytes() - rejected = cli.tools_add( - project, - "mcp", - "system.ai.missing_service", - "--name", - "broken_mcp", + rejected = cli.run( + "--output json tools add mcp system.ai.missing_service --name broken_mcp " + f"--source {project.path}", check=False, - json_output=True, ) assert rejected.returncode != 0 and manifest.read_bytes() == before, ( "agentbricks tools add accepted an unavailable MCP service or changed agent.toml" @@ -238,17 +236,29 @@ def reject_unavailable_mcp(cli: AgentbricksCli, project: Project) -> None: assert code in {"NOT_FOUND", "RESOURCE_DOES_NOT_EXIST"}, ( f"Unexpected MCP validation error: {rejected.stderr}" ) - cli.tools_remove(project, "mcp", "system.ai.missing_service") - assert not any(tool["id"] == "broken_mcp" for tool in cli.manifest(project).get("tools", [])), ( + cli.run(f"tools remove mcp system.ai.missing_service --source {project.path}") + assert not any(tool["id"] == "broken_mcp" for tool in project.manifest().get("tools", [])), ( "agentbricks tools remove left the broken MCP binding in agent.toml" ) -def invoke_and_check(agent: Agent, tool: Tool) -> None: - tool.check(agent.invoke(tool.prompt, f"{agent.runtime}-{tool.name}")) +def invoke_and_check(endpoint: AgentEndpoint, tool: Tool) -> None: + tool.check(endpoint.invoke(tool.prompt, f"{endpoint.runtime}-{tool.name}")) -def assert_endpoint_invoke(agentbricks_cli: AgentbricksCli, agent: Agent) -> None: - """The CLI's own invoke command reaches the running agent and gets a reply.""" - output = agentbricks_cli.endpoint_invoke(agent, "Reply with the single word OK.") +def assert_endpoint_invoke(cli: AgentbricksCli, endpoint: AgentEndpoint) -> None: + """`agentbricks endpoint invoke`, the call the CLI tells users to run after dev/deploy, gets a reply.""" + invocation_id = str(uuid.uuid4()) + body = json.dumps( + { + "id": invocation_id, + "session_id": invocation_id, + "input": [{"role": "user", "content": "Reply with the single word OK."}], + } + ) + target = endpoint.app_name or f"--url {endpoint.base_url} --no-auth" + output = cli.run( + f"endpoint invoke {target} --path /api/invocations --json {shlex.quote(body)}", + retry_unavailable=True, + ).stdout assert "ok" in output.lower(), f"endpoint invoke returned no reply: {output[:500]}" diff --git a/integrations/agentbricks/tests/suites/tracing/test_tracing.py b/integrations/agentbricks/tests/suites/tracing/test_tracing.py index ac1859a8f..85d807af0 100644 --- a/integrations/agentbricks/tests/suites/tracing/test_tracing.py +++ b/integrations/agentbricks/tests/suites/tracing/test_tracing.py @@ -6,7 +6,7 @@ from __future__ import annotations import pytest -from harness.cli import AgentbricksCli +from harness.cli import AgentbricksCli, AgentEndpoint from harness.run import TOOL_RESOURCE_PREFIX from harness.workspace import Workspace @@ -15,7 +15,8 @@ def test_tracing_experiment_bound( agentbricks_cli: AgentbricksCli, workspace_client: Workspace ) -> None: project = agentbricks_cli.new_project("cli") - agentbricks_cli.deploy(project) + agentbricks_cli.run(f"deploy {project.app_name} --source {project.path}", timeout=2400) + AgentEndpoint.deployed(workspace_client, project.app_name) assert project.experiment_name, "agent.toml binds no tracing experiment" # Deploy attaches the experiment as an App resource that Agent Bricks does not own. From c3274703785bf644d7497b86f9281f55d684334d Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Wed, 14 Oct 2026 18:45:52 +0000 Subject: [PATCH 17/45] [agentbricks] E2E suites: tests run init themselves; project fixture owns cleanup Co-authored-by: Isaac --- .../agentbricks/tests/suites/README.md | 11 +- .../agentbricks/tests/suites/conftest.py | 102 ++++++----- .../tests/suites/harness/cleanup.py | 10 +- .../agentbricks/tests/suites/harness/cli.py | 163 +----------------- .../tests/suites/harness/project.py | 143 +++++++++++++++ .../suites/memory_store/test_memory_store.py | 23 ++- .../session_store/test_session_store.py | 16 +- .../tests/suites/tools/test_deploy.py | 5 +- .../tests/suites/tools/test_dev.py | 12 +- .../tests/suites/tools/test_grants.py | 5 +- .../tests/suites/tools/test_selection.py | 5 +- .../tests/suites/tools/tool_definitions.py | 3 +- .../tests/suites/tracing/test_tracing.py | 13 +- 13 files changed, 284 insertions(+), 227 deletions(-) create mode 100644 integrations/agentbricks/tests/suites/harness/project.py diff --git a/integrations/agentbricks/tests/suites/README.md b/integrations/agentbricks/tests/suites/README.md index b0cfc037d..8119871da 100644 --- a/integrations/agentbricks/tests/suites/README.md +++ b/integrations/agentbricks/tests/suites/README.md @@ -29,9 +29,9 @@ The run id and output directory (logs, generated projects, `apps.jsonl`) are pri Every resource a run creates carries its run id, and is deleted at three points: -1. **When each test ends**, by the fixture that made it: `agentbricks_cli` deletes the test's Apps and - everything their deploys created (`harness.cleanup.delete_app_resources`); `scratch_schema` deletes - the test's schema. +1. **When each test ends**, by the fixture that made it: `project` deletes the test's App, its stores, + and everything its deploys created (`harness.cleanup.delete_app_resources`); `scratch_schema` + deletes the test's schema. 2. **When the run ends**, `pytest_sessionfinish` deletes whatever a crashed test or a failed teardown left (`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: @@ -50,6 +50,7 @@ Every resource a run creates carries its run id, and is deleted at three points: | `harness/run.py` | `RunConfig`, the run id, and the name prefixes every resource carries | | `harness/target.py` | The `--workspace live\|fake` seam the CLI and fixtures talk to | | `harness/cli.py` | Runs `agentbricks` commands (`cli.run("tools add ...")`), and the endpoints of the agents they start | +| `harness/project.py` | A test's project: run-scoped path and App name, its files, and `pin_runtime` | | `harness/workspace.py` | Databricks SDK helpers for what is not under test: SQL, UC, Apps, grants | | `harness/cleanup.py` | Deleting what the suites create, per App and per run | | `harness/logs.py` | Test-log output; every command is echoed so a run's log replays it | @@ -62,5 +63,7 @@ that you can follow through its imports. - **A tool:** add one `ToolDefinition(name, requires, build)` in `tools/tool_definitions.py`. `requires` names the fixtures it needs; tools whose fixtures skip are left out and logged. -- **A suite:** add a folder with `test_*.py` files. Add a `conftest.py` only for resources the CLI +- **A suite:** add a folder with `test_*.py` files. A test takes the `project` fixture, runs + `agentbricks_cli.run(f"init ... {project.path}")`, then `pin_runtime(project, ...)` so the project + runs the checkout's runtime; teardown fails a test that skips the pin. Add a `conftest.py` only for resources the CLI under test does not create itself; create them per test and delete them in the fixture's teardown. diff --git a/integrations/agentbricks/tests/suites/conftest.py b/integrations/agentbricks/tests/suites/conftest.py index 7504bda90..8cba524c7 100644 --- a/integrations/agentbricks/tests/suites/conftest.py +++ b/integrations/agentbricks/tests/suites/conftest.py @@ -29,6 +29,7 @@ from harness.cleanup import delete_app_resources, delete_run_leftovers, record_app_identity from harness.cli import AgentbricksCli from harness.logs import log, log_command +from harness.project import Project, is_pinned from harness.run import RUN_ID_PATTERN, E2EError, RunConfig, run_output_dir from harness.target import ( WORKSPACE_KINDS, @@ -230,45 +231,68 @@ def pytest_runtest_logreport(report: pytest.TestReport) -> None: @pytest.fixture def agentbricks_cli( request: pytest.FixtureRequest, target_workspace: TargetWorkspace, run_config: RunConfig -) -> Iterator[AgentbricksCli]: - """The CLI wrapper for one test; deletes every App and store the test created when it ends.""" - cli = AgentbricksCli(target_workspace, run_config, request.node.name) - yield cli - if cli.projects and target_workspace.client is not None: - _delete_projects(cli, target_workspace.client, failed=request.node.nodeid in _failed_tests) - - -def _delete_projects(cli: AgentbricksCli, workspace: Workspace, *, failed: bool) -> None: - for project in cli.projects: - name = project.app_name - # Whatever deploy got as far as creating; read before anything is deleted, because the - # principal and branch come from the App and its Runtime Store. - identity = workspace.app_identity(name) - deployed = identity != AppIdentity() - if deployed: - # Leftover cleanup retries the Lakebase deletes from this record if they fail here. - record_app_identity(cli.run_config.output, name, identity) - if failed and identity.service_principal: - # The App is deleted next, so a failing test must capture its logs while it exists. - workspace.app_logs(name, cli.logs_dir / f"deploy-runtime-{name}.log") - try: - failures = delete_app_resources( - workspace, - name, - delete_store=lambda kind, store: cli.run( - f"{kind} stores delete {store} --yes", timeout=600, check=False - ), - memory_store=project.memory_store_name, - session_store=project.session_store_name, - has_app=deployed, - # Every managed deploy creates a Runtime Store, with or without memory/session stores. - runtime_store=deployed, - identity=identity, - ) - except Exception as exc: - failures = [str(exc)] - for failure in failures: - log(f"cleanup warning | {name} | {failure}") +) -> AgentbricksCli: + return AgentbricksCli(target_workspace, run_config, request.node.name) + + +@pytest.fixture +def project( + request: pytest.FixtureRequest, + agentbricks_cli: AgentbricksCli, + target_workspace: TargetWorkspace, + run_config: RunConfig, +) -> Iterator[Project]: + """A run-scoped path and App name for the test to `agentbricks init`; deleted when it ends.""" + project = Project.allocate(run_config) + workspace = target_workspace.client + if workspace is not None: + # Teardown deletes whatever App has this name, so it must not belong to anyone else. + workspace.assert_app_absent(project.app_name) + yield project + failed = request.node.nodeid in _failed_tests + if workspace is not None: + _delete_project(agentbricks_cli, workspace, project, 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 is_pinned(project): + raise E2EError( + f"{project.path} was never pinned to the runtime under test; call pin_runtime after init." + ) + + +def _delete_project( + cli: AgentbricksCli, workspace: Workspace, project: Project, *, failed: bool +) -> None: + name = project.app_name + # Whatever deploy got as far as creating; read before anything is deleted, because the + # principal and branch come from the App and its Runtime Store. + identity = workspace.app_identity(name) + deployed = identity != AppIdentity() + if deployed: + # Leftover cleanup retries the Lakebase deletes from this record if they fail here. + record_app_identity(cli.run_config.output, name, identity) + if failed and identity.service_principal: + # The App is deleted next, so a failing test must capture its logs while it exists. + workspace.app_logs(name, cli.logs_dir / f"deploy-runtime-{name}.log") + # Deploy creates the stores agent.toml declares; memory stores are addressed by resource name. + memory_store = project.declared_store("memory") + try: + failures = delete_app_resources( + workspace, + name, + delete_store=lambda kind, store: cli.run( + f"{kind} stores delete {store} --yes", timeout=600, check=False + ), + memory_store=f"memory-stores/{memory_store}" if memory_store else None, + session_store=project.declared_store("session"), + has_app=deployed, + # Every managed deploy creates a Runtime Store, with or without memory/session stores. + runtime_store=deployed, + identity=identity, + ) + except Exception as exc: + failures = [str(exc)] + for failure in failures: + log(f"cleanup warning | {name} | {failure}") # Deleting what the whole run left behind diff --git a/integrations/agentbricks/tests/suites/harness/cleanup.py b/integrations/agentbricks/tests/suites/harness/cleanup.py index a27affb7e..76b000bb7 100644 --- a/integrations/agentbricks/tests/suites/harness/cleanup.py +++ b/integrations/agentbricks/tests/suites/harness/cleanup.py @@ -16,6 +16,7 @@ import dataclasses import json import pathlib +import re import subprocess from collections.abc import Callable, Collection from typing import Any @@ -25,6 +26,9 @@ from harness.workspace import AppIdentity, Workspace _APPS_RECORD = "apps.jsonl" +_STORE_MISSING = re.compile( + r"NOT_FOUND|RESOURCE_DOES_NOT_EXIST|does not exist|not found", re.IGNORECASE +) # Identities recorded by per-test cleanup @@ -130,8 +134,10 @@ def attempt(label: str, action: Callable[..., object], *args: Any) -> bool: def delete_cli_store(kind: str, store: str) -> None: assert delete_store is not None result = delete_store(kind, store) - if result.returncode != 0: - raise E2EError((result.stderr or result.stdout).strip()) + output = (result.stderr or result.stdout).strip() + # A store the test declared but never deployed was never created. + if result.returncode != 0 and not _STORE_MISSING.search(output): + raise E2EError(output) # Read before any delete: the principal and branch come from the App and its Runtime Store. known = (identity or AppIdentity()).merged( diff --git a/integrations/agentbricks/tests/suites/harness/cli.py b/integrations/agentbricks/tests/suites/harness/cli.py index 8cf314b76..9a2f6b028 100644 --- a/integrations/agentbricks/tests/suites/harness/cli.py +++ b/integrations/agentbricks/tests/suites/harness/cli.py @@ -3,8 +3,7 @@ Every CLI 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. The only non-subprocess part is ``AgentEndpoint.invoke``, which POSTs to a running agent. Nothing in this -module knows about any particular suite. It records the projects it creates in ``projects`` but -deletes nothing: the ``agentbricks_cli`` fixture does. +module knows about any particular suite or deletes anything: the ``project`` fixture does. """ from __future__ import annotations @@ -26,71 +25,18 @@ import urllib.error import urllib.request import uuid -from collections.abc import Callable, Iterator, Mapping, Sequence -from typing import Any, Protocol - -import tomli -import tomlkit +from collections.abc import Callable, Iterator, Mapping +from typing import Any from harness.logs import last_lines, last_nonempty_line, log, log_command -from harness.run import APP_PREFIX, E2E_MODEL, FRAMEWORKS, E2EError, RunConfig, now, project_prefix +from harness.run import E2E_MODEL, E2EError, RunConfig, now from harness.target import TargetWorkspace from harness.workspace import Workspace -_DIRECT_HEADER = ( - 'schema_version = 1\n\n[agent]\nframework = "{framework}"\nserver = "agentbricks"\n' -) # The Apps gateway's answers while an App is not up, e.g. restarting after a redeploy. _GATEWAY_UNAVAILABLE = re.compile(r"HTTP 50[234]\b") -class ToolSpec(Protocol): - """What ``write_manifest`` needs from a tool: its [[tools]] TOML and any user-owned files.""" - - toml: str - files: Mapping[str, str] - - -@dataclasses.dataclass -class Project: - authoring: str - path: pathlib.Path - app_name: str - memory_store_config: str | None = None - session_store_config: str | None = None - memory_resource: dict[str, Any] | None = None - session_resource: dict[str, Any] | None = None - memory_store_name: str | None = None - session_store_name: str | None = None - experiment_name: str | None = None - - def manifest(self) -> dict[str, Any]: - return tomli.loads((self.path / "agent.toml").read_text(encoding="utf-8")) - - 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_manifest(self, tool_specs: Sequence[ToolSpec]) -> None: - """Direct authoring: write agent.toml with exactly these tools plus a tracing binding.""" - sections: list[str] = [_DIRECT_HEADER.format(framework=FRAMEWORKS[0])] - sections.extend(spec.toml for spec in tool_specs if spec.toml) - # Follows default_experiment_name's shape so direct authoring carries the same tracing - # binding `agentbricks init` gives CLI authoring. - slug = re.sub(r"[^a-z0-9-]+", "-", self.path.name.lower()).strip("-") or "agent" - self.experiment_name = f"/Shared/agentbricks_traces/{slug}" - sections.append(f'[tracing]\nexperiment_name = "{self.experiment_name}"\n') - target = self.path / "agent.toml" - log(f"# write {target}: direct authoring; no agentbricks tools command") - target.write_text("\n".join(sections), encoding="utf-8") - for spec in tool_specs: - for relative_path, content in spec.files.items(): - self.write_file(relative_path, content) - - @dataclasses.dataclass class BackgroundCommand: """A CLI command left running by ``AgentbricksCli.start``, e.g. ``agentbricks dev``.""" @@ -206,7 +152,7 @@ def warm_up(self, deadline_seconds: float = 180.0) -> None: class AgentbricksCli: - """Runs ``agentbricks`` commands against a target workspace and authors projects with them.""" + """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 @@ -215,7 +161,6 @@ def __init__(self, target_workspace: TargetWorkspace, run_config: RunConfig, tes self.logs_dir = ( run_config.output / "logs" / re.sub(r"[^A-Za-z0-9_.-]+", "-", test_name).strip("-") ) - self.projects: list[Project] = [] # 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. self._env = {**os.environ, **target_workspace.env, "AGENTBRICKS_DISABLE_AGENT_HINT": "1"} @@ -330,104 +275,6 @@ def _run_once( log(text) return subprocess.CompletedProcess(argv, process.returncode, stdout, stderr) - # Authoring - - def new_project(self, authoring: str) -> Project: - """`agentbricks init` a uniquely named project and pin the runtime under test. - - `init` binds a session store, a memory store and a tracing experiment by default, and CLI - authoring keeps them, so their stores are created here. Direct authoring replaces the - whole manifest with ``Project.write_manifest``. - """ - name = f"{project_prefix(self.run_config.run_id)}{uuid.uuid4().hex[:6]}" - path = self.run_config.output / "projects" / name - path.parent.mkdir(parents=True, exist_ok=True) - project = Project(authoring, path, f"{APP_PREFIX}{name}") - workspace = self.target_workspace.client - if workspace is not None: - # Teardown deletes whatever App has this name, so it must not belong to anyone else. - workspace.assert_app_absent(project.app_name) - self.projects.append(project) - self.run(f"init --framework {FRAMEWORKS[0]} {path}", timeout=600) - if self.run_config.bridge_sha: - self._pin_bridge_sources(path) - else: - self._pin_project_wheel(path) - if authoring == "cli": - self._create_declared_stores(project) - return project - - def _create_declared_stores(self, project: Project) -> None: - """Create the stores `init` declared in agent.toml, which deploy expects to exist.""" - manifest = project.manifest() - project.memory_store_config = manifest.get("memory_store", {}).get("name") - project.session_store_config = manifest.get("session_store", {}).get("name") - project.experiment_name = manifest.get("tracing", {}).get("experiment_name") - if project.memory_store_config: - result = self.run( - f"--output json memory stores create --name {project.memory_store_config}" - ) - project.memory_resource = json.loads(result.stdout) - name = project.memory_resource.get("name") - if isinstance(name, str) and name: - project.memory_store_name = name - if project.session_store_config: - result = self.run( - f"--output json sessions stores create --name {project.session_store_config}" - ) - project.session_resource = json.loads(result.stdout) - project.session_store_name = project.session_store_config - - def _pin_project_wheel(self, project: pathlib.Path) -> None: - wheel = self.run_config.wheel - assert wheel is not 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") - - def _pin_bridge_sources(self, project: pathlib.Path) -> None: - bridge_sha = self.run_config.bridge_sha - pyproject = project / "pyproject.toml" - document = tomlkit.parse(pyproject.read_text(encoding="utf-8")) - dependencies = document["project"]["dependencies"] - if not any( - str(dependency).startswith("databricks-langchain") for dependency in dependencies - ): - dependencies.append("databricks-langchain>=0.17.0") - if "tool" not in document: - document["tool"] = tomlkit.table() - if "uv" not in document["tool"]: - document["tool"]["uv"] = tomlkit.table() - if "sources" not in document["tool"]["uv"]: - document["tool"]["uv"]["sources"] = tomlkit.table() - for package, subdirectory in ( - ("databricks-agentbricks", "integrations/agentbricks"), - ("databricks-langchain", "integrations/langchain"), - ): - source = tomlkit.inline_table() - source.update( - { - "git": "https://github.com/databricks/databricks-ai-bridge.git", - "rev": bridge_sha, - "subdirectory": subdirectory, - } - ) - document["tool"]["uv"]["sources"][package] = source - log(f"# write {pyproject}: pin Agent Bricks and LangChain to {bridge_sha}") - pyproject.write_text(tomlkit.dumps(document), encoding="utf-8") - def _agentbricks_bin() -> str: """The CLI of the interpreter running the tests, i.e. the source checkout's own install.""" diff --git a/integrations/agentbricks/tests/suites/harness/project.py b/integrations/agentbricks/tests/suites/harness/project.py new file mode 100644 index 000000000..06bd4bfe9 --- /dev/null +++ b/integrations/agentbricks/tests/suites/harness/project.py @@ -0,0 +1,143 @@ +"""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 re +import shutil +import uuid +from collections.abc import Mapping, Sequence +from typing import Any, Protocol + +import tomli +import tomlkit + +from harness.logs import log +from harness.run import APP_PREFIX, FRAMEWORKS, RunConfig, project_prefix + +_DIRECT_HEADER = ( + 'schema_version = 1\n\n[agent]\nframework = "{framework}"\nserver = "agentbricks"\n' +) + + +class ToolSpec(Protocol): + """What ``write_manifest`` needs from a tool: its [[tools]] TOML and any user-owned files.""" + + toml: str + files: Mapping[str, str] + + +@dataclasses.dataclass(frozen=True) +class Project: + path: pathlib.Path + app_name: str + + @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) + return cls(path, f"{APP_PREFIX}{name}") + + def manifest(self) -> dict[str, Any]: + return tomli.loads((self.path / "agent.toml").read_text(encoding="utf-8")) + + def declared_store(self, kind: str) -> str | None: + """The name agent.toml gives the ``memory`` or ``session`` store, if it declares one.""" + if not (self.path / "agent.toml").exists(): + return None + return self.manifest().get(f"{kind}_store", {}).get("name") + + 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_manifest(self, tool_specs: Sequence[ToolSpec]) -> None: + """Direct authoring: write agent.toml with exactly these tools plus a tracing binding.""" + sections: list[str] = [_DIRECT_HEADER.format(framework=FRAMEWORKS[0])] + sections.extend(spec.toml for spec in tool_specs if spec.toml) + # Follows default_experiment_name's shape so direct authoring carries the same tracing + # binding `agentbricks init` gives CLI authoring. + slug = re.sub(r"[^a-z0-9-]+", "-", self.path.name.lower()).strip("-") or "agent" + sections.append(f'[tracing]\nexperiment_name = "/Shared/agentbricks_traces/{slug}"\n') + target = self.path / "agent.toml" + log(f"# write {target}: direct authoring; no agentbricks tools command") + target.write_text("\n".join(sections), encoding="utf-8") + for spec in tool_specs: + for relative_path, content in spec.files.items(): + self.write_file(relative_path, content) + + +def pin_runtime(project: Project, run_config: RunConfig) -> None: + """Point the project's databricks-agentbricks dependency at the runtime under test. + + That is the wheel built from this checkout, or the bridge commit given by ``--bridge-sha``. + """ + if run_config.bridge_sha: + _pin_bridge_sources(project.path, run_config.bridge_sha) + else: + assert run_config.wheel is not None + _pin_project_wheel(project.path, run_config.wheel) + + +def is_pinned(project: Project) -> bool: + pyproject = project.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 _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") + + +def _pin_bridge_sources(project: pathlib.Path, bridge_sha: str) -> None: + pyproject = project / "pyproject.toml" + document = tomlkit.parse(pyproject.read_text(encoding="utf-8")) + dependencies = document["project"]["dependencies"] + if not any(str(dependency).startswith("databricks-langchain") for dependency in dependencies): + dependencies.append("databricks-langchain>=0.17.0") + if "tool" not in document: + document["tool"] = tomlkit.table() + if "uv" not in document["tool"]: + document["tool"]["uv"] = tomlkit.table() + if "sources" not in document["tool"]["uv"]: + document["tool"]["uv"]["sources"] = tomlkit.table() + for package, subdirectory in ( + ("databricks-agentbricks", "integrations/agentbricks"), + ("databricks-langchain", "integrations/langchain"), + ): + source = tomlkit.inline_table() + source.update( + { + "git": "https://github.com/databricks/databricks-ai-bridge.git", + "rev": bridge_sha, + "subdirectory": subdirectory, + } + ) + document["tool"]["uv"]["sources"][package] = source + log(f"# write {pyproject}: pin Agent Bricks and LangChain to {bridge_sha}") + pyproject.write_text(tomlkit.dumps(document), encoding="utf-8") diff --git a/integrations/agentbricks/tests/suites/memory_store/test_memory_store.py b/integrations/agentbricks/tests/suites/memory_store/test_memory_store.py index 682be0d80..8f7b6000c 100644 --- a/integrations/agentbricks/tests/suites/memory_store/test_memory_store.py +++ b/integrations/agentbricks/tests/suites/memory_store/test_memory_store.py @@ -1,24 +1,33 @@ -"""STUB: an `init` agent (default bindings); the memory store the CLI declared was created. +"""STUB: an `init` agent (default bindings); deploy created the memory store agent.toml declares. Real case still to write: the agent writes a memory in one invocation and recalls it in another. """ from __future__ import annotations +import json + import pytest from harness.cli import AgentbricksCli, AgentEndpoint +from harness.project import Project, pin_runtime from harness.workspace import Workspace -def test_memory_store_created(agentbricks_cli: AgentbricksCli, workspace_client: Workspace) -> None: - project = agentbricks_cli.new_project("cli") +def test_memory_store_created( + agentbricks_cli: AgentbricksCli, project: Project, workspace_client: Workspace +) -> None: + agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) + pin_runtime(project, agentbricks_cli.run_config) + declared = project.declared_store("memory") + assert declared, "agentbricks init declared no memory store" agentbricks_cli.run(f"deploy {project.app_name} --source {project.path}", timeout=2400) AgentEndpoint.deployed(workspace_client, project.app_name) - resource = project.memory_resource or {} - assert resource.get("display_name") == project.memory_store_config, resource - assert str(resource.get("name")).startswith("memory-stores/"), resource - assert project.memory_store_name == resource["name"] + store = json.loads( + agentbricks_cli.run(f"--output json memory stores get memory-stores/{declared}").stdout + ) + assert store.get("display_name") == declared, store + assert store.get("name") == f"memory-stores/{declared}", store @pytest.mark.xfail(reason="STUB: no memory write/recall case written yet", run=False) diff --git a/integrations/agentbricks/tests/suites/session_store/test_session_store.py b/integrations/agentbricks/tests/suites/session_store/test_session_store.py index 56a7d9709..e793778f0 100644 --- a/integrations/agentbricks/tests/suites/session_store/test_session_store.py +++ b/integrations/agentbricks/tests/suites/session_store/test_session_store.py @@ -1,4 +1,4 @@ -"""STUB: an `init` agent (default bindings); the session store the CLI declared was created. +"""STUB: an `init` agent (default bindings); deploy created the session store agent.toml declares. Real case still to write: conversation state persists across invocations of one session_id and is isolated between sessions. @@ -6,20 +6,26 @@ from __future__ import annotations +import json + import pytest from harness.cli import AgentbricksCli, AgentEndpoint +from harness.project import Project, pin_runtime from harness.workspace import Workspace def test_session_store_created( - agentbricks_cli: AgentbricksCli, workspace_client: Workspace + agentbricks_cli: AgentbricksCli, project: Project, workspace_client: Workspace ) -> None: - project = agentbricks_cli.new_project("cli") + agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) + pin_runtime(project, agentbricks_cli.run_config) + declared = project.declared_store("session") + assert declared, "agentbricks init declared no session store" agentbricks_cli.run(f"deploy {project.app_name} --source {project.path}", timeout=2400) AgentEndpoint.deployed(workspace_client, project.app_name) - resource = project.session_resource or {} - assert resource.get("session_store_name") == project.session_store_config, resource + store = json.loads(agentbricks_cli.run(f"--output json sessions stores get {declared}").stdout) + assert declared in json.dumps(store), store @pytest.mark.xfail(reason="STUB: no session-persistence case written yet", run=False) diff --git a/integrations/agentbricks/tests/suites/tools/test_deploy.py b/integrations/agentbricks/tests/suites/tools/test_deploy.py index 0f268fe02..a592f72e2 100644 --- a/integrations/agentbricks/tests/suites/tools/test_deploy.py +++ b/integrations/agentbricks/tests/suites/tools/test_deploy.py @@ -10,18 +10,21 @@ from harness.cli import AgentbricksCli, AgentEndpoint from harness.logs import log +from harness.project import Project, pin_runtime from harness.workspace import Workspace from tool_definitions import Tool, assert_endpoint_invoke, bind, invoke_and_check def test_deploy( agentbricks_cli: AgentbricksCli, + project: Project, workspace_client: Workspace, tools: tuple[Tool, ...], authoring: str, subtests, ) -> None: - project = agentbricks_cli.new_project(authoring) + agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) + pin_runtime(project, agentbricks_cli.run_config) app = project.app_name def assert_granted(tool: Tool) -> None: diff --git a/integrations/agentbricks/tests/suites/tools/test_dev.py b/integrations/agentbricks/tests/suites/tools/test_dev.py index 6f99b7b9f..dfc81a8a6 100644 --- a/integrations/agentbricks/tests/suites/tools/test_dev.py +++ b/integrations/agentbricks/tests/suites/tools/test_dev.py @@ -11,7 +11,8 @@ from collections.abc import Iterator import pytest -from harness.cli import AgentbricksCli, AgentEndpoint, Project, free_port +from harness.cli import AgentbricksCli, AgentEndpoint, free_port +from harness.project import Project, pin_runtime from tool_definitions import ( Tool, assert_endpoint_invoke, @@ -36,9 +37,14 @@ def dev(cli: AgentbricksCli, project: Project) -> Iterator[AgentEndpoint]: def test_dev( - agentbricks_cli: AgentbricksCli, tools: tuple[Tool, ...], authoring: str, subtests + agentbricks_cli: AgentbricksCli, + project: Project, + tools: tuple[Tool, ...], + authoring: str, + subtests, ) -> None: - project = agentbricks_cli.new_project(authoring) + agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) + pin_runtime(project, agentbricks_cli.run_config) if authoring == "cli": with subtests.test(step="rejects-unavailable-mcp"): diff --git a/integrations/agentbricks/tests/suites/tools/test_grants.py b/integrations/agentbricks/tests/suites/tools/test_grants.py index 92df4cc70..216ab0051 100644 --- a/integrations/agentbricks/tests/suites/tools/test_grants.py +++ b/integrations/agentbricks/tests/suites/tools/test_grants.py @@ -3,6 +3,7 @@ from __future__ import annotations from harness.cli import AgentbricksCli, AgentEndpoint +from harness.project import Project, pin_runtime from harness.run import TOOL_RESOURCE_PREFIX from harness.workspace import Workspace from tool_definitions import bind, uc_function_tool @@ -11,6 +12,7 @@ def test_grants( agentbricks_cli: AgentbricksCli, + project: Project, workspace_client: Workspace, uc_function: UcFunction, authoring: str, @@ -18,7 +20,8 @@ def test_grants( ) -> None: tool = uc_function_tool(uc_function) # `init` and direct authoring both bind tracing, a resource Agent Bricks does not own. - project = agentbricks_cli.new_project(authoring) + agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) + pin_runtime(project, agentbricks_cli.run_config) app = project.app_name if authoring == "cli": bind(agentbricks_cli, project, tool) diff --git a/integrations/agentbricks/tests/suites/tools/test_selection.py b/integrations/agentbricks/tests/suites/tools/test_selection.py index 24b9fde80..f10f1f5fb 100644 --- a/integrations/agentbricks/tests/suites/tools/test_selection.py +++ b/integrations/agentbricks/tests/suites/tools/test_selection.py @@ -7,6 +7,7 @@ from __future__ import annotations from harness.cli import AgentbricksCli, AgentEndpoint +from harness.project import Project, pin_runtime from harness.workspace import Workspace from tool_definitions import PYTHON_MARKER, UC_MARKER, Tool, bind, excerpt from uc_objects import UcFunction @@ -20,12 +21,14 @@ def test_selects_uc_function_not_python_marker( agentbricks_cli: AgentbricksCli, + project: Project, workspace_client: Workspace, tools: tuple[Tool, ...], uc_function: UcFunction, authoring: str, ) -> None: - project = agentbricks_cli.new_project(authoring) + agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) + pin_runtime(project, agentbricks_cli.run_config) app = project.app_name if authoring == "cli": for tool in tools: diff --git a/integrations/agentbricks/tests/suites/tools/tool_definitions.py b/integrations/agentbricks/tests/suites/tools/tool_definitions.py index 255d9aabd..2532c4bb2 100644 --- a/integrations/agentbricks/tests/suites/tools/tool_definitions.py +++ b/integrations/agentbricks/tests/suites/tools/tool_definitions.py @@ -8,7 +8,8 @@ import uuid from collections.abc import Callable, Mapping -from harness.cli import AgentbricksCli, AgentEndpoint, Project +from harness.cli import AgentbricksCli, AgentEndpoint +from harness.project import Project from harness.workspace import GrantTuple from uc_objects import UcFunction diff --git a/integrations/agentbricks/tests/suites/tracing/test_tracing.py b/integrations/agentbricks/tests/suites/tracing/test_tracing.py index 85d807af0..d60935509 100644 --- a/integrations/agentbricks/tests/suites/tracing/test_tracing.py +++ b/integrations/agentbricks/tests/suites/tracing/test_tracing.py @@ -7,17 +7,20 @@ import pytest from harness.cli import AgentbricksCli, AgentEndpoint +from harness.project import Project, pin_runtime from harness.run import TOOL_RESOURCE_PREFIX from harness.workspace import Workspace def test_tracing_experiment_bound( - agentbricks_cli: AgentbricksCli, workspace_client: Workspace + agentbricks_cli: AgentbricksCli, project: Project, workspace_client: Workspace ) -> None: - project = agentbricks_cli.new_project("cli") + agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) + pin_runtime(project, agentbricks_cli.run_config) + experiment_name = project.manifest().get("tracing", {}).get("experiment_name") agentbricks_cli.run(f"deploy {project.app_name} --source {project.path}", timeout=2400) AgentEndpoint.deployed(workspace_client, project.app_name) - assert project.experiment_name, "agent.toml binds no tracing experiment" + assert experiment_name, "agent.toml binds no tracing experiment" # Deploy attaches the experiment as an App resource that Agent Bricks does not own. experiments = [ @@ -25,9 +28,9 @@ def test_tracing_experiment_bound( ] assert experiments, "No experiment App resource among the non-tool resources" assert not any(r["name"].startswith(TOOL_RESOURCE_PREFIX) for r in experiments) - expected_id = workspace_client.experiment_id(project.experiment_name) + expected_id = workspace_client.experiment_id(experiment_name) assert any(str(r["experiment"].get("experiment_id")) == expected_id for r in experiments), ( - f"{project.experiment_name} ({expected_id}) not in {experiments}" + f"{experiment_name} ({expected_id}) not in {experiments}" ) From 8452bb40883a4bfa76c2749c90a3bc9e38dbbbc3 Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Wed, 14 Oct 2026 18:54:37 +0000 Subject: [PATCH 18/45] [agentbricks] E2E suites: move AgentEndpoint to harness/agent_endpoint.py Co-authored-by: Isaac --- .../agentbricks/tests/suites/README.md | 3 +- .../tests/suites/harness/agent_endpoint.py | 166 ++++++++++++++++++ .../agentbricks/tests/suites/harness/cli.py | 163 +---------------- .../suites/memory_store/test_memory_store.py | 3 +- .../session_store/test_session_store.py | 3 +- .../tests/suites/tools/test_deploy.py | 3 +- .../tests/suites/tools/test_dev.py | 3 +- .../tests/suites/tools/test_grants.py | 3 +- .../tests/suites/tools/test_selection.py | 3 +- .../tests/suites/tools/tool_definitions.py | 3 +- .../tests/suites/tracing/test_tracing.py | 3 +- 11 files changed, 189 insertions(+), 167 deletions(-) create mode 100644 integrations/agentbricks/tests/suites/harness/agent_endpoint.py diff --git a/integrations/agentbricks/tests/suites/README.md b/integrations/agentbricks/tests/suites/README.md index 8119871da..73e1f38e4 100644 --- a/integrations/agentbricks/tests/suites/README.md +++ b/integrations/agentbricks/tests/suites/README.md @@ -49,7 +49,8 @@ Every resource a run creates carries its run id, and is deleted at three points: | `conftest.py` | All pytest wiring: options, hooks, and the fixtures every suite shares | | `harness/run.py` | `RunConfig`, the run id, and the name prefixes every resource carries | | `harness/target.py` | The `--workspace live\|fake` seam the CLI and fixtures talk to | -| `harness/cli.py` | Runs `agentbricks` commands (`cli.run("tools add ...")`), and the endpoints of the agents they start | +| `harness/cli.py` | Runs `agentbricks` commands: `cli.run("tools add ...")`, and `cli.start("dev ...")` in the background | +| `harness/agent_endpoint.py` | `AgentEndpoint`: a running dev server or deployed App, and invoking it over HTTP | | `harness/project.py` | A test's project: run-scoped path and App name, its files, and `pin_runtime` | | `harness/workspace.py` | Databricks SDK helpers for what is not under test: SQL, UC, Apps, grants | | `harness/cleanup.py` | Deleting what the suites create, per App and per run | diff --git a/integrations/agentbricks/tests/suites/harness/agent_endpoint.py b/integrations/agentbricks/tests/suites/harness/agent_endpoint.py new file mode 100644 index 000000000..189f4643c --- /dev/null +++ b/integrations/agentbricks/tests/suites/harness/agent_endpoint.py @@ -0,0 +1,166 @@ +"""A running agent, a dev server or a deployed App, and how a test calls it over HTTP. + +``AgentEndpoint.invoke`` POSTs straight to the agent's /api/invocations, so it is the one test step +that does not go through the CLI; ``agentbricks endpoint invoke`` is checked separately. +""" + +from __future__ import annotations + +import concurrent.futures +import json +import time +import urllib.error +import urllib.request +import uuid +from collections.abc import Callable +from typing import Any + +from harness.cli import BackgroundCommand +from harness.logs import last_lines, last_nonempty_line, log +from harness.run import E2E_MODEL, E2EError, now +from harness.workspace import Workspace + + +class AgentEndpoint: + """Where a running agent answers, a dev server or a deployed App, and how to call it.""" + + def __init__( + self, + base_url: str, + headers: Callable[[], dict[str, str]], + *, + app_name: str | None = None, + app: dict[str, Any] | None = None, + ): + self.base_url = base_url + self.app_name = app_name + self.app = app + self._headers = headers + + @classmethod + def local(cls, server: BackgroundCommand, port: int) -> AgentEndpoint: + """The `agentbricks dev` server on ``port``, once it answers; fails if the server exits.""" + started = time.monotonic() + next_tick = 60.0 + while True: + if server.process.poll() is not None: + raise E2EError( + f"{server.label} exited {server.process.returncode}: " + f"{last_lines(server.log_path, 30)}" + ) + try: + with urllib.request.urlopen(f"http://127.0.0.1:{port}/", timeout=5): + break + except urllib.error.HTTPError as exc: + if exc.code < 500: + break + except (urllib.error.URLError, TimeoutError): + pass + elapsed = time.monotonic() - started + if elapsed > 1200: + raise E2EError( + f"{server.label} did not become reachable: {last_lines(server.log_path, 30)}" + ) + if elapsed >= next_tick: + log( + f"tick {now():%H:%M} | {server.label} | starting | " + f"{last_nonempty_line(server.log_path)}" + ) + next_tick += 60 + time.sleep(5) + return cls(f"http://127.0.0.1:{port}", lambda: {}) + + @classmethod + def deployed(cls, workspace: Workspace, app_name: str) -> AgentEndpoint: + """The App once it is up and has answered once.""" + workspace.check_app_auth() + app = workspace.wait_for_app(app_name) + url = str(app.get("url") or "").rstrip("/") + if not url: + raise E2EError(f"App {app_name} has no URL: {app}") + endpoint = cls(url, lambda: workspace.app_headers, app_name=app_name, app=app) + endpoint.warm_up() + return endpoint + + @property + def runtime(self) -> str: + return "deploy" if self.app_name else "dev" + + def invoke(self, prompt: str, label: str = "invoke") -> str: + """POST the prompt to /api/invocations and return the serialized response.""" + url = f"{self.base_url}/api/invocations" + last: Exception | None = None + invocation_id = str(uuid.uuid4()) + body = { + "id": invocation_id, + # session_id is a top-level invocation field; the adapter rejects it nested in input. + "session_id": invocation_id, + "input": { + "model": E2E_MODEL, + "messages": [{"role": "user", "content": prompt}], + }, + } + for attempt in range(1, 4): + try: + response = _monitored( + label, lambda: _http_json(url, body, self._headers()), timeout=360 + ) + return json.dumps(response, sort_keys=True, default=str) + except Exception as exc: + last = exc + log(f"attempt {attempt}/3 | {label} | {exc}") + if attempt < 3: + time.sleep(15) + raise E2EError(f"{label} failed after 3 attempts: {last}") + + def warm_up(self, deadline_seconds: float = 180.0) -> None: + """Invoke until the agent answers once. A new App's session store can drop the first calls, + and a redeployed one restarts.""" + started = time.monotonic() + while True: + try: + self.invoke("Reply with the single word OK.", label="warm-up") + return + except E2EError: + if time.monotonic() - started > deadline_seconds: + raise + + +def _monitored( + label: str, operation: Callable[[], dict[str, Any]], *, timeout: float +) -> dict[str, Any]: + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(operation) + started = time.monotonic() + while True: + try: + return future.result( + timeout=min(60, max(1, timeout - (time.monotonic() - started))) + ) + except concurrent.futures.TimeoutError: + elapsed = time.monotonic() - started + log(f"tick {now():%H:%M} | {label} | running | {elapsed:.0f}s") + if elapsed >= timeout: + raise E2EError(f"{label} timed out after {timeout:.0f}s") from None + + +def _http_json(url: str, body: dict[str, Any], headers: dict[str, str]) -> dict[str, Any]: + request = urllib.request.Request( + url, + data=json.dumps(body).encode(), + headers={"Content-Type": "application/json", **headers}, + method="POST", + ) + try: + with urllib.request.urlopen(request, timeout=340) as response: + payload = response.read().decode() + except urllib.error.HTTPError as exc: + detail = exc.read().decode(errors="replace") + raise E2EError(f"HTTP {exc.code} from {url}: {detail}") from exc + try: + value = json.loads(payload) + except json.JSONDecodeError as exc: + raise E2EError(f"Invalid JSON from {url}: {payload[:2000]}") from exc + if not isinstance(value, dict): + raise E2EError(f"Expected object response from {url}, got {type(value).__name__}") + return value diff --git a/integrations/agentbricks/tests/suites/harness/cli.py b/integrations/agentbricks/tests/suites/harness/cli.py index 9a2f6b028..a3db5ba6c 100644 --- a/integrations/agentbricks/tests/suites/harness/cli.py +++ b/integrations/agentbricks/tests/suites/harness/cli.py @@ -1,17 +1,14 @@ """The ``agentbricks`` CLI and the agents it produces: the system under test. Every CLI 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. The only -non-subprocess part is ``AgentEndpoint.invoke``, which POSTs to a running agent. Nothing in this -module knows about any particular suite or deletes anything: the ``project`` fixture does. +(``cli.run("tools add ...")``), so a regression in the CLI fails the test that makes it. Nothing in +this module knows about any particular suite or deletes anything: the ``project`` fixture does. """ from __future__ import annotations -import concurrent.futures import contextlib import dataclasses -import json import os import pathlib import re @@ -22,16 +19,11 @@ import subprocess import sys import time -import urllib.error -import urllib.request -import uuid -from collections.abc import Callable, Iterator, Mapping -from typing import Any +from collections.abc import Iterator, Mapping -from harness.logs import last_lines, last_nonempty_line, log, log_command -from harness.run import E2E_MODEL, E2EError, RunConfig, now +from harness.logs import last_nonempty_line, log, log_command +from harness.run import E2EError, RunConfig, now from harness.target import TargetWorkspace -from harness.workspace import Workspace # The Apps gateway's answers while an App is not up, e.g. restarting after a redeploy. _GATEWAY_UNAVAILABLE = re.compile(r"HTTP 50[234]\b") @@ -46,111 +38,6 @@ class BackgroundCommand: log_path: pathlib.Path -class AgentEndpoint: - """Where a running agent answers, a dev server or a deployed App, and how to call it.""" - - def __init__( - self, - base_url: str, - headers: Callable[[], dict[str, str]], - *, - app_name: str | None = None, - app: dict[str, Any] | None = None, - ): - self.base_url = base_url - self.app_name = app_name - self.app = app - self._headers = headers - - @classmethod - def local(cls, server: BackgroundCommand, port: int) -> AgentEndpoint: - """The `agentbricks dev` server on ``port``, once it answers; fails if the server exits.""" - started = time.monotonic() - next_tick = 60.0 - while True: - if server.process.poll() is not None: - raise E2EError( - f"{server.label} exited {server.process.returncode}: " - f"{last_lines(server.log_path, 30)}" - ) - try: - with urllib.request.urlopen(f"http://127.0.0.1:{port}/", timeout=5): - break - except urllib.error.HTTPError as exc: - if exc.code < 500: - break - except (urllib.error.URLError, TimeoutError): - pass - elapsed = time.monotonic() - started - if elapsed > 1200: - raise E2EError( - f"{server.label} did not become reachable: {last_lines(server.log_path, 30)}" - ) - if elapsed >= next_tick: - log( - f"tick {now():%H:%M} | {server.label} | starting | " - f"{last_nonempty_line(server.log_path)}" - ) - next_tick += 60 - time.sleep(5) - return cls(f"http://127.0.0.1:{port}", lambda: {}) - - @classmethod - def deployed(cls, workspace: Workspace, app_name: str) -> AgentEndpoint: - """The App once it is up and has answered once.""" - workspace.check_app_auth() - app = workspace.wait_for_app(app_name) - url = str(app.get("url") or "").rstrip("/") - if not url: - raise E2EError(f"App {app_name} has no URL: {app}") - endpoint = cls(url, lambda: workspace.app_headers, app_name=app_name, app=app) - endpoint.warm_up() - return endpoint - - @property - def runtime(self) -> str: - return "deploy" if self.app_name else "dev" - - def invoke(self, prompt: str, label: str = "invoke") -> str: - """POST the prompt to /api/invocations and return the serialized response.""" - url = f"{self.base_url}/api/invocations" - last: Exception | None = None - invocation_id = str(uuid.uuid4()) - body = { - "id": invocation_id, - # session_id is a top-level invocation field; the adapter rejects it nested in input. - "session_id": invocation_id, - "input": { - "model": E2E_MODEL, - "messages": [{"role": "user", "content": prompt}], - }, - } - for attempt in range(1, 4): - try: - response = _monitored( - label, lambda: _http_json(url, body, self._headers()), timeout=360 - ) - return json.dumps(response, sort_keys=True, default=str) - except Exception as exc: - last = exc - log(f"attempt {attempt}/3 | {label} | {exc}") - if attempt < 3: - time.sleep(15) - raise E2EError(f"{label} failed after 3 attempts: {last}") - - def warm_up(self, deadline_seconds: float = 180.0) -> None: - """Invoke until the agent answers once. A new App's session store can drop the first calls, - and a redeployed one restarts.""" - started = time.monotonic() - while True: - try: - self.invoke("Reply with the single word OK.", label="warm-up") - return - except E2EError: - if time.monotonic() - started > deadline_seconds: - raise - - class AgentbricksCli: """Runs ``agentbricks`` commands against a target workspace, logging each to its own files.""" @@ -305,43 +192,3 @@ def free_port() -> int: with socket.socket() as probe: probe.bind(("127.0.0.1", 0)) return probe.getsockname()[1] - - -def _monitored( - label: str, operation: Callable[[], dict[str, Any]], *, timeout: float -) -> dict[str, Any]: - with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor: - future = executor.submit(operation) - started = time.monotonic() - while True: - try: - return future.result( - timeout=min(60, max(1, timeout - (time.monotonic() - started))) - ) - except concurrent.futures.TimeoutError: - elapsed = time.monotonic() - started - log(f"tick {now():%H:%M} | {label} | running | {elapsed:.0f}s") - if elapsed >= timeout: - raise E2EError(f"{label} timed out after {timeout:.0f}s") from None - - -def _http_json(url: str, body: dict[str, Any], headers: dict[str, str]) -> dict[str, Any]: - request = urllib.request.Request( - url, - data=json.dumps(body).encode(), - headers={"Content-Type": "application/json", **headers}, - method="POST", - ) - try: - with urllib.request.urlopen(request, timeout=340) as response: - payload = response.read().decode() - except urllib.error.HTTPError as exc: - detail = exc.read().decode(errors="replace") - raise E2EError(f"HTTP {exc.code} from {url}: {detail}") from exc - try: - value = json.loads(payload) - except json.JSONDecodeError as exc: - raise E2EError(f"Invalid JSON from {url}: {payload[:2000]}") from exc - if not isinstance(value, dict): - raise E2EError(f"Expected object response from {url}, got {type(value).__name__}") - return value diff --git a/integrations/agentbricks/tests/suites/memory_store/test_memory_store.py b/integrations/agentbricks/tests/suites/memory_store/test_memory_store.py index 8f7b6000c..d0a776338 100644 --- a/integrations/agentbricks/tests/suites/memory_store/test_memory_store.py +++ b/integrations/agentbricks/tests/suites/memory_store/test_memory_store.py @@ -8,7 +8,8 @@ import json import pytest -from harness.cli import AgentbricksCli, AgentEndpoint +from harness.agent_endpoint import AgentEndpoint +from harness.cli import AgentbricksCli from harness.project import Project, pin_runtime from harness.workspace import Workspace diff --git a/integrations/agentbricks/tests/suites/session_store/test_session_store.py b/integrations/agentbricks/tests/suites/session_store/test_session_store.py index e793778f0..2e6e64673 100644 --- a/integrations/agentbricks/tests/suites/session_store/test_session_store.py +++ b/integrations/agentbricks/tests/suites/session_store/test_session_store.py @@ -9,7 +9,8 @@ import json import pytest -from harness.cli import AgentbricksCli, AgentEndpoint +from harness.agent_endpoint import AgentEndpoint +from harness.cli import AgentbricksCli from harness.project import Project, pin_runtime from harness.workspace import Workspace diff --git a/integrations/agentbricks/tests/suites/tools/test_deploy.py b/integrations/agentbricks/tests/suites/tools/test_deploy.py index a592f72e2..fc4cc9e01 100644 --- a/integrations/agentbricks/tests/suites/tools/test_deploy.py +++ b/integrations/agentbricks/tests/suites/tools/test_deploy.py @@ -8,7 +8,8 @@ from __future__ import annotations -from harness.cli import AgentbricksCli, AgentEndpoint +from harness.agent_endpoint import AgentEndpoint +from harness.cli import AgentbricksCli from harness.logs import log from harness.project import Project, pin_runtime from harness.workspace import Workspace diff --git a/integrations/agentbricks/tests/suites/tools/test_dev.py b/integrations/agentbricks/tests/suites/tools/test_dev.py index dfc81a8a6..292921a9d 100644 --- a/integrations/agentbricks/tests/suites/tools/test_dev.py +++ b/integrations/agentbricks/tests/suites/tools/test_dev.py @@ -11,7 +11,8 @@ from collections.abc import Iterator import pytest -from harness.cli import AgentbricksCli, AgentEndpoint, free_port +from harness.agent_endpoint import AgentEndpoint +from harness.cli import AgentbricksCli, free_port from harness.project import Project, pin_runtime from tool_definitions import ( Tool, diff --git a/integrations/agentbricks/tests/suites/tools/test_grants.py b/integrations/agentbricks/tests/suites/tools/test_grants.py index 216ab0051..7baaf8fd9 100644 --- a/integrations/agentbricks/tests/suites/tools/test_grants.py +++ b/integrations/agentbricks/tests/suites/tools/test_grants.py @@ -2,7 +2,8 @@ from __future__ import annotations -from harness.cli import AgentbricksCli, AgentEndpoint +from harness.agent_endpoint import AgentEndpoint +from harness.cli import AgentbricksCli from harness.project import Project, pin_runtime from harness.run import TOOL_RESOURCE_PREFIX from harness.workspace import Workspace diff --git a/integrations/agentbricks/tests/suites/tools/test_selection.py b/integrations/agentbricks/tests/suites/tools/test_selection.py index f10f1f5fb..a2616bd60 100644 --- a/integrations/agentbricks/tests/suites/tools/test_selection.py +++ b/integrations/agentbricks/tests/suites/tools/test_selection.py @@ -6,7 +6,8 @@ from __future__ import annotations -from harness.cli import AgentbricksCli, AgentEndpoint +from harness.agent_endpoint import AgentEndpoint +from harness.cli import AgentbricksCli from harness.project import Project, pin_runtime from harness.workspace import Workspace from tool_definitions import PYTHON_MARKER, UC_MARKER, Tool, bind, excerpt diff --git a/integrations/agentbricks/tests/suites/tools/tool_definitions.py b/integrations/agentbricks/tests/suites/tools/tool_definitions.py index 2532c4bb2..caead37b9 100644 --- a/integrations/agentbricks/tests/suites/tools/tool_definitions.py +++ b/integrations/agentbricks/tests/suites/tools/tool_definitions.py @@ -8,7 +8,8 @@ import uuid from collections.abc import Callable, Mapping -from harness.cli import AgentbricksCli, AgentEndpoint +from harness.agent_endpoint import AgentEndpoint +from harness.cli import AgentbricksCli from harness.project import Project from harness.workspace import GrantTuple from uc_objects import UcFunction diff --git a/integrations/agentbricks/tests/suites/tracing/test_tracing.py b/integrations/agentbricks/tests/suites/tracing/test_tracing.py index d60935509..6e2bc12c2 100644 --- a/integrations/agentbricks/tests/suites/tracing/test_tracing.py +++ b/integrations/agentbricks/tests/suites/tracing/test_tracing.py @@ -6,7 +6,8 @@ from __future__ import annotations import pytest -from harness.cli import AgentbricksCli, AgentEndpoint +from harness.agent_endpoint import AgentEndpoint +from harness.cli import AgentbricksCli from harness.project import Project, pin_runtime from harness.run import TOOL_RESOURCE_PREFIX from harness.workspace import Workspace From 4016b0141f7d8de0038d3939a4eb31f0bbaa63a6 Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Wed, 14 Oct 2026 18:57:06 +0000 Subject: [PATCH 19/45] [agentbricks] E2E suites: rename harness/cli.py to agentbricks_cli.py; make pin_runtime a Project method Co-authored-by: Isaac --- .../agentbricks/tests/suites/README.md | 6 +-- .../agentbricks/tests/suites/conftest.py | 6 +-- .../tests/suites/harness/agent_endpoint.py | 2 +- .../harness/{cli.py => agentbricks_cli.py} | 0 .../tests/suites/harness/project.py | 38 +++++++++---------- .../suites/memory_store/test_memory_store.py | 6 +-- .../session_store/test_session_store.py | 6 +-- .../tests/suites/tools/test_deploy.py | 6 +-- .../tests/suites/tools/test_dev.py | 6 +-- .../tests/suites/tools/test_grants.py | 6 +-- .../tests/suites/tools/test_selection.py | 6 +-- .../tests/suites/tools/tool_definitions.py | 2 +- .../tests/suites/tracing/test_tracing.py | 6 +-- 13 files changed, 47 insertions(+), 49 deletions(-) rename integrations/agentbricks/tests/suites/harness/{cli.py => agentbricks_cli.py} (100%) diff --git a/integrations/agentbricks/tests/suites/README.md b/integrations/agentbricks/tests/suites/README.md index 73e1f38e4..9d0891fba 100644 --- a/integrations/agentbricks/tests/suites/README.md +++ b/integrations/agentbricks/tests/suites/README.md @@ -49,9 +49,9 @@ Every resource a run creates carries its run id, and is deleted at three points: | `conftest.py` | All pytest wiring: options, hooks, and the fixtures every suite shares | | `harness/run.py` | `RunConfig`, the run id, and the name prefixes every resource carries | | `harness/target.py` | The `--workspace live\|fake` seam the CLI and fixtures talk to | -| `harness/cli.py` | Runs `agentbricks` commands: `cli.run("tools add ...")`, and `cli.start("dev ...")` in the background | +| `harness/agentbricks_cli.py` | Runs `agentbricks` commands: `cli.run("tools add ...")`, and `cli.start("dev ...")` in the background | | `harness/agent_endpoint.py` | `AgentEndpoint`: a running dev server or deployed App, and invoking it over HTTP | -| `harness/project.py` | A test's project: run-scoped path and App name, its files, and `pin_runtime` | +| `harness/project.py` | A test's project: run-scoped path and App name, its files, and `pin_runtime()` | | `harness/workspace.py` | Databricks SDK helpers for what is not under test: SQL, UC, Apps, grants | | `harness/cleanup.py` | Deleting what the suites create, per App and per run | | `harness/logs.py` | Test-log output; every command is echoed so a run's log replays it | @@ -65,6 +65,6 @@ that you can follow through its imports. - **A tool:** add one `ToolDefinition(name, requires, build)` in `tools/tool_definitions.py`. `requires` names the fixtures it needs; tools whose fixtures skip are left out and logged. - **A suite:** add a folder with `test_*.py` files. A test takes the `project` fixture, runs - `agentbricks_cli.run(f"init ... {project.path}")`, then `pin_runtime(project, ...)` so the project + `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. Add a `conftest.py` only for resources the CLI under test does not create itself; create them per test and delete them in the fixture's teardown. diff --git a/integrations/agentbricks/tests/suites/conftest.py b/integrations/agentbricks/tests/suites/conftest.py index 8cba524c7..f4ca78567 100644 --- a/integrations/agentbricks/tests/suites/conftest.py +++ b/integrations/agentbricks/tests/suites/conftest.py @@ -26,10 +26,10 @@ from typing import Any import pytest +from harness.agentbricks_cli import AgentbricksCli from harness.cleanup import delete_app_resources, delete_run_leftovers, record_app_identity -from harness.cli import AgentbricksCli from harness.logs import log, log_command -from harness.project import Project, is_pinned +from harness.project import Project from harness.run import RUN_ID_PATTERN, E2EError, RunConfig, run_output_dir from harness.target import ( WORKSPACE_KINDS, @@ -253,7 +253,7 @@ def project( if workspace is not None: _delete_project(agentbricks_cli, workspace, project, 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 is_pinned(project): + 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." ) diff --git a/integrations/agentbricks/tests/suites/harness/agent_endpoint.py b/integrations/agentbricks/tests/suites/harness/agent_endpoint.py index 189f4643c..df935fa03 100644 --- a/integrations/agentbricks/tests/suites/harness/agent_endpoint.py +++ b/integrations/agentbricks/tests/suites/harness/agent_endpoint.py @@ -15,7 +15,7 @@ from collections.abc import Callable from typing import Any -from harness.cli import BackgroundCommand +from harness.agentbricks_cli import BackgroundCommand from harness.logs import last_lines, last_nonempty_line, log from harness.run import E2E_MODEL, E2EError, now from harness.workspace import Workspace diff --git a/integrations/agentbricks/tests/suites/harness/cli.py b/integrations/agentbricks/tests/suites/harness/agentbricks_cli.py similarity index 100% rename from integrations/agentbricks/tests/suites/harness/cli.py rename to integrations/agentbricks/tests/suites/harness/agentbricks_cli.py diff --git a/integrations/agentbricks/tests/suites/harness/project.py b/integrations/agentbricks/tests/suites/harness/project.py index 06bd4bfe9..fa219e46e 100644 --- a/integrations/agentbricks/tests/suites/harness/project.py +++ b/integrations/agentbricks/tests/suites/harness/project.py @@ -36,6 +36,9 @@ class ToolSpec(Protocol): class Project: path: pathlib.Path app_name: str + # The runtime under test: the wheel built from this checkout, or a bridge commit. + wheel: pathlib.Path | None = None + bridge_sha: str | None = None @classmethod def allocate(cls, run_config: RunConfig) -> Project: @@ -43,7 +46,21 @@ def allocate(cls, run_config: RunConfig) -> Project: 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) - return cls(path, f"{APP_PREFIX}{name}") + return cls(path, f"{APP_PREFIX}{name}", run_config.wheel, run_config.bridge_sha) + + def pin_runtime(self) -> None: + """Point the project's databricks-agentbricks dependency at the runtime under test.""" + if self.bridge_sha: + _pin_bridge_sources(self.path, self.bridge_sha) + else: + assert self.wheel is not None + _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 manifest(self) -> dict[str, Any]: return tomli.loads((self.path / "agent.toml").read_text(encoding="utf-8")) @@ -77,25 +94,6 @@ def write_manifest(self, tool_specs: Sequence[ToolSpec]) -> None: self.write_file(relative_path, content) -def pin_runtime(project: Project, run_config: RunConfig) -> None: - """Point the project's databricks-agentbricks dependency at the runtime under test. - - That is the wheel built from this checkout, or the bridge commit given by ``--bridge-sha``. - """ - if run_config.bridge_sha: - _pin_bridge_sources(project.path, run_config.bridge_sha) - else: - assert run_config.wheel is not None - _pin_project_wheel(project.path, run_config.wheel) - - -def is_pinned(project: Project) -> bool: - pyproject = project.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 _pin_project_wheel(project: pathlib.Path, wheel: pathlib.Path) -> None: vendor_dir = project / "agentbricks_e2e_wheels" vendor_dir.mkdir() diff --git a/integrations/agentbricks/tests/suites/memory_store/test_memory_store.py b/integrations/agentbricks/tests/suites/memory_store/test_memory_store.py index d0a776338..88d9ba651 100644 --- a/integrations/agentbricks/tests/suites/memory_store/test_memory_store.py +++ b/integrations/agentbricks/tests/suites/memory_store/test_memory_store.py @@ -9,8 +9,8 @@ import pytest from harness.agent_endpoint import AgentEndpoint -from harness.cli import AgentbricksCli -from harness.project import Project, pin_runtime +from harness.agentbricks_cli import AgentbricksCli +from harness.project import Project from harness.workspace import Workspace @@ -18,7 +18,7 @@ def test_memory_store_created( agentbricks_cli: AgentbricksCli, project: Project, workspace_client: Workspace ) -> None: agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) - pin_runtime(project, agentbricks_cli.run_config) + project.pin_runtime() declared = project.declared_store("memory") assert declared, "agentbricks init declared no memory store" agentbricks_cli.run(f"deploy {project.app_name} --source {project.path}", timeout=2400) diff --git a/integrations/agentbricks/tests/suites/session_store/test_session_store.py b/integrations/agentbricks/tests/suites/session_store/test_session_store.py index 2e6e64673..a2ebde65b 100644 --- a/integrations/agentbricks/tests/suites/session_store/test_session_store.py +++ b/integrations/agentbricks/tests/suites/session_store/test_session_store.py @@ -10,8 +10,8 @@ import pytest from harness.agent_endpoint import AgentEndpoint -from harness.cli import AgentbricksCli -from harness.project import Project, pin_runtime +from harness.agentbricks_cli import AgentbricksCli +from harness.project import Project from harness.workspace import Workspace @@ -19,7 +19,7 @@ def test_session_store_created( agentbricks_cli: AgentbricksCli, project: Project, workspace_client: Workspace ) -> None: agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) - pin_runtime(project, agentbricks_cli.run_config) + project.pin_runtime() declared = project.declared_store("session") assert declared, "agentbricks init declared no session store" agentbricks_cli.run(f"deploy {project.app_name} --source {project.path}", timeout=2400) diff --git a/integrations/agentbricks/tests/suites/tools/test_deploy.py b/integrations/agentbricks/tests/suites/tools/test_deploy.py index fc4cc9e01..e8bea0fea 100644 --- a/integrations/agentbricks/tests/suites/tools/test_deploy.py +++ b/integrations/agentbricks/tests/suites/tools/test_deploy.py @@ -9,9 +9,9 @@ from __future__ import annotations from harness.agent_endpoint import AgentEndpoint -from harness.cli import AgentbricksCli +from harness.agentbricks_cli import AgentbricksCli from harness.logs import log -from harness.project import Project, pin_runtime +from harness.project import Project from harness.workspace import Workspace from tool_definitions import Tool, assert_endpoint_invoke, bind, invoke_and_check @@ -25,7 +25,7 @@ def test_deploy( subtests, ) -> None: agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) - pin_runtime(project, agentbricks_cli.run_config) + project.pin_runtime() app = project.app_name def assert_granted(tool: Tool) -> None: diff --git a/integrations/agentbricks/tests/suites/tools/test_dev.py b/integrations/agentbricks/tests/suites/tools/test_dev.py index 292921a9d..a23717900 100644 --- a/integrations/agentbricks/tests/suites/tools/test_dev.py +++ b/integrations/agentbricks/tests/suites/tools/test_dev.py @@ -12,8 +12,8 @@ import pytest from harness.agent_endpoint import AgentEndpoint -from harness.cli import AgentbricksCli, free_port -from harness.project import Project, pin_runtime +from harness.agentbricks_cli import AgentbricksCli, free_port +from harness.project import Project from tool_definitions import ( Tool, assert_endpoint_invoke, @@ -45,7 +45,7 @@ def test_dev( subtests, ) -> None: agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) - pin_runtime(project, agentbricks_cli.run_config) + project.pin_runtime() if authoring == "cli": with subtests.test(step="rejects-unavailable-mcp"): diff --git a/integrations/agentbricks/tests/suites/tools/test_grants.py b/integrations/agentbricks/tests/suites/tools/test_grants.py index 7baaf8fd9..8bb88efd3 100644 --- a/integrations/agentbricks/tests/suites/tools/test_grants.py +++ b/integrations/agentbricks/tests/suites/tools/test_grants.py @@ -3,8 +3,8 @@ from __future__ import annotations from harness.agent_endpoint import AgentEndpoint -from harness.cli import AgentbricksCli -from harness.project import Project, pin_runtime +from harness.agentbricks_cli import AgentbricksCli +from harness.project import Project from harness.run import TOOL_RESOURCE_PREFIX from harness.workspace import Workspace from tool_definitions import bind, uc_function_tool @@ -22,7 +22,7 @@ def test_grants( tool = uc_function_tool(uc_function) # `init` and direct authoring both bind tracing, a resource Agent Bricks does not own. agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) - pin_runtime(project, agentbricks_cli.run_config) + project.pin_runtime() app = project.app_name if authoring == "cli": bind(agentbricks_cli, project, tool) diff --git a/integrations/agentbricks/tests/suites/tools/test_selection.py b/integrations/agentbricks/tests/suites/tools/test_selection.py index a2616bd60..40a35c2c2 100644 --- a/integrations/agentbricks/tests/suites/tools/test_selection.py +++ b/integrations/agentbricks/tests/suites/tools/test_selection.py @@ -7,8 +7,8 @@ from __future__ import annotations from harness.agent_endpoint import AgentEndpoint -from harness.cli import AgentbricksCli -from harness.project import Project, pin_runtime +from harness.agentbricks_cli import AgentbricksCli +from harness.project import Project from harness.workspace import Workspace from tool_definitions import PYTHON_MARKER, UC_MARKER, Tool, bind, excerpt from uc_objects import UcFunction @@ -29,7 +29,7 @@ def test_selects_uc_function_not_python_marker( authoring: str, ) -> None: agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) - pin_runtime(project, agentbricks_cli.run_config) + project.pin_runtime() app = project.app_name if authoring == "cli": for tool in tools: diff --git a/integrations/agentbricks/tests/suites/tools/tool_definitions.py b/integrations/agentbricks/tests/suites/tools/tool_definitions.py index caead37b9..46253b55b 100644 --- a/integrations/agentbricks/tests/suites/tools/tool_definitions.py +++ b/integrations/agentbricks/tests/suites/tools/tool_definitions.py @@ -9,7 +9,7 @@ from collections.abc import Callable, Mapping from harness.agent_endpoint import AgentEndpoint -from harness.cli import AgentbricksCli +from harness.agentbricks_cli import AgentbricksCli from harness.project import Project from harness.workspace import GrantTuple from uc_objects import UcFunction diff --git a/integrations/agentbricks/tests/suites/tracing/test_tracing.py b/integrations/agentbricks/tests/suites/tracing/test_tracing.py index 6e2bc12c2..75d6f371a 100644 --- a/integrations/agentbricks/tests/suites/tracing/test_tracing.py +++ b/integrations/agentbricks/tests/suites/tracing/test_tracing.py @@ -7,8 +7,8 @@ import pytest from harness.agent_endpoint import AgentEndpoint -from harness.cli import AgentbricksCli -from harness.project import Project, pin_runtime +from harness.agentbricks_cli import AgentbricksCli +from harness.project import Project from harness.run import TOOL_RESOURCE_PREFIX from harness.workspace import Workspace @@ -17,7 +17,7 @@ def test_tracing_experiment_bound( agentbricks_cli: AgentbricksCli, project: Project, workspace_client: Workspace ) -> None: agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) - pin_runtime(project, agentbricks_cli.run_config) + project.pin_runtime() experiment_name = project.manifest().get("tracing", {}).get("experiment_name") agentbricks_cli.run(f"deploy {project.app_name} --source {project.path}", timeout=2400) AgentEndpoint.deployed(workspace_client, project.app_name) From 916b0dcc875a977b682415b9e04434906bc385c5 Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Wed, 14 Oct 2026 20:06:29 +0000 Subject: [PATCH 20/45] [agentbricks] E2E suites: write every CLI step in tests as a command; drop unused App waits Co-authored-by: Isaac --- .../tests/suites/harness/project.py | 7 ++- .../suites/memory_store/test_memory_store.py | 9 +-- .../session_store/test_session_store.py | 11 ++-- .../tests/suites/tools/test_deploy.py | 29 +++++++-- .../tests/suites/tools/test_dev.py | 62 +++++++++++-------- .../tests/suites/tools/test_grants.py | 7 +-- .../tests/suites/tools/test_selection.py | 7 ++- .../tests/suites/tools/test_tools_add.py | 34 ++++++++++ .../tests/suites/tools/tool_definitions.py | 56 ++++------------- .../tests/suites/tracing/test_tracing.py | 2 - 10 files changed, 129 insertions(+), 95 deletions(-) create mode 100644 integrations/agentbricks/tests/suites/tools/test_tools_add.py diff --git a/integrations/agentbricks/tests/suites/harness/project.py b/integrations/agentbricks/tests/suites/harness/project.py index fa219e46e..2d109267a 100644 --- a/integrations/agentbricks/tests/suites/harness/project.py +++ b/integrations/agentbricks/tests/suites/harness/project.py @@ -78,6 +78,10 @@ def write_file(self, relative_path: str, content: str) -> None: 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 write_manifest(self, tool_specs: Sequence[ToolSpec]) -> None: """Direct authoring: write agent.toml with exactly these tools plus a tracing binding.""" sections: list[str] = [_DIRECT_HEADER.format(framework=FRAMEWORKS[0])] @@ -90,8 +94,7 @@ def write_manifest(self, tool_specs: Sequence[ToolSpec]) -> None: log(f"# write {target}: direct authoring; no agentbricks tools command") target.write_text("\n".join(sections), encoding="utf-8") for spec in tool_specs: - for relative_path, content in spec.files.items(): - self.write_file(relative_path, content) + self.write_files(spec.files) def _pin_project_wheel(project: pathlib.Path, wheel: pathlib.Path) -> None: diff --git a/integrations/agentbricks/tests/suites/memory_store/test_memory_store.py b/integrations/agentbricks/tests/suites/memory_store/test_memory_store.py index 88d9ba651..6827cc694 100644 --- a/integrations/agentbricks/tests/suites/memory_store/test_memory_store.py +++ b/integrations/agentbricks/tests/suites/memory_store/test_memory_store.py @@ -8,21 +8,18 @@ import json import pytest -from harness.agent_endpoint import AgentEndpoint from harness.agentbricks_cli import AgentbricksCli from harness.project import Project -from harness.workspace import Workspace -def test_memory_store_created( - agentbricks_cli: AgentbricksCli, project: Project, workspace_client: Workspace -) -> None: +# Deploy needs a live workspace. +@pytest.mark.usefixtures("workspace_client") +def test_memory_store_created(agentbricks_cli: AgentbricksCli, project: Project) -> None: agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) project.pin_runtime() declared = project.declared_store("memory") assert declared, "agentbricks init declared no memory store" agentbricks_cli.run(f"deploy {project.app_name} --source {project.path}", timeout=2400) - AgentEndpoint.deployed(workspace_client, project.app_name) store = json.loads( agentbricks_cli.run(f"--output json memory stores get memory-stores/{declared}").stdout diff --git a/integrations/agentbricks/tests/suites/session_store/test_session_store.py b/integrations/agentbricks/tests/suites/session_store/test_session_store.py index a2ebde65b..c3981e232 100644 --- a/integrations/agentbricks/tests/suites/session_store/test_session_store.py +++ b/integrations/agentbricks/tests/suites/session_store/test_session_store.py @@ -9,24 +9,21 @@ import json import pytest -from harness.agent_endpoint import AgentEndpoint from harness.agentbricks_cli import AgentbricksCli from harness.project import Project -from harness.workspace import Workspace -def test_session_store_created( - agentbricks_cli: AgentbricksCli, project: Project, workspace_client: Workspace -) -> None: +# Deploy needs a live workspace. +@pytest.mark.usefixtures("workspace_client") +def test_session_store_created(agentbricks_cli: AgentbricksCli, project: Project) -> None: agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) project.pin_runtime() declared = project.declared_store("session") assert declared, "agentbricks init declared no session store" agentbricks_cli.run(f"deploy {project.app_name} --source {project.path}", timeout=2400) - AgentEndpoint.deployed(workspace_client, project.app_name) store = json.loads(agentbricks_cli.run(f"--output json sessions stores get {declared}").stdout) - assert declared in json.dumps(store), store + assert store.get("session_store_name") == declared, store @pytest.mark.xfail(reason="STUB: no session-persistence case written yet", run=False) diff --git a/integrations/agentbricks/tests/suites/tools/test_deploy.py b/integrations/agentbricks/tests/suites/tools/test_deploy.py index e8bea0fea..ff7c96751 100644 --- a/integrations/agentbricks/tests/suites/tools/test_deploy.py +++ b/integrations/agentbricks/tests/suites/tools/test_deploy.py @@ -13,7 +13,15 @@ from harness.logs import log from harness.project import Project from harness.workspace import Workspace -from tool_definitions import Tool, assert_endpoint_invoke, bind, invoke_and_check +from tool_definitions import ( + Tool, + assert_declared, + assert_replied_ok, + invocation_json, + invoke_and_check, +) + +OK_PROMPT = "Reply with the single word OK." def test_deploy( @@ -49,7 +57,10 @@ def grant_transitive_if_needed(tool: Tool) -> None: if authoring == "cli": for index, tool in enumerate(tools): with subtests.test(tool=tool.name): - bind(agentbricks_cli, project, tool) + if tool.add_command: + agentbricks_cli.run(f"tools add {tool.add_command} --source {project.path}") + project.write_files(tool.files) + assert_declared(project, tool) agentbricks_cli.run(f"deploy {app} --source {project.path}", timeout=2400) endpoint = AgentEndpoint.deployed(workspace_client, app) wait_for_grant(tool) @@ -58,7 +69,12 @@ def grant_transitive_if_needed(tool: Tool) -> None: for added in tools[: index + 1]: assert_granted(added) with subtests.test(step="endpoint-invoke"): - assert_endpoint_invoke(agentbricks_cli, endpoint) + assert_replied_ok( + agentbricks_cli.run( + f"endpoint invoke {app} --path /api/invocations --json {invocation_json(OK_PROMPT)}", + retry_unavailable=True, + ).stdout + ) return project.write_manifest(tools) @@ -71,4 +87,9 @@ def grant_transitive_if_needed(tool: Tool) -> None: invoke_and_check(endpoint, tool) assert_granted(tool) with subtests.test(step="endpoint-invoke"): - assert_endpoint_invoke(agentbricks_cli, endpoint) + assert_replied_ok( + agentbricks_cli.run( + f"endpoint invoke {app} --path /api/invocations --json {invocation_json(OK_PROMPT)}", + retry_unavailable=True, + ).stdout + ) diff --git a/integrations/agentbricks/tests/suites/tools/test_dev.py b/integrations/agentbricks/tests/suites/tools/test_dev.py index a23717900..de0406034 100644 --- a/integrations/agentbricks/tests/suites/tools/test_dev.py +++ b/integrations/agentbricks/tests/suites/tools/test_dev.py @@ -3,38 +3,29 @@ CLI authoring is incremental: add one tool, restart dev, verify it, then add the next on the same project. Direct authoring writes every tool into agent.toml and runs dev once. A subtest per tool reports each tool independently and the sequence continues past a failure. + +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 contextlib -from collections.abc import Iterator - import pytest from harness.agent_endpoint import AgentEndpoint from harness.agentbricks_cli import AgentbricksCli, free_port from harness.project import Project from tool_definitions import ( Tool, - assert_endpoint_invoke, - bind, + assert_declared, + assert_replied_ok, + invocation_json, invoke_and_check, - reject_unavailable_mcp, ) # 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") - -@contextlib.contextmanager -def dev(cli: AgentbricksCli, project: Project) -> Iterator[AgentEndpoint]: - """`agentbricks dev` the project for the duration of the block.""" - port = free_port() - # The template server binds $PORT, so an inherited PORT would override --app-port. - with cli.start( - f"dev --source {project.path} --app-port {port}", env={"PORT": str(port)} - ) as server: - yield AgentEndpoint.local(server, port) +OK_PROMPT = "Reply with the single word OK." def test_dev( @@ -48,21 +39,42 @@ def test_dev( project.pin_runtime() if authoring == "cli": - with subtests.test(step="rejects-unavailable-mcp"): - reject_unavailable_mcp(agentbricks_cli, project) for tool in tools: with subtests.test(tool=tool.name): - bind(agentbricks_cli, project, tool) - with dev(agentbricks_cli, project) as endpoint: - invoke_and_check(endpoint, tool) - with subtests.test(step="endpoint-invoke"), dev(agentbricks_cli, project) as endpoint: - assert_endpoint_invoke(agentbricks_cli, endpoint) + if tool.add_command: + agentbricks_cli.run(f"tools add {tool.add_command} --source {project.path}") + project.write_files(tool.files) + assert_declared(project, tool) + port = free_port() + with agentbricks_cli.start( + f"dev --source {project.path} --app-port {port}", env={"PORT": str(port)} + ) as server: + invoke_and_check(AgentEndpoint.local(server, port), tool) + with subtests.test(step="endpoint-invoke"): + port = free_port() + with agentbricks_cli.start( + f"dev --source {project.path} --app-port {port}", env={"PORT": str(port)} + ) as server: + endpoint = AgentEndpoint.local(server, port) + output = agentbricks_cli.run( + f"endpoint invoke --url {endpoint.base_url} --no-auth " + f"--path /api/invocations --json {invocation_json(OK_PROMPT)}" + ).stdout + assert_replied_ok(output) return project.write_manifest(tools) - with dev(agentbricks_cli, project) as endpoint: + port = free_port() + with agentbricks_cli.start( + f"dev --source {project.path} --app-port {port}", env={"PORT": str(port)} + ) as server: + endpoint = AgentEndpoint.local(server, port) for tool in tools: with subtests.test(tool=tool.name): invoke_and_check(endpoint, tool) with subtests.test(step="endpoint-invoke"): - assert_endpoint_invoke(agentbricks_cli, endpoint) + output = agentbricks_cli.run( + f"endpoint invoke --url {endpoint.base_url} --no-auth " + f"--path /api/invocations --json {invocation_json(OK_PROMPT)}" + ).stdout + assert_replied_ok(output) diff --git a/integrations/agentbricks/tests/suites/tools/test_grants.py b/integrations/agentbricks/tests/suites/tools/test_grants.py index 8bb88efd3..edc8ccae7 100644 --- a/integrations/agentbricks/tests/suites/tools/test_grants.py +++ b/integrations/agentbricks/tests/suites/tools/test_grants.py @@ -2,12 +2,11 @@ from __future__ import annotations -from harness.agent_endpoint import AgentEndpoint from harness.agentbricks_cli import AgentbricksCli from harness.project import Project from harness.run import TOOL_RESOURCE_PREFIX from harness.workspace import Workspace -from tool_definitions import bind, uc_function_tool +from tool_definitions import assert_declared, uc_function_tool from uc_objects import UcFunction @@ -25,11 +24,11 @@ def test_grants( project.pin_runtime() app = project.app_name if authoring == "cli": - bind(agentbricks_cli, project, tool) + agentbricks_cli.run(f"tools add {tool.add_command} --source {project.path}") + assert_declared(project, tool) else: project.write_manifest([tool]) agentbricks_cli.run(f"deploy {app} --source {project.path}", timeout=2400) - AgentEndpoint.deployed(workspace_client, app) with subtests.test(check="transitive-function-not-auto-granted"): # Raises if the hidden nested function is already granted or listed as an App resource. diff --git a/integrations/agentbricks/tests/suites/tools/test_selection.py b/integrations/agentbricks/tests/suites/tools/test_selection.py index 40a35c2c2..9318f832e 100644 --- a/integrations/agentbricks/tests/suites/tools/test_selection.py +++ b/integrations/agentbricks/tests/suites/tools/test_selection.py @@ -10,7 +10,7 @@ from harness.agentbricks_cli import AgentbricksCli from harness.project import Project from harness.workspace import Workspace -from tool_definitions import PYTHON_MARKER, UC_MARKER, Tool, bind, excerpt +from tool_definitions import PYTHON_MARKER, UC_MARKER, Tool, assert_declared, excerpt from uc_objects import UcFunction # Names the Unity Catalog function without its tool id, and the Python helper only to exclude it. @@ -33,7 +33,10 @@ def test_selects_uc_function_not_python_marker( app = project.app_name if authoring == "cli": for tool in tools: - bind(agentbricks_cli, project, tool) + if tool.add_command: + agentbricks_cli.run(f"tools add {tool.add_command} --source {project.path}") + project.write_files(tool.files) + assert_declared(project, tool) else: project.write_manifest(tools) agentbricks_cli.run(f"deploy {app} --source {project.path}", timeout=2400) diff --git a/integrations/agentbricks/tests/suites/tools/test_tools_add.py b/integrations/agentbricks/tests/suites/tools/test_tools_add.py new file mode 100644 index 000000000..f4708aa22 --- /dev/null +++ b/integrations/agentbricks/tests/suites/tools/test_tools_add.py @@ -0,0 +1,34 @@ +"""`agentbricks tools add` validates what it adds before it touches agent.toml.""" + +from __future__ import annotations + +import json + +import pytest +from harness.agentbricks_cli import AgentbricksCli +from harness.project import Project + + +@pytest.mark.usefixtures("workspace_client") +def test_rejects_unavailable_mcp(agentbricks_cli: AgentbricksCli, project: Project) -> None: + """An MCP service that does not exist is refused, and agent.toml is left untouched.""" + agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) + project.pin_runtime() + before = (project.path / "agent.toml").read_bytes() + + rejected = agentbricks_cli.run( + "--output json tools add mcp system.ai.missing_service --name broken_mcp " + f"--source {project.path}", + check=False, + ) + + assert rejected.returncode != 0, "tools add accepted an unavailable MCP service" + assert (project.path / "agent.toml").read_bytes() == before, ( + "a refused tools add changed agent.toml" + ) + code = json.loads(rejected.stderr).get("error", {}).get("code") + assert code in {"NOT_FOUND", "RESOURCE_DOES_NOT_EXIST"}, rejected.stderr + + agentbricks_cli.run(f"tools remove mcp system.ai.missing_service --source {project.path}") + declared = [tool["id"] for tool in project.manifest().get("tools", [])] + assert "broken_mcp" not in declared, "tools remove left the broken MCP binding in agent.toml" diff --git a/integrations/agentbricks/tests/suites/tools/tool_definitions.py b/integrations/agentbricks/tests/suites/tools/tool_definitions.py index 46253b55b..af4a702ca 100644 --- a/integrations/agentbricks/tests/suites/tools/tool_definitions.py +++ b/integrations/agentbricks/tests/suites/tools/tool_definitions.py @@ -9,7 +9,6 @@ from collections.abc import Callable, Mapping from harness.agent_endpoint import AgentEndpoint -from harness.agentbricks_cli import AgentbricksCli from harness.project import Project from harness.workspace import GrantTuple from uc_objects import UcFunction @@ -204,12 +203,8 @@ def genie_tool(genie_space: str) -> Tool: ) -def bind(cli: AgentbricksCli, project: Project, tool: Tool) -> None: - """CLI authoring: add one tool and confirm agent.toml declares it as intended.""" - if tool.add_command is not None: - cli.run(f"tools add {tool.add_command} --source {project.path}") - for relative_path, content in tool.files.items(): - project.write_file(relative_path, content) +def assert_declared(project: Project, tool: Tool) -> None: + """After CLI authoring, agent.toml declares the tool under its id, App-authed if it must be.""" if tool.tool_id is None: return declared = {t["id"]: t for t in project.manifest().get("tools", [])} @@ -222,45 +217,20 @@ def bind(cli: AgentbricksCli, project: Project, tool: Tool) -> None: ) -def reject_unavailable_mcp(cli: AgentbricksCli, project: Project) -> None: - """`tools add` must refuse an MCP service that does not exist and leave agent.toml untouched.""" - manifest = project.path / "agent.toml" - before = manifest.read_bytes() - rejected = cli.run( - "--output json tools add mcp system.ai.missing_service --name broken_mcp " - f"--source {project.path}", - check=False, - ) - assert rejected.returncode != 0 and manifest.read_bytes() == before, ( - "agentbricks tools add accepted an unavailable MCP service or changed agent.toml" - ) - code = json.loads(rejected.stderr).get("error", {}).get("code") - assert code in {"NOT_FOUND", "RESOURCE_DOES_NOT_EXIST"}, ( - f"Unexpected MCP validation error: {rejected.stderr}" - ) - cli.run(f"tools remove mcp system.ai.missing_service --source {project.path}") - assert not any(tool["id"] == "broken_mcp" for tool in project.manifest().get("tools", [])), ( - "agentbricks tools remove left the broken MCP binding in agent.toml" - ) - - def invoke_and_check(endpoint: AgentEndpoint, tool: Tool) -> None: tool.check(endpoint.invoke(tool.prompt, f"{endpoint.runtime}-{tool.name}")) -def assert_endpoint_invoke(cli: AgentbricksCli, endpoint: AgentEndpoint) -> None: - """`agentbricks endpoint invoke`, the call the CLI tells users to run after dev/deploy, gets a reply.""" +def invocation_json(prompt: str) -> str: + """A shell-quoted /api/invocations body for `agentbricks endpoint invoke --json`.""" invocation_id = str(uuid.uuid4()) - body = json.dumps( - { - "id": invocation_id, - "session_id": invocation_id, - "input": [{"role": "user", "content": "Reply with the single word OK."}], - } - ) - target = endpoint.app_name or f"--url {endpoint.base_url} --no-auth" - output = cli.run( - f"endpoint invoke {target} --path /api/invocations --json {shlex.quote(body)}", - retry_unavailable=True, - ).stdout + body = { + "id": invocation_id, + "session_id": invocation_id, + "input": [{"role": "user", "content": prompt}], + } + return shlex.quote(json.dumps(body)) + + +def assert_replied_ok(output: str) -> None: assert "ok" in output.lower(), f"endpoint invoke returned no reply: {output[:500]}" diff --git a/integrations/agentbricks/tests/suites/tracing/test_tracing.py b/integrations/agentbricks/tests/suites/tracing/test_tracing.py index 75d6f371a..b7bc8ab56 100644 --- a/integrations/agentbricks/tests/suites/tracing/test_tracing.py +++ b/integrations/agentbricks/tests/suites/tracing/test_tracing.py @@ -6,7 +6,6 @@ from __future__ import annotations import pytest -from harness.agent_endpoint import AgentEndpoint from harness.agentbricks_cli import AgentbricksCli from harness.project import Project from harness.run import TOOL_RESOURCE_PREFIX @@ -20,7 +19,6 @@ def test_tracing_experiment_bound( project.pin_runtime() experiment_name = project.manifest().get("tracing", {}).get("experiment_name") agentbricks_cli.run(f"deploy {project.app_name} --source {project.path}", timeout=2400) - AgentEndpoint.deployed(workspace_client, project.app_name) assert experiment_name, "agent.toml binds no tracing experiment" # Deploy attaches the experiment as an App resource that Agent Bricks does not own. From fcc2f79184f066da51079f81639d11b5f2ea8301 Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Wed, 14 Oct 2026 20:11:18 +0000 Subject: [PATCH 21/45] [agentbricks] E2E suites: explicit direct-authoring branch in test_dev and test_deploy Co-authored-by: Isaac --- .../tests/suites/tools/test_deploy.py | 27 ++++++--------- .../tests/suites/tools/test_dev.py | 34 +++++++++---------- 2 files changed, 27 insertions(+), 34 deletions(-) diff --git a/integrations/agentbricks/tests/suites/tools/test_deploy.py b/integrations/agentbricks/tests/suites/tools/test_deploy.py index ff7c96751..e7180512e 100644 --- a/integrations/agentbricks/tests/suites/tools/test_deploy.py +++ b/integrations/agentbricks/tests/suites/tools/test_deploy.py @@ -68,24 +68,17 @@ def grant_transitive_if_needed(tool: Tool) -> None: invoke_and_check(endpoint, tool) for added in tools[: index + 1]: assert_granted(added) - with subtests.test(step="endpoint-invoke"): - assert_replied_ok( - agentbricks_cli.run( - f"endpoint invoke {app} --path /api/invocations --json {invocation_json(OK_PROMPT)}", - retry_unavailable=True, - ).stdout - ) - return + elif authoring == "direct": + project.write_manifest(tools) + agentbricks_cli.run(f"deploy {app} --source {project.path}", timeout=2400) + endpoint = AgentEndpoint.deployed(workspace_client, app) + for tool in tools: + with subtests.test(tool=tool.name): + wait_for_grant(tool) + grant_transitive_if_needed(tool) + invoke_and_check(endpoint, tool) + assert_granted(tool) - project.write_manifest(tools) - agentbricks_cli.run(f"deploy {app} --source {project.path}", timeout=2400) - endpoint = AgentEndpoint.deployed(workspace_client, app) - for tool in tools: - with subtests.test(tool=tool.name): - wait_for_grant(tool) - grant_transitive_if_needed(tool) - invoke_and_check(endpoint, tool) - assert_granted(tool) with subtests.test(step="endpoint-invoke"): assert_replied_ok( agentbricks_cli.run( diff --git a/integrations/agentbricks/tests/suites/tools/test_dev.py b/integrations/agentbricks/tests/suites/tools/test_dev.py index de0406034..bac52166f 100644 --- a/integrations/agentbricks/tests/suites/tools/test_dev.py +++ b/integrations/agentbricks/tests/suites/tools/test_dev.py @@ -61,20 +61,20 @@ def test_dev( f"--path /api/invocations --json {invocation_json(OK_PROMPT)}" ).stdout assert_replied_ok(output) - return - - project.write_manifest(tools) - port = free_port() - with agentbricks_cli.start( - f"dev --source {project.path} --app-port {port}", env={"PORT": str(port)} - ) as server: - endpoint = AgentEndpoint.local(server, port) - for tool in tools: - with subtests.test(tool=tool.name): - invoke_and_check(endpoint, tool) - with subtests.test(step="endpoint-invoke"): - output = agentbricks_cli.run( - f"endpoint invoke --url {endpoint.base_url} --no-auth " - f"--path /api/invocations --json {invocation_json(OK_PROMPT)}" - ).stdout - assert_replied_ok(output) + elif authoring == "direct": + project.write_manifest(tools) + port = free_port() + with agentbricks_cli.start( + f"dev --source {project.path} --app-port {port}", env={"PORT": str(port)} + ) as server: + endpoint = AgentEndpoint.local(server, port) + for tool in tools: + with subtests.test(tool=tool.name): + invoke_and_check(endpoint, tool) + # Reuses the running server rather than restarting dev for this one call. + with subtests.test(step="endpoint-invoke"): + output = agentbricks_cli.run( + f"endpoint invoke --url {endpoint.base_url} --no-auth " + f"--path /api/invocations --json {invocation_json(OK_PROMPT)}" + ).stdout + assert_replied_ok(output) From eb41b70c8baf3448e23ded9cc87babaaeec0d1a1 Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Thu, 15 Oct 2026 17:30:48 +0000 Subject: [PATCH 22/45] [agentbricks] E2E suites: move harness next to the suites; WorkspaceKind enum Co-authored-by: Isaac --- integrations/agentbricks/pyproject.toml | 2 ++ .../tests/{suites => }/harness/__init__.py | 0 .../{suites => }/harness/agent_endpoint.py | 0 .../{suites => }/harness/agentbricks_cli.py | 0 .../tests/{suites => }/harness/cleanup.py | 0 .../tests/{suites => }/harness/logs.py | 0 .../tests/{suites => }/harness/project.py | 0 .../tests/{suites => }/harness/run.py | 0 .../tests/{suites => }/harness/target.py | 21 ++++++++++----- .../tests/{suites => }/harness/workspace.py | 0 .../agentbricks/tests/suites/README.md | 27 ++++++++++--------- .../agentbricks/tests/suites/conftest.py | 7 ++--- 12 files changed, 35 insertions(+), 22 deletions(-) rename integrations/agentbricks/tests/{suites => }/harness/__init__.py (100%) rename integrations/agentbricks/tests/{suites => }/harness/agent_endpoint.py (100%) rename integrations/agentbricks/tests/{suites => }/harness/agentbricks_cli.py (100%) rename integrations/agentbricks/tests/{suites => }/harness/cleanup.py (100%) rename integrations/agentbricks/tests/{suites => }/harness/logs.py (100%) rename integrations/agentbricks/tests/{suites => }/harness/project.py (100%) rename integrations/agentbricks/tests/{suites => }/harness/run.py (100%) rename integrations/agentbricks/tests/{suites => }/harness/target.py (88%) rename integrations/agentbricks/tests/{suites => }/harness/workspace.py (100%) diff --git a/integrations/agentbricks/pyproject.toml b/integrations/agentbricks/pyproject.toml index d9b7ac5ab..efc755f99 100644 --- a/integrations/agentbricks/pyproject.toml +++ b/integrations/agentbricks/pyproject.toml @@ -96,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/suites/harness/__init__.py b/integrations/agentbricks/tests/harness/__init__.py similarity index 100% rename from integrations/agentbricks/tests/suites/harness/__init__.py rename to integrations/agentbricks/tests/harness/__init__.py diff --git a/integrations/agentbricks/tests/suites/harness/agent_endpoint.py b/integrations/agentbricks/tests/harness/agent_endpoint.py similarity index 100% rename from integrations/agentbricks/tests/suites/harness/agent_endpoint.py rename to integrations/agentbricks/tests/harness/agent_endpoint.py diff --git a/integrations/agentbricks/tests/suites/harness/agentbricks_cli.py b/integrations/agentbricks/tests/harness/agentbricks_cli.py similarity index 100% rename from integrations/agentbricks/tests/suites/harness/agentbricks_cli.py rename to integrations/agentbricks/tests/harness/agentbricks_cli.py diff --git a/integrations/agentbricks/tests/suites/harness/cleanup.py b/integrations/agentbricks/tests/harness/cleanup.py similarity index 100% rename from integrations/agentbricks/tests/suites/harness/cleanup.py rename to integrations/agentbricks/tests/harness/cleanup.py diff --git a/integrations/agentbricks/tests/suites/harness/logs.py b/integrations/agentbricks/tests/harness/logs.py similarity index 100% rename from integrations/agentbricks/tests/suites/harness/logs.py rename to integrations/agentbricks/tests/harness/logs.py diff --git a/integrations/agentbricks/tests/suites/harness/project.py b/integrations/agentbricks/tests/harness/project.py similarity index 100% rename from integrations/agentbricks/tests/suites/harness/project.py rename to integrations/agentbricks/tests/harness/project.py diff --git a/integrations/agentbricks/tests/suites/harness/run.py b/integrations/agentbricks/tests/harness/run.py similarity index 100% rename from integrations/agentbricks/tests/suites/harness/run.py rename to integrations/agentbricks/tests/harness/run.py diff --git a/integrations/agentbricks/tests/suites/harness/target.py b/integrations/agentbricks/tests/harness/target.py similarity index 88% rename from integrations/agentbricks/tests/suites/harness/target.py rename to integrations/agentbricks/tests/harness/target.py index 6590cfd2f..c21b71144 100644 --- a/integrations/agentbricks/tests/suites/harness/target.py +++ b/integrations/agentbricks/tests/harness/target.py @@ -6,6 +6,7 @@ from __future__ import annotations +import enum from collections.abc import Mapping, Sequence from typing import Protocol @@ -14,7 +15,13 @@ from harness.run import RunConfig from harness.workspace import Workspace -WORKSPACE_KINDS = ("live", "fake") + +class WorkspaceKind(str, enum.Enum): + LIVE = "live" + FAKE = "fake" + + def __str__(self) -> str: + return self.value class TargetWorkspaceUnavailable(Exception): @@ -22,7 +29,7 @@ class TargetWorkspaceUnavailable(Exception): class TargetWorkspace(Protocol): - name: str + name: WorkspaceKind @property def env(self) -> Mapping[str, str]: @@ -41,7 +48,7 @@ def client(self) -> Workspace | None: class LiveWorkspace: - name = "live" + name = WorkspaceKind.LIVE def __init__(self, profile: str | None, run_config: RunConfig): self.profile = profile @@ -72,7 +79,7 @@ def client(self) -> Workspace: class FakeWorkspace: """Seam for the in-process fake: DATABRICKS_HOST/DATABRICKS_TOKEN env, no CLI args, no Workspace.""" - name = "fake" + name = WorkspaceKind.FAKE env: Mapping[str, str] = {} cli_args: Sequence[str] = () client: Workspace | None = None @@ -89,7 +96,9 @@ def require_workspace(target_workspace: TargetWorkspace, what: str) -> Workspace return workspace -def make_target_workspace(name: str, profile: str | None, run_config: RunConfig) -> TargetWorkspace: - if name == "fake": +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/suites/harness/workspace.py b/integrations/agentbricks/tests/harness/workspace.py similarity index 100% rename from integrations/agentbricks/tests/suites/harness/workspace.py rename to integrations/agentbricks/tests/harness/workspace.py diff --git a/integrations/agentbricks/tests/suites/README.md b/integrations/agentbricks/tests/suites/README.md index 9d0891fba..1a778cabb 100644 --- a/integrations/agentbricks/tests/suites/README.md +++ b/integrations/agentbricks/tests/suites/README.md @@ -46,19 +46,20 @@ Every resource a run creates carries its run id, and is deleted at three points: | Path | What it is | |---|---| -| `conftest.py` | All pytest wiring: options, hooks, and the fixtures every suite shares | -| `harness/run.py` | `RunConfig`, the run id, and the name prefixes every resource carries | -| `harness/target.py` | The `--workspace live\|fake` seam the CLI and fixtures talk to | -| `harness/agentbricks_cli.py` | Runs `agentbricks` commands: `cli.run("tools add ...")`, and `cli.start("dev ...")` in the background | -| `harness/agent_endpoint.py` | `AgentEndpoint`: a running dev server or deployed App, and invoking it over HTTP | -| `harness/project.py` | A test's project: run-scoped path and App name, its files, and `pin_runtime()` | -| `harness/workspace.py` | Databricks SDK helpers for what is not under test: SQL, UC, Apps, grants | -| `harness/cleanup.py` | Deleting what the suites create, per App and per run | -| `harness/logs.py` | Test-log output; every command is echoed so a run's log replays it | -| `/` | One folder per feature: `tools`, `memory_store`, `session_store`, `tracing` | - -`conftest.py` files are the only code pytest calls by name; everything in `harness/` is plain Python -that you can follow through its imports. +| `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 | +| `tests/harness/agent_endpoint.py` | `AgentEndpoint`: a running dev server or deployed App, and invoking it over HTTP | +| `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: SQL, UC, Apps, grants | +| `tests/harness/cleanup.py` | Deleting what the suites create, per App and per run | +| `tests/harness/logs.py` | Test-log output; every command is echoed so a run's log replays it | +| `tests/suites//` | One folder per feature: `tools`, `memory_store`, `session_store`, `tracing` | + +`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 diff --git a/integrations/agentbricks/tests/suites/conftest.py b/integrations/agentbricks/tests/suites/conftest.py index f4ca78567..56b4f6493 100644 --- a/integrations/agentbricks/tests/suites/conftest.py +++ b/integrations/agentbricks/tests/suites/conftest.py @@ -32,9 +32,9 @@ from harness.project import Project from harness.run import RUN_ID_PATTERN, E2EError, RunConfig, run_output_dir from harness.target import ( - WORKSPACE_KINDS, TargetWorkspace, TargetWorkspaceUnavailable, + WorkspaceKind, make_target_workspace, require_workspace, ) @@ -53,8 +53,9 @@ def pytest_addoption(parser: pytest.Parser) -> None: group = parser.getgroup("agentbricks", "Agent Bricks E2E suites") group.addoption( "--workspace", - choices=WORKSPACE_KINDS, - default="live", + type=WorkspaceKind, + choices=list(WorkspaceKind), + default=WorkspaceKind.LIVE, help="live runs against a Databricks workspace; fake skips whatever needs one.", ) group.addoption( From 84b1519bd19dd17e220910cba7daef34a92f3ccd Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Thu, 15 Oct 2026 18:49:12 +0000 Subject: [PATCH 23/45] [agentbricks] E2E suites: reach the agent and App resources through the CLI Co-authored-by: Isaac --- .../tests/harness/agent_endpoint.py | 166 ------------------ .../tests/harness/agentbricks_cli.py | 57 ++++-- .../agentbricks/tests/harness/cleanup.py | 75 ++++---- integrations/agentbricks/tests/harness/run.py | 1 - .../agentbricks/tests/harness/target.py | 1 - .../agentbricks/tests/harness/workspace.py | 144 +++++---------- .../agentbricks/tests/suites/README.md | 5 +- .../agentbricks/tests/suites/conftest.py | 20 +-- .../tests/suites/tools/test_deploy.py | 69 ++++---- .../tests/suites/tools/test_dev.py | 48 ++--- .../tests/suites/tools/test_grants.py | 11 +- .../tests/suites/tools/test_selection.py | 23 ++- .../tests/suites/tools/tool_definitions.py | 13 +- .../tests/suites/tracing/test_tracing.py | 11 +- 14 files changed, 219 insertions(+), 425 deletions(-) delete mode 100644 integrations/agentbricks/tests/harness/agent_endpoint.py diff --git a/integrations/agentbricks/tests/harness/agent_endpoint.py b/integrations/agentbricks/tests/harness/agent_endpoint.py deleted file mode 100644 index df935fa03..000000000 --- a/integrations/agentbricks/tests/harness/agent_endpoint.py +++ /dev/null @@ -1,166 +0,0 @@ -"""A running agent, a dev server or a deployed App, and how a test calls it over HTTP. - -``AgentEndpoint.invoke`` POSTs straight to the agent's /api/invocations, so it is the one test step -that does not go through the CLI; ``agentbricks endpoint invoke`` is checked separately. -""" - -from __future__ import annotations - -import concurrent.futures -import json -import time -import urllib.error -import urllib.request -import uuid -from collections.abc import Callable -from typing import Any - -from harness.agentbricks_cli import BackgroundCommand -from harness.logs import last_lines, last_nonempty_line, log -from harness.run import E2E_MODEL, E2EError, now -from harness.workspace import Workspace - - -class AgentEndpoint: - """Where a running agent answers, a dev server or a deployed App, and how to call it.""" - - def __init__( - self, - base_url: str, - headers: Callable[[], dict[str, str]], - *, - app_name: str | None = None, - app: dict[str, Any] | None = None, - ): - self.base_url = base_url - self.app_name = app_name - self.app = app - self._headers = headers - - @classmethod - def local(cls, server: BackgroundCommand, port: int) -> AgentEndpoint: - """The `agentbricks dev` server on ``port``, once it answers; fails if the server exits.""" - started = time.monotonic() - next_tick = 60.0 - while True: - if server.process.poll() is not None: - raise E2EError( - f"{server.label} exited {server.process.returncode}: " - f"{last_lines(server.log_path, 30)}" - ) - try: - with urllib.request.urlopen(f"http://127.0.0.1:{port}/", timeout=5): - break - except urllib.error.HTTPError as exc: - if exc.code < 500: - break - except (urllib.error.URLError, TimeoutError): - pass - elapsed = time.monotonic() - started - if elapsed > 1200: - raise E2EError( - f"{server.label} did not become reachable: {last_lines(server.log_path, 30)}" - ) - if elapsed >= next_tick: - log( - f"tick {now():%H:%M} | {server.label} | starting | " - f"{last_nonempty_line(server.log_path)}" - ) - next_tick += 60 - time.sleep(5) - return cls(f"http://127.0.0.1:{port}", lambda: {}) - - @classmethod - def deployed(cls, workspace: Workspace, app_name: str) -> AgentEndpoint: - """The App once it is up and has answered once.""" - workspace.check_app_auth() - app = workspace.wait_for_app(app_name) - url = str(app.get("url") or "").rstrip("/") - if not url: - raise E2EError(f"App {app_name} has no URL: {app}") - endpoint = cls(url, lambda: workspace.app_headers, app_name=app_name, app=app) - endpoint.warm_up() - return endpoint - - @property - def runtime(self) -> str: - return "deploy" if self.app_name else "dev" - - def invoke(self, prompt: str, label: str = "invoke") -> str: - """POST the prompt to /api/invocations and return the serialized response.""" - url = f"{self.base_url}/api/invocations" - last: Exception | None = None - invocation_id = str(uuid.uuid4()) - body = { - "id": invocation_id, - # session_id is a top-level invocation field; the adapter rejects it nested in input. - "session_id": invocation_id, - "input": { - "model": E2E_MODEL, - "messages": [{"role": "user", "content": prompt}], - }, - } - for attempt in range(1, 4): - try: - response = _monitored( - label, lambda: _http_json(url, body, self._headers()), timeout=360 - ) - return json.dumps(response, sort_keys=True, default=str) - except Exception as exc: - last = exc - log(f"attempt {attempt}/3 | {label} | {exc}") - if attempt < 3: - time.sleep(15) - raise E2EError(f"{label} failed after 3 attempts: {last}") - - def warm_up(self, deadline_seconds: float = 180.0) -> None: - """Invoke until the agent answers once. A new App's session store can drop the first calls, - and a redeployed one restarts.""" - started = time.monotonic() - while True: - try: - self.invoke("Reply with the single word OK.", label="warm-up") - return - except E2EError: - if time.monotonic() - started > deadline_seconds: - raise - - -def _monitored( - label: str, operation: Callable[[], dict[str, Any]], *, timeout: float -) -> dict[str, Any]: - with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor: - future = executor.submit(operation) - started = time.monotonic() - while True: - try: - return future.result( - timeout=min(60, max(1, timeout - (time.monotonic() - started))) - ) - except concurrent.futures.TimeoutError: - elapsed = time.monotonic() - started - log(f"tick {now():%H:%M} | {label} | running | {elapsed:.0f}s") - if elapsed >= timeout: - raise E2EError(f"{label} timed out after {timeout:.0f}s") from None - - -def _http_json(url: str, body: dict[str, Any], headers: dict[str, str]) -> dict[str, Any]: - request = urllib.request.Request( - url, - data=json.dumps(body).encode(), - headers={"Content-Type": "application/json", **headers}, - method="POST", - ) - try: - with urllib.request.urlopen(request, timeout=340) as response: - payload = response.read().decode() - except urllib.error.HTTPError as exc: - detail = exc.read().decode(errors="replace") - raise E2EError(f"HTTP {exc.code} from {url}: {detail}") from exc - try: - value = json.loads(payload) - except json.JSONDecodeError as exc: - raise E2EError(f"Invalid JSON from {url}: {payload[:2000]}") from exc - if not isinstance(value, dict): - raise E2EError(f"Expected object response from {url}, got {type(value).__name__}") - return value diff --git a/integrations/agentbricks/tests/harness/agentbricks_cli.py b/integrations/agentbricks/tests/harness/agentbricks_cli.py index a3db5ba6c..69e9ecf04 100644 --- a/integrations/agentbricks/tests/harness/agentbricks_cli.py +++ b/integrations/agentbricks/tests/harness/agentbricks_cli.py @@ -1,7 +1,8 @@ -"""The ``agentbricks`` CLI and the agents it produces: the system under test. +"""The ``agentbricks`` CLI: the system under test. -Every CLI 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. Nothing in +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. """ @@ -19,14 +20,17 @@ import subprocess import sys import time +import urllib.error +import urllib.request from collections.abc import Iterator, Mapping -from harness.logs import last_nonempty_line, log, log_command +from harness.logs import last_lines, last_nonempty_line, log, log_command from harness.run import E2EError, RunConfig, now from harness.target import TargetWorkspace -# The Apps gateway's answers while an App is not up, e.g. restarting after a redeploy. -_GATEWAY_UNAVAILABLE = re.compile(r"HTTP 50[234]\b") +# `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") @dataclasses.dataclass @@ -37,6 +41,35 @@ class BackgroundCommand: 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.""" @@ -59,12 +92,12 @@ def run( *, timeout: float = 300, check: bool = True, - retry_unavailable: bool = False, + retry_server_errors: bool = False, ) -> subprocess.CompletedProcess[str]: """Run ``agentbricks `` to completion. - ``retry_unavailable`` retries only the Apps gateway's "App not up" answers (a staging App - can restart under the test); any other failure is the CLI's. + ``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): @@ -72,12 +105,12 @@ def run( output = result.stderr or result.stdout if ( result.returncode == 0 - or not retry_unavailable + or not retry_server_errors or attempt == 3 - or _GATEWAY_UNAVAILABLE.search(output) is None + or _SERVER_ERROR.search(output) is None ): break - log(f"attempt {attempt}/3 | {command.split()[0]} | App unavailable; retrying") + 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}") diff --git a/integrations/agentbricks/tests/harness/cleanup.py b/integrations/agentbricks/tests/harness/cleanup.py index 76b000bb7..96e46de17 100644 --- a/integrations/agentbricks/tests/harness/cleanup.py +++ b/integrations/agentbricks/tests/harness/cleanup.py @@ -2,8 +2,8 @@ Two entry points, both safe to repeat (a missing resource counts as deleted): -- ``delete_app_resources``: everything one App's deploy created. The ``agentbricks_cli`` fixture - calls it for each project when its test ends. +- ``delete_app_resources``: everything one App's deploy created. The ``project`` fixture calls + it when its test ends. - ``delete_run_leftovers``: whatever a crashed test or a failed teardown left. It runs once after every live run, and on its own with ``pytest tests/suites --cleanup-only --run-id `` after a run that was killed before it could. It finds resources by the run's name prefix (Apps, source @@ -17,18 +17,16 @@ import json import pathlib import re -import subprocess from collections.abc import Callable, Collection from typing import Any +from harness.agentbricks_cli import AgentbricksCli from harness.logs import log from harness.run import APP_PREFIX, E2EError, project_prefix, scratch_schema_prefix from harness.workspace import AppIdentity, Workspace _APPS_RECORD = "apps.jsonl" -_STORE_MISSING = re.compile( - r"NOT_FOUND|RESOURCE_DOES_NOT_EXIST|does not exist|not found", re.IGNORECASE -) +_MISSING = re.compile(r"NOT_FOUND|RESOURCE_DOES_NOT_EXIST|does not exist|not found", re.IGNORECASE) # Identities recorded by per-test cleanup @@ -104,22 +102,21 @@ def delete_lakebase_leftovers( def delete_app_resources( + cli: AgentbricksCli, workspace: Workspace, name: str, *, - delete_store: Callable[[str, str], subprocess.CompletedProcess[str]] | None = None, memory_store: str | None = None, session_store: str | None = None, - has_app: bool = True, - runtime_store: bool = False, identity: AppIdentity | None = None, ) -> list[str]: - """Delete one project's stores and, if it has an App, everything the App's deploy created. + """Delete one project's stores and everything its deploy created, as far as any of it exists. - ``delete_store(kind, store)`` is the agentbricks CLI's store delete, which has no SDK surface. - ``identity`` is what was recorded earlier; it fills in whatever the live lookup misses. - Every step runs regardless of earlier failures except those that need the App gone. Returns a - description of each failure; an empty list means everything was deleted. + The App, its Runtime Store and the stores go through the CLI. The rest are workarounds for what + no public command deletes yet: the synced source folder, and the Lakebase databases and roles a + failed deploy leaves. ``identity`` is what was recorded earlier; it fills in whatever the live + lookup misses. Every step runs regardless of earlier failures. Returns a description of each + failure; an empty list means everything was deleted. """ failures: list[str] = [] @@ -131,29 +128,31 @@ def attempt(label: str, action: Callable[..., object], *args: Any) -> bool: return False return True - def delete_cli_store(kind: str, store: str) -> None: - assert delete_store is not None - result = delete_store(kind, store) + def run(command: str) -> None: + result = cli.run(command, timeout=1200, check=False) output = (result.stderr or result.stdout).strip() # A store the test declared but never deployed was never created. - if result.returncode != 0 and not _STORE_MISSING.search(output): + if result.returncode != 0 and not _MISSING.search(output): raise E2EError(output) # Read before any delete: the principal and branch come from the App and its Runtime Store. - known = (identity or AppIdentity()).merged( - workspace.app_identity(name) if has_app else AppIdentity() - ) - for kind, store in (("memory", memory_store), ("sessions", session_store)): - if store and delete_store: - attempt(f"{kind} store {store}", delete_cli_store, kind, store) - if not has_app: + known = (identity or AppIdentity()).merged(workspace.app_identity(name)) + if memory_store: + attempt(f"memory store {memory_store}", run, f"memory stores delete {memory_store} --yes") + if session_store: + attempt( + f"session store {session_store}", run, f"sessions stores delete {session_store} --yes" + ) + if known == AppIdentity(): return failures attempt(f"catalog grant {name}", workspace.revoke_catalog_grant, name) - if runtime_store: + if known.service_principal: + # Deletes the Runtime Store first, while the App's principal still resolves, then the App. + attempt(f"deployment {name}", run, f"deployments delete {name} --yes") + else: attempt(f"runtime store {name}", workspace.delete_runtime_store, name) - app_deleted = attempt(f"app {name}", workspace.delete_app, name) - if app_deleted: + if attempt(f"app {name}", workspace.wait_until_app_gone, name): attempt(f"deployment source {name}", workspace.delete_deployment_source, name) if known.branch: failures += delete_lakebase_leftovers( @@ -169,12 +168,12 @@ def delete_cli_store(kind: str, store: str) -> None: def delete_run_leftovers( - workspace: Workspace, run_id: str, output: pathlib.Path, catalog: str + 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 Apps", lambda: _delete_leaked_apps(workspace, run_id, output)), + ("leaked Apps", lambda: _delete_leaked_apps(cli, workspace, run_id, output)), ("leaked Lakebase databases and roles", lambda: _delete_leaked_lakebase(workspace, output)), ("leaked source folders", lambda: _delete_leaked_deployment_sources(workspace, run_id)), ("leaked schemas", lambda: _delete_leaked_schemas(workspace, catalog, run_id)), @@ -191,17 +190,19 @@ def _app_prefix(run_id: str) -> str: return f"{APP_PREFIX}{project_prefix(run_id)}" -def _delete_leaked_apps(workspace: Workspace, run_id: str, output: pathlib.Path) -> list[str]: +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, found by the run's name prefix.""" prefix = _app_prefix(run_id) recorded = recorded_app_identities(output) + deployments = json.loads(cli.run("--output json deployments list").stdout or "[]") failures: list[str] = [] - for app in list(workspace.client.apps.list()): - if app.name and app.name.startswith(prefix): - log(f"deleting leaked App {app.name}") - failures += delete_app_resources( - workspace, app.name, runtime_store=True, identity=recorded.get(app.name) - ) + for deployment in deployments: + name = str(deployment.get("name") or "") + if name.startswith(prefix): + log(f"deleting leaked App {name}") + failures += delete_app_resources(cli, workspace, name, identity=recorded.get(name)) return failures diff --git a/integrations/agentbricks/tests/harness/run.py b/integrations/agentbricks/tests/harness/run.py index c73b5aa00..725423f7c 100644 --- a/integrations/agentbricks/tests/harness/run.py +++ b/integrations/agentbricks/tests/harness/run.py @@ -52,7 +52,6 @@ class RunConfig: run_id: str output: pathlib.Path - app_auth_profile: str | None # Operator override; None lets the workspace pick the first running warehouse. warehouse_id: str | None preprovisioned_app_catalog_access: bool diff --git a/integrations/agentbricks/tests/harness/target.py b/integrations/agentbricks/tests/harness/target.py index c21b71144..ea892082c 100644 --- a/integrations/agentbricks/tests/harness/target.py +++ b/integrations/agentbricks/tests/harness/target.py @@ -69,7 +69,6 @@ def client(self) -> Workspace: if self._client is None: self._client = Workspace( self.profile, - app_auth_profile=self._run_config.app_auth_profile, warehouse_id=self._run_config.warehouse_id, preprovisioned_app_catalog_access=self._run_config.preprovisioned_app_catalog_access, ) diff --git a/integrations/agentbricks/tests/harness/workspace.py b/integrations/agentbricks/tests/harness/workspace.py index bc72cf1b1..dd4792a06 100644 --- a/integrations/agentbricks/tests/harness/workspace.py +++ b/integrations/agentbricks/tests/harness/workspace.py @@ -1,8 +1,10 @@ -"""Databricks SDK helpers for the E2E suites and fixtures: the workspace, which is not under test. +"""Databricks SDK helpers for what is not under test: the prerequisites the suites need, and checks +of results the CLI has no command to show. -``Workspace`` holds one operation per method (create, read, delete one thing); deciding what to -delete and in which order is ``harness.cleanup``. Everything goes through the SDK except -``app_logs``, since the SDK has no Apps log API. +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, UC privilege checks, 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 @@ -53,55 +55,18 @@ def __init__( self, profile: str | None, *, - app_auth_profile: str | None = None, warehouse_id: str | None = None, preprovisioned_app_catalog_access: bool = False, ): self.profile = profile - self.app_auth_profile = app_auth_profile or profile self.warehouse_id = warehouse_id self.preprovisioned_app_catalog_access = preprovisioned_app_catalog_access self.client = WorkspaceClient(profile=profile) - self._app_auth_client: WorkspaceClient | None = None - self._app_auth_checked = False self._warehouse_started = False self._user_name: str | None = None # App name -> (catalog, principal) of each USE CATALOG this suite granted, for cleanup. self._catalog_grants: dict[str, tuple[str, str]] = {} - # App auth: Databricks Apps /api routes need OAuth, so a PAT is rejected. - - def _app_auth_label(self) -> str: - return ( - f"App auth profile {self.app_auth_profile!r}" - if self.app_auth_profile - else "ambient environment credentials" - ) - - def _app_client(self) -> WorkspaceClient: - if self._app_auth_client is None: - self._app_auth_client = WorkspaceClient(profile=self.app_auth_profile) - return self._app_auth_client - - def check_app_auth(self) -> None: - if self._app_auth_checked: - return - if self._app_client().config.auth_type == "pat": - raise E2EError( - f"{self._app_auth_label()} uses a PAT. " - "Databricks Apps /api routes require OAuth; run `databricks auth login` " - "for a profile on the same workspace." - ) - self._app_auth_checked = True - - @property - def app_headers(self) -> dict[str, str]: - # Resolved per call so OAuth tokens refresh during a long run. - authorization = self._app_client().config.authenticate().get("Authorization") - if not authorization: - raise E2EError(f"Could not resolve credentials from {self._app_auth_label()}.") - return {"Authorization": authorization} - # Existing resources the suite depends on def require_catalog(self, name: str) -> None: @@ -215,23 +180,6 @@ def assert_app_absent(self, name: str) -> None: raise E2EError(f"Could not verify App {name} is absent: {exc}") from exc raise E2EError(f"App {name} already exists; refusing to deploy over it.") - def wait_for_app(self, name: str) -> dict[str, Any]: - """The App once it is ACTIVE with a URL.""" - started = time.monotonic() - next_tick = 0.0 - while time.monotonic() - started < 1200: - app = self.app(name) - compute = app.get("compute_status", {}) - state = compute.get("state") if isinstance(compute, dict) else None - if state == "ACTIVE" and app.get("url"): - return app - elapsed = time.monotonic() - started - if elapsed >= next_tick: - log(f"tick {now():%H:%M} | app-{name} | {state or 'UNKNOWN'}") - next_tick += 60 - time.sleep(15) - raise E2EError(f"App {name} did not become ACTIVE.") - 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"] @@ -253,20 +201,19 @@ def app_logs(self, name: str, log_path: pathlib.Path) -> pathlib.Path | None: # Grants - def granted(self, app_name: str, grant: GrantTuple) -> bool: - """Whether Agent Bricks granted the App ``grant``, as (kind, name, type, permission).""" - return grant in app_resource_tuples(tool_resources(self.app(app_name))) - - def wait_until_effective(self, app_name: str, grant: GrantTuple, timeout: float = 120) -> float: + def wait_until_effective( + self, deployment: dict[str, Any], grant: GrantTuple, timeout: float = 120 + ) -> float: """Wait until a UC ``grant`` deploy added is effective for the App; returns the seconds waited. - Deploy reconciles the App's UC resources before returning, but the privilege can become - effective for its principal a little later; a tool called in between is denied. + ``deployment`` is ``agentbricks deployments get --output json``. Deploy reconciles the App's + UC resources before returning, but the privilege can become effective for its principal a + little later; a tool called in between is denied. """ kind, full_name, securable_type, privilege = grant if kind != "uc_securable": return 0.0 - principal = self._principal(self.app(app_name)) + principal = _principal(deployment) started = time.monotonic() while True: held = effective_privileges( @@ -276,26 +223,20 @@ def wait_until_effective(self, app_name: str, grant: GrantTuple, timeout: float return time.monotonic() - started if time.monotonic() - started > timeout: raise E2EError( - f"{privilege} on {full_name} not effective for {app_name} after {timeout:.0f}s: {held}" + f"{privilege} on {full_name} not effective for {deployment.get('name')} " + f"after {timeout:.0f}s: {held}" ) time.sleep(5) - def granted_tuples(self, app_name: str) -> set[GrantTuple]: - return app_resource_tuples(tool_resources(self.app(app_name))) - def experiment_id(self, name: str) -> str: experiment = self.client.experiments.get_by_name(name).experiment if experiment is None or not experiment.experiment_id: raise E2EError(f"Experiment {name!r} does not exist.") return str(experiment.experiment_id) - def non_tool_resources(self, app_name: str) -> list[dict[str, Any]]: - return non_tool_resources(self.app(app_name)) - - def transitive_state(self, app_name: str, function: str) -> dict[str, Any]: + def transitive_state(self, deployment: dict[str, Any], function: str) -> dict[str, Any]: """The App principal's privileges on a nested function before any manual grant.""" - app = self.app(app_name) - principal = self._principal(app) + principal = _principal(deployment) direct = direct_privileges(self.client, SecurableType.FUNCTION, function, principal) if "EXECUTE" in direct: raise E2EError( @@ -309,7 +250,7 @@ def transitive_state(self, app_name: str, function: str) -> dict[str, Any]: ) if any( resource.get("uc_securable", {}).get("securable_full_name") == function - for resource in tool_resources(app) + for resource in tool_resources(deployment) ): raise E2EError("The transitive function appeared in Agent Bricks-owned Apps resources.") return { @@ -330,9 +271,10 @@ def revoke_catalog_grant(self, app_name: str) -> None: ) del self._catalog_grants[app_name] - def grant_transitive(self, app_name: str, function: str) -> dict[str, Any]: + def grant_transitive(self, deployment: dict[str, Any], function: str) -> dict[str, Any]: """Grant the App principal EXECUTE on a nested function and confirm it is direct and effective.""" - principal = self._principal(self.app(app_name)) + app_name = str(deployment.get("name")) + principal = _principal(deployment) catalog, schema, function_name = function.split(".") quoted_principal = f"`{principal.replace('`', '``')}`" statements = [] @@ -368,24 +310,11 @@ def grant_transitive(self, app_name: str, function: str) -> dict[str, Any]: "effective_privileges": effective, } - @staticmethod - def _principal(app: dict[str, Any]) -> str: - principal = app.get("service_principal_client_id") - if not principal: - raise E2EError(f"App response has no service_principal_client_id: {app}") - return str(principal) - - # Cleanup - - def delete_app(self, name: str) -> None: - """Delete the App; a missing App counts as deleted.""" - try: - self.client.apps.delete(name) - except NotFound: - return - self._wait_for_app_deleted(name) + # Workarounds: cleanup that no public command does yet. `agentbricks deployments delete` + # removes the App and its Runtime Store; these cover what it leaves behind. - def _wait_for_app_deleted(self, name: str, timeout: float = 1200) -> None: + def wait_until_app_gone(self, name: str, timeout: float = 1200) -> None: + """Block until the App is gone; `apps delete` returns while it is still deleting.""" started = time.monotonic() next_tick = 0.0 while True: @@ -411,7 +340,10 @@ def runtime_store(self, app_name: str) -> dict[str, Any]: ) def delete_runtime_store(self, app_name: str) -> None: - """Delete the deploy-created Runtime Store; a missing store counts as deleted.""" + """Delete a Runtime Store whose App is gone, which `deployments delete` cannot reach. + + A missing store counts as deleted. + """ try: self.client.api_client.do("DELETE", f"/api/2.0/agents/runtime-stores/{app_name}") except NotFound: @@ -535,7 +467,23 @@ def _wait(operation: Any) -> None: wait() -# Grant helpers +# Reading `agentbricks deployments get --output json` + + +def _principal(deployment: dict[str, Any]) -> str: + principal = deployment.get("service_principal_client_id") + if not principal: + raise E2EError(f"Deployment has no service_principal_client_id: {deployment}") + return str(principal) + + +def granted(deployment: dict[str, Any], grant: GrantTuple) -> bool: + """Whether deploy attached ``grant``, as (kind, name, type, permission), to the App.""" + return grant in granted_tuples(deployment) + + +def granted_tuples(deployment: dict[str, Any]) -> set[GrantTuple]: + return app_resource_tuples(tool_resources(deployment)) def _split_resources(app: dict[str, Any]) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: diff --git a/integrations/agentbricks/tests/suites/README.md b/integrations/agentbricks/tests/suites/README.md index 1a778cabb..a0e77f666 100644 --- a/integrations/agentbricks/tests/suites/README.md +++ b/integrations/agentbricks/tests/suites/README.md @@ -50,10 +50,9 @@ Every resource a run creates carries its run id, and is deleted at three points: | `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 | -| `tests/harness/agent_endpoint.py` | `AgentEndpoint`: a running dev server or deployed App, and invoking it over HTTP | | `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: SQL, UC, Apps, grants | -| `tests/harness/cleanup.py` | Deleting what the suites create, per App and per run | +| `tests/harness/workspace.py` | Databricks SDK helpers for what is not under test: prerequisites (UC objects, warehouse, Genie), UC privilege checks, and cleanup no CLI command does yet | +| `tests/harness/cleanup.py` | Deleting what the suites create, per App and per run; the App and its Runtime Store go through `agentbricks deployments delete` | | `tests/harness/logs.py` | Test-log output; every command is echoed so a run's log replays it | | `tests/suites//` | One folder per feature: `tools`, `memory_store`, `session_store`, `tracing` | diff --git a/integrations/agentbricks/tests/suites/conftest.py b/integrations/agentbricks/tests/suites/conftest.py index 56b4f6493..a913d73e9 100644 --- a/integrations/agentbricks/tests/suites/conftest.py +++ b/integrations/agentbricks/tests/suites/conftest.py @@ -62,10 +62,6 @@ def pytest_addoption(parser: pytest.Parser) -> None: "--databricks-profile", help="Databricks CLI profile; default is ambient credentials (DATABRICKS_CONFIG_PROFILE).", ) - group.addoption( - "--app-auth-profile", - help="OAuth profile for deployed App /api calls; defaults to --databricks-profile.", - ) group.addoption( "--catalog", default="main", @@ -151,7 +147,6 @@ def _base_run_config(config: pytest.Config) -> RunConfig: run = RunConfig( run_id=run_id, output=output, - app_auth_profile=option("app_auth_profile") or option("databricks_profile"), warehouse_id=option("warehouse_id"), preprovisioned_app_catalog_access=bool(option("preprovisioned_app_catalog_access")), bridge_sha=option("bridge_sha"), @@ -278,16 +273,11 @@ def _delete_project( memory_store = project.declared_store("memory") try: failures = delete_app_resources( + cli, workspace, name, - delete_store=lambda kind, store: cli.run( - f"{kind} stores delete {store} --yes", timeout=600, check=False - ), memory_store=f"memory-stores/{memory_store}" if memory_store else None, session_store=project.declared_store("session"), - has_app=deployed, - # Every managed deploy creates a Runtime Store, with or without memory/session stores. - runtime_store=deployed, identity=identity, ) except Exception as exc: @@ -317,7 +307,8 @@ def pytest_sessionfinish(session: pytest.Session) -> None: log("Resources retained (--keep-resources); skipping leftover cleanup.") return try: - workspace = _build_target_workspace(config).client + target_workspace = _build_target_workspace(config) + workspace = target_workspace.client except TargetWorkspaceUnavailable: workspace = None if workspace is None: @@ -326,7 +317,10 @@ def pytest_sessionfinish(session: pytest.Session) -> None: 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") - failures = delete_run_leftovers(workspace, run.run_id, run.output, config.getoption("catalog")) + 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: diff --git a/integrations/agentbricks/tests/suites/tools/test_deploy.py b/integrations/agentbricks/tests/suites/tools/test_deploy.py index e7180512e..860af2d51 100644 --- a/integrations/agentbricks/tests/suites/tools/test_deploy.py +++ b/integrations/agentbricks/tests/suites/tools/test_deploy.py @@ -4,24 +4,20 @@ After every redeploy the grants of all tools added so far must still hold, which is the idempotency check. Direct authoring writes every tool into agent.toml and deploys once. A subtest per tool reports each tool independently and the sequence continues past a failure. + +`agentbricks deploy` returns once the App is running its new source, so the next step can use it. """ from __future__ import annotations -from harness.agent_endpoint import AgentEndpoint +import json +from typing import Any + from harness.agentbricks_cli import AgentbricksCli from harness.logs import log from harness.project import Project -from harness.workspace import Workspace -from tool_definitions import ( - Tool, - assert_declared, - assert_replied_ok, - invocation_json, - invoke_and_check, -) - -OK_PROMPT = "Reply with the single word OK." +from harness.workspace import Workspace, granted, granted_tuples +from tool_definitions import Tool, assert_declared, invocation_json def test_deploy( @@ -36,23 +32,28 @@ def test_deploy( project.pin_runtime() app = project.app_name - def assert_granted(tool: Tool) -> None: + def assert_granted(tool: Tool, deployment: dict[str, Any]) -> None: if tool.grant is not None: - assert workspace_client.granted(app, tool.grant), ( + assert granted(deployment, tool.grant), ( f"{tool.name}: {tool.grant} not granted to {app}: " - f"{sorted(workspace_client.granted_tuples(app))}" + f"{sorted(granted_tuples(deployment))}" ) - def wait_for_grant(tool: Tool) -> None: + def prepare_and_check(tool: Tool, deployment: dict[str, Any]) -> None: if tool.grant is not None: - waited = workspace_client.wait_until_effective(app, tool.grant) + waited = workspace_client.wait_until_effective(deployment, tool.grant) log(f"{tool.name}: {tool.grant[3]} effective for {app} after {waited:.0f}s") - - def grant_transitive_if_needed(tool: Tool) -> None: - # The nested function is deliberately outside Agent Bricks' grants; the tool only works - # once this manual grant is in place. if tool.transitive_function: - workspace_client.grant_transitive(app, tool.transitive_function) + # The nested function is deliberately outside Agent Bricks' grants; the tool only + # works once this manual grant is in place. + workspace_client.grant_transitive(deployment, tool.transitive_function) + output = agentbricks_cli.run( + f"--output json endpoint invoke {app} --path /api/invocations " + f"--json {invocation_json(tool.prompt)}", + timeout=360, + retry_server_errors=True, + ).stdout + tool.check(output) if authoring == "cli": for index, tool in enumerate(tools): @@ -62,27 +63,17 @@ def grant_transitive_if_needed(tool: Tool) -> None: project.write_files(tool.files) assert_declared(project, tool) agentbricks_cli.run(f"deploy {app} --source {project.path}", timeout=2400) - endpoint = AgentEndpoint.deployed(workspace_client, app) - wait_for_grant(tool) - grant_transitive_if_needed(tool) - invoke_and_check(endpoint, tool) + deployment = json.loads( + agentbricks_cli.run(f"--output json deployments get {app}").stdout + ) + prepare_and_check(tool, deployment) for added in tools[: index + 1]: - assert_granted(added) + assert_granted(added, deployment) elif authoring == "direct": project.write_manifest(tools) agentbricks_cli.run(f"deploy {app} --source {project.path}", timeout=2400) - endpoint = AgentEndpoint.deployed(workspace_client, app) + deployment = json.loads(agentbricks_cli.run(f"--output json deployments get {app}").stdout) for tool in tools: with subtests.test(tool=tool.name): - wait_for_grant(tool) - grant_transitive_if_needed(tool) - invoke_and_check(endpoint, tool) - assert_granted(tool) - - with subtests.test(step="endpoint-invoke"): - assert_replied_ok( - agentbricks_cli.run( - f"endpoint invoke {app} --path /api/invocations --json {invocation_json(OK_PROMPT)}", - retry_unavailable=True, - ).stdout - ) + prepare_and_check(tool, deployment) + assert_granted(tool, deployment) diff --git a/integrations/agentbricks/tests/suites/tools/test_dev.py b/integrations/agentbricks/tests/suites/tools/test_dev.py index bac52166f..38e9679fc 100644 --- a/integrations/agentbricks/tests/suites/tools/test_dev.py +++ b/integrations/agentbricks/tests/suites/tools/test_dev.py @@ -2,7 +2,8 @@ CLI authoring is incremental: add one tool, restart dev, verify it, then add the next on the same project. Direct authoring writes every tool into agent.toml and runs dev once. A subtest per tool -reports each tool independently and the sequence continues past a failure. +reports each tool independently and the sequence continues past a failure. Every check calls the +dev server with `agentbricks endpoint invoke`, the command the CLI tells users to run. 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. @@ -11,22 +12,13 @@ from __future__ import annotations import pytest -from harness.agent_endpoint import AgentEndpoint from harness.agentbricks_cli import AgentbricksCli, free_port from harness.project import Project -from tool_definitions import ( - Tool, - assert_declared, - assert_replied_ok, - invocation_json, - invoke_and_check, -) +from tool_definitions import Tool, assert_declared, invocation_json # 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") -OK_PROMPT = "Reply with the single word OK." - def test_dev( agentbricks_cli: AgentbricksCli, @@ -38,6 +30,15 @@ def test_dev( agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) project.pin_runtime() + def check(tool: Tool, port: int) -> None: + output = agentbricks_cli.run( + f"--output json endpoint invoke --url http://127.0.0.1:{port} --no-auth " + f"--path /api/invocations --json {invocation_json(tool.prompt)}", + timeout=360, + retry_server_errors=True, + ).stdout + tool.check(output) + if authoring == "cli": for tool in tools: with subtests.test(tool=tool.name): @@ -49,32 +50,15 @@ def test_dev( with agentbricks_cli.start( f"dev --source {project.path} --app-port {port}", env={"PORT": str(port)} ) as server: - invoke_and_check(AgentEndpoint.local(server, port), tool) - with subtests.test(step="endpoint-invoke"): - port = free_port() - with agentbricks_cli.start( - f"dev --source {project.path} --app-port {port}", env={"PORT": str(port)} - ) as server: - endpoint = AgentEndpoint.local(server, port) - output = agentbricks_cli.run( - f"endpoint invoke --url {endpoint.base_url} --no-auth " - f"--path /api/invocations --json {invocation_json(OK_PROMPT)}" - ).stdout - assert_replied_ok(output) + server.wait_until_serving(port) + check(tool, port) elif authoring == "direct": project.write_manifest(tools) port = free_port() with agentbricks_cli.start( f"dev --source {project.path} --app-port {port}", env={"PORT": str(port)} ) as server: - endpoint = AgentEndpoint.local(server, port) + server.wait_until_serving(port) for tool in tools: with subtests.test(tool=tool.name): - invoke_and_check(endpoint, tool) - # Reuses the running server rather than restarting dev for this one call. - with subtests.test(step="endpoint-invoke"): - output = agentbricks_cli.run( - f"endpoint invoke --url {endpoint.base_url} --no-auth " - f"--path /api/invocations --json {invocation_json(OK_PROMPT)}" - ).stdout - assert_replied_ok(output) + check(tool, port) diff --git a/integrations/agentbricks/tests/suites/tools/test_grants.py b/integrations/agentbricks/tests/suites/tools/test_grants.py index edc8ccae7..ddfc2f6ed 100644 --- a/integrations/agentbricks/tests/suites/tools/test_grants.py +++ b/integrations/agentbricks/tests/suites/tools/test_grants.py @@ -2,10 +2,12 @@ from __future__ import annotations +import json + from harness.agentbricks_cli import AgentbricksCli from harness.project import Project from harness.run import TOOL_RESOURCE_PREFIX -from harness.workspace import Workspace +from harness.workspace import Workspace, non_tool_resources from tool_definitions import assert_declared, uc_function_tool from uc_objects import UcFunction @@ -29,17 +31,18 @@ def test_grants( else: project.write_manifest([tool]) agentbricks_cli.run(f"deploy {app} --source {project.path}", timeout=2400) + deployment = json.loads(agentbricks_cli.run(f"--output json deployments get {app}").stdout) with subtests.test(check="transitive-function-not-auto-granted"): # Raises if the hidden nested function is already granted or listed as an App resource. - workspace_client.transitive_state(app, uc_function.nested_function) + workspace_client.transitive_state(deployment, uc_function.nested_function) with subtests.test(check="manual-grant-is-direct-and-effective"): - granted = workspace_client.grant_transitive(app, uc_function.nested_function) + granted = workspace_client.grant_transitive(deployment, uc_function.nested_function) assert "EXECUTE" in granted["direct_privileges"] assert "EXECUTE" in granted["effective_privileges"] with subtests.test(check="non-tool-resource-preserved"): - unrelated = workspace_client.non_tool_resources(app) + unrelated = non_tool_resources(deployment) assert unrelated, "No non-tool App resource; the tracing experiment should be attached" assert not any(r["name"].startswith(TOOL_RESOURCE_PREFIX) for r in unrelated) diff --git a/integrations/agentbricks/tests/suites/tools/test_selection.py b/integrations/agentbricks/tests/suites/tools/test_selection.py index 9318f832e..9533a049b 100644 --- a/integrations/agentbricks/tests/suites/tools/test_selection.py +++ b/integrations/agentbricks/tests/suites/tools/test_selection.py @@ -6,11 +6,19 @@ from __future__ import annotations -from harness.agent_endpoint import AgentEndpoint +import json + from harness.agentbricks_cli import AgentbricksCli from harness.project import Project from harness.workspace import Workspace -from tool_definitions import PYTHON_MARKER, UC_MARKER, Tool, assert_declared, excerpt +from tool_definitions import ( + PYTHON_MARKER, + UC_MARKER, + Tool, + assert_declared, + excerpt, + invocation_json, +) from uc_objects import UcFunction # Names the Unity Catalog function without its tool id, and the Python helper only to exclude it. @@ -40,10 +48,15 @@ def test_selects_uc_function_not_python_marker( else: project.write_manifest(tools) agentbricks_cli.run(f"deploy {app} --source {project.path}", timeout=2400) - endpoint = AgentEndpoint.deployed(workspace_client, app) + deployment = json.loads(agentbricks_cli.run(f"--output json deployments get {app}").stdout) # The nested function is deliberately outside Agent Bricks' grants. - workspace_client.grant_transitive(app, uc_function.nested_function) - serialized = endpoint.invoke(PROMPT, "deploy-selection") + workspace_client.grant_transitive(deployment, uc_function.nested_function) + serialized = agentbricks_cli.run( + f"--output json endpoint invoke {app} --path /api/invocations " + f"--json {invocation_json(PROMPT)}", + timeout=360, + retry_server_errors=True, + ).stdout assert UC_MARKER in serialized, f"The UC function tool was not called: {excerpt(serialized)}" assert PYTHON_MARKER not in serialized, ( diff --git a/integrations/agentbricks/tests/suites/tools/tool_definitions.py b/integrations/agentbricks/tests/suites/tools/tool_definitions.py index af4a702ca..64398ccb9 100644 --- a/integrations/agentbricks/tests/suites/tools/tool_definitions.py +++ b/integrations/agentbricks/tests/suites/tools/tool_definitions.py @@ -8,8 +8,8 @@ import uuid from collections.abc import Callable, Mapping -from harness.agent_endpoint import AgentEndpoint from harness.project import Project +from harness.run import E2E_MODEL from harness.workspace import GrantTuple from uc_objects import UcFunction @@ -217,20 +217,13 @@ def assert_declared(project: Project, tool: Tool) -> None: ) -def invoke_and_check(endpoint: AgentEndpoint, tool: Tool) -> None: - tool.check(endpoint.invoke(tool.prompt, f"{endpoint.runtime}-{tool.name}")) - - def invocation_json(prompt: str) -> str: """A shell-quoted /api/invocations body for `agentbricks endpoint invoke --json`.""" invocation_id = str(uuid.uuid4()) body = { "id": invocation_id, + # session_id is a top-level invocation field; the adapter rejects it nested in input. "session_id": invocation_id, - "input": [{"role": "user", "content": prompt}], + "input": {"model": E2E_MODEL, "messages": [{"role": "user", "content": prompt}]}, } return shlex.quote(json.dumps(body)) - - -def assert_replied_ok(output: str) -> None: - assert "ok" in output.lower(), f"endpoint invoke returned no reply: {output[:500]}" diff --git a/integrations/agentbricks/tests/suites/tracing/test_tracing.py b/integrations/agentbricks/tests/suites/tracing/test_tracing.py index b7bc8ab56..55f707c1b 100644 --- a/integrations/agentbricks/tests/suites/tracing/test_tracing.py +++ b/integrations/agentbricks/tests/suites/tracing/test_tracing.py @@ -5,11 +5,13 @@ from __future__ import annotations +import json + import pytest from harness.agentbricks_cli import AgentbricksCli from harness.project import Project from harness.run import TOOL_RESOURCE_PREFIX -from harness.workspace import Workspace +from harness.workspace import Workspace, non_tool_resources def test_tracing_experiment_bound( @@ -20,11 +22,12 @@ def test_tracing_experiment_bound( experiment_name = project.manifest().get("tracing", {}).get("experiment_name") agentbricks_cli.run(f"deploy {project.app_name} --source {project.path}", timeout=2400) assert experiment_name, "agent.toml binds no tracing experiment" + deployment = json.loads( + agentbricks_cli.run(f"--output json deployments get {project.app_name}").stdout + ) # Deploy attaches the experiment as an App resource that Agent Bricks does not own. - experiments = [ - r for r in workspace_client.non_tool_resources(project.app_name) if "experiment" in r - ] + experiments = [r for r in non_tool_resources(deployment) if "experiment" in r] assert experiments, "No experiment App resource among the non-tool resources" assert not any(r["name"].startswith(TOOL_RESOURCE_PREFIX) for r in experiments) expected_id = workspace_client.experiment_id(experiment_name) From 715306091a462960d9b6f560eb3ce7aaf0d660db Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Thu, 15 Oct 2026 20:01:16 +0000 Subject: [PATCH 24/45] [agentbricks] E2E suites: per-test cleanup through the CLI only; the sweep deletes the rest Co-authored-by: Isaac --- .../agentbricks/tests/harness/cleanup.py | 104 ++++++++++-------- .../agentbricks/tests/harness/workspace.py | 17 +-- .../agentbricks/tests/suites/README.md | 16 +-- .../agentbricks/tests/suites/conftest.py | 4 +- 4 files changed, 73 insertions(+), 68 deletions(-) diff --git a/integrations/agentbricks/tests/harness/cleanup.py b/integrations/agentbricks/tests/harness/cleanup.py index 96e46de17..e7bfbbb3c 100644 --- a/integrations/agentbricks/tests/harness/cleanup.py +++ b/integrations/agentbricks/tests/harness/cleanup.py @@ -2,13 +2,15 @@ Two entry points, both safe to repeat (a missing resource counts as deleted): -- ``delete_app_resources``: everything one App's deploy created. The ``project`` fixture calls - it when its test ends. -- ``delete_run_leftovers``: whatever a crashed test or a failed teardown left. It runs once after - every live run, and on its own with ``pytest tests/suites --cleanup-only --run-id `` after a - run that was killed before it could. It finds resources by the run's name prefix (Apps, source - folders, schemas) and by the App identities per-test cleanup recorded in ``/apps.jsonl`` - (Lakebase databases and roles, which outlive their App), so it touches nothing outside the run. +- ``delete_app_resources``: one App and its stores, deleted the way a user would, through the CLI. + The ``project`` fixture calls it when its test ends. +- ``delete_run_leftovers``: everything else, including whatever the CLI leaves behind (synced + source folders, Runtime Stores whose App is gone, Lakebase databases and roles) and whatever a + crashed test or a failed teardown left. It runs once after every live run, and on its own with + ``pytest tests/suites --cleanup-only --run-id `` after a run that was killed before it could. + It finds resources by the run's name prefix (Apps, source folders, schemas) and by the App + identities per-test cleanup recorded in ``/apps.jsonl`` (Runtime Stores, Lakebase + databases and roles, which outlive their App), so it touches nothing outside the run. """ from __future__ import annotations @@ -20,6 +22,8 @@ from collections.abc import Callable, Collection from typing import Any +from databricks.sdk.errors import NotFound + from harness.agentbricks_cli import AgentbricksCli from harness.logs import log from harness.run import APP_PREFIX, E2EError, project_prefix, scratch_schema_prefix @@ -35,6 +39,8 @@ 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) @@ -105,28 +111,24 @@ def delete_app_resources( cli: AgentbricksCli, workspace: Workspace, name: str, + identity: AppIdentity, *, memory_store: str | None = None, session_store: str | None = None, - identity: AppIdentity | None = None, ) -> list[str]: - """Delete one project's stores and everything its deploy created, as far as any of it exists. + """Delete one project's stores and App through the CLI, and the catalog grant its test made. - The App, its Runtime Store and the stores go through the CLI. The rest are workarounds for what - no public command deletes yet: the synced source folder, and the Lakebase databases and roles a - failed deploy leaves. ``identity`` is what was recorded earlier; it fills in whatever the live - lookup misses. Every step runs regardless of earlier failures. Returns a description of each - failure; an empty list means everything was deleted. + ``identity`` is the App as it was before any delete; record it first, because + ``delete_run_leftovers`` needs it for what the CLI leaves behind. Every step runs regardless of + earlier failures. Returns a description of each failure; an empty list means all succeeded. """ failures: list[str] = [] - def attempt(label: str, action: Callable[..., object], *args: Any) -> bool: + def attempt(label: str, action: Callable[..., object], *args: Any) -> None: try: action(*args) except Exception as exc: failures.append(f"{label}: {exc}") - return False - return True def run(command: str) -> None: result = cli.run(command, timeout=1200, check=False) @@ -135,32 +137,16 @@ def run(command: str) -> None: if result.returncode != 0 and not _MISSING.search(output): raise E2EError(output) - # Read before any delete: the principal and branch come from the App and its Runtime Store. - known = (identity or AppIdentity()).merged(workspace.app_identity(name)) if memory_store: attempt(f"memory store {memory_store}", run, f"memory stores delete {memory_store} --yes") if session_store: attempt( f"session store {session_store}", run, f"sessions stores delete {session_store} --yes" ) - if known == AppIdentity(): - return failures - + # Before the App goes: the grant is revoked from the App's principal. attempt(f"catalog grant {name}", workspace.revoke_catalog_grant, name) - if known.service_principal: - # Deletes the Runtime Store first, while the App's principal still resolves, then the App. + if identity.service_principal: attempt(f"deployment {name}", run, f"deployments delete {name} --yes") - else: - attempt(f"runtime store {name}", workspace.delete_runtime_store, name) - if attempt(f"app {name}", workspace.wait_until_app_gone, name): - attempt(f"deployment source {name}", workspace.delete_deployment_source, name) - if known.branch: - failures += delete_lakebase_leftovers( - workspace, - known.branch, - principals=[known.service_principal or ""], - database_ids=[known.database_id or ""], - ) return failures @@ -173,7 +159,9 @@ def delete_run_leftovers( """Delete what a whole run left behind; returns a description of each failure.""" failures: list[str] = [] steps = ( + # First: the rest outlives the App, and some of it can't be deleted while the App exists. ("leaked Apps", lambda: _delete_leaked_apps(cli, workspace, run_id, output)), + ("orphaned Runtime Stores", lambda: _delete_orphaned_runtime_stores(workspace, output)), ("leaked Lakebase databases and roles", lambda: _delete_leaked_lakebase(workspace, output)), ("leaked source folders", lambda: _delete_leaked_deployment_sources(workspace, run_id)), ("leaked schemas", lambda: _delete_leaked_schemas(workspace, catalog, run_id)), @@ -193,16 +181,46 @@ def _app_prefix(run_id: str) -> str: 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, found by the run's name prefix.""" + """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) - recorded = recorded_app_identities(output) - deployments = json.loads(cli.run("--output json deployments list").stdout or "[]") - failures: list[str] = [] - for deployment in deployments: + for deployment in json.loads(cli.run("--output json deployments list").stdout or "[]"): name = str(deployment.get("name") or "") - if name.startswith(prefix): - log(f"deleting leaked App {name}") - failures += delete_app_resources(cli, workspace, name, identity=recorded.get(name)) + if not name.startswith(prefix): + continue + log(f"deleting leaked App {name}") + # The principal stops resolving once the App is gone, and the later steps need it. + record_app_identity(output, name, workspace.app_identity(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()}") + failures: list[str] = [] + for app in workspace.client.apps.list(): + if app.name and app.name.startswith(prefix): + try: + workspace.wait_until_app_gone(app.name) + except Exception as exc: + failures.append(f"App {app.name}: {exc}") + return failures + + +def _delete_orphaned_runtime_stores(workspace: Workspace, output: pathlib.Path) -> list[str]: + """Delete the Runtime Stores of this run's Apps that outlived them. + + `deployments delete` refuses once the App is gone, because it checks the store's owner against + the App's principal. + """ + failures: list[str] = [] + for name in recorded_app_identities(output): + try: + workspace.client.api_client.do("DELETE", f"/api/2.0/agents/runtime-stores/{name}") + except NotFound: + pass + except Exception as exc: + failures.append(f"runtime store {name}: {exc}") return failures diff --git a/integrations/agentbricks/tests/harness/workspace.py b/integrations/agentbricks/tests/harness/workspace.py index dd4792a06..df9844fb5 100644 --- a/integrations/agentbricks/tests/harness/workspace.py +++ b/integrations/agentbricks/tests/harness/workspace.py @@ -310,8 +310,7 @@ def grant_transitive(self, deployment: dict[str, Any], function: str) -> dict[st "effective_privileges": effective, } - # Workarounds: cleanup that no public command does yet. `agentbricks deployments delete` - # removes the App and its Runtime Store; these cover what it leaves behind. + # Workarounds for the end-of-run cleanup: what `agentbricks deployments delete` leaves behind. def wait_until_app_gone(self, name: str, timeout: float = 1200) -> None: """Block until the App is gone; `apps delete` returns while it is still deleting.""" @@ -339,16 +338,6 @@ def runtime_store(self, app_name: str) -> dict[str, Any]: self.client.api_client.do("GET", f"/api/2.0/agents/runtime-stores/{app_name}"), ) - def delete_runtime_store(self, app_name: str) -> None: - """Delete a Runtime Store whose App is gone, which `deployments delete` cannot reach. - - A missing store counts as deleted. - """ - try: - self.client.api_client.do("DELETE", f"/api/2.0/agents/runtime-stores/{app_name}") - except NotFound: - pass - def app_identity(self, app_name: str) -> AppIdentity: """The App's principal and Lakebase location as far as they still exist; never raises.""" principal: str | None = None @@ -439,10 +428,6 @@ def user_name(self) -> str: def _deployments_dir(self) -> str: return f"/Workspace/Users/{self.user_name}/agentbricks_deployments" - def delete_deployment_source(self, app_name: str) -> None: - """Delete the App's synced source folder; a missing folder counts as deleted.""" - self.delete_workspace_path(f"{self._deployments_dir()}/{app_name}") - def deployment_sources_with_prefix(self, prefix: str) -> list[str]: try: entries = list(self.client.workspace.list(self._deployments_dir())) diff --git a/integrations/agentbricks/tests/suites/README.md b/integrations/agentbricks/tests/suites/README.md index a0e77f666..e8dd79769 100644 --- a/integrations/agentbricks/tests/suites/README.md +++ b/integrations/agentbricks/tests/suites/README.md @@ -20,7 +20,7 @@ uv run pytest tests/suites -n 6 --dist loadgroup \ - 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 "`. - `--workspace fake` runs against no workspace; everything that needs one is skipped. -- `--keep-resources` keeps leftovers for debugging; each test still deletes its own resources. +- `--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 )`. @@ -29,11 +29,13 @@ The run id and output directory (logs, generated projects, `apps.jsonl`) are pri Every resource a run creates carries its run id, and is deleted at three points: -1. **When each test ends**, by the fixture that made it: `project` deletes the test's App, its stores, - and everything its deploys created (`harness.cleanup.delete_app_resources`); `scratch_schema` - deletes the test's schema. -2. **When the run ends**, `pytest_sessionfinish` deletes whatever a crashed test or a failed teardown - left (`harness.cleanup.delete_run_leftovers`). A failure there fails the run. +1. **When each test ends**, by the fixture that made it: `project` deletes the test's App and its + stores through the CLI, as a user would (`harness.cleanup.delete_app_resources`); + `scratch_schema` deletes the test's schema. +2. **When the run ends**, `pytest_sessionfinish` deletes what the CLI leaves behind (synced source + folders, Runtime Stores whose App is gone, Lakebase databases and roles) and whatever a crashed + test or a failed teardown left (`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 @@ -52,7 +54,7 @@ Every resource a run creates carries its run id, and is deleted at three points: | `tests/harness/agentbricks_cli.py` | Runs `agentbricks` commands: `cli.run("tools add ...")`, and `cli.start("dev ...")` in the background | | `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), UC privilege checks, and cleanup no CLI command does yet | -| `tests/harness/cleanup.py` | Deleting what the suites create, per App and per run; the App and its Runtime Store go through `agentbricks deployments delete` | +| `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//` | One folder per feature: `tools`, `memory_store`, `session_store`, `tracing` | diff --git a/integrations/agentbricks/tests/suites/conftest.py b/integrations/agentbricks/tests/suites/conftest.py index a913d73e9..01cbf1843 100644 --- a/integrations/agentbricks/tests/suites/conftest.py +++ b/integrations/agentbricks/tests/suites/conftest.py @@ -264,7 +264,7 @@ def _delete_project( identity = workspace.app_identity(name) deployed = identity != AppIdentity() if deployed: - # Leftover cleanup retries the Lakebase deletes from this record if they fail here. + # Leftover cleanup deletes what the CLI leaves behind from this record. record_app_identity(cli.run_config.output, name, identity) if failed and identity.service_principal: # The App is deleted next, so a failing test must capture its logs while it exists. @@ -276,9 +276,9 @@ def _delete_project( cli, workspace, name, + identity, memory_store=f"memory-stores/{memory_store}" if memory_store else None, session_store=project.declared_store("session"), - identity=identity, ) except Exception as exc: failures = [str(exc)] From d3f7331e5d14fa71586dfdf330b5e417dcbaf5c4 Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Thu, 15 Oct 2026 20:31:12 +0000 Subject: [PATCH 25/45] [agentbricks] E2E suites: run per-test cleanup commands directly in the project fixture Co-authored-by: Isaac --- .../agentbricks/tests/harness/cleanup.py | 68 +++---------------- .../agentbricks/tests/suites/README.md | 3 +- .../agentbricks/tests/suites/conftest.py | 34 +++++----- 3 files changed, 29 insertions(+), 76 deletions(-) diff --git a/integrations/agentbricks/tests/harness/cleanup.py b/integrations/agentbricks/tests/harness/cleanup.py index e7bfbbb3c..9acfd6328 100644 --- a/integrations/agentbricks/tests/harness/cleanup.py +++ b/integrations/agentbricks/tests/harness/cleanup.py @@ -1,16 +1,14 @@ """Deleting what the E2E suites create: per App when a test ends, and per run afterwards. -Two entry points, both safe to repeat (a missing resource counts as deleted): - -- ``delete_app_resources``: one App and its stores, deleted the way a user would, through the CLI. - The ``project`` fixture calls it when its test ends. -- ``delete_run_leftovers``: everything else, including whatever the CLI leaves behind (synced - source folders, Runtime Stores whose App is gone, Lakebase databases and roles) and whatever a - crashed test or a failed teardown left. It runs once after every live run, and on its own with - ``pytest tests/suites --cleanup-only --run-id `` after a run that was killed before it could. - It finds resources by the run's name prefix (Apps, source folders, schemas) and by the App - identities per-test cleanup recorded in ``/apps.jsonl`` (Runtime Stores, Lakebase - databases and roles, which outlive their App), so it touches nothing outside the run. +The ``project`` fixture deletes each test's App and stores through the CLI when the test ends, after +recording the App's identity with ``record_app_identity``. ``delete_run_leftovers`` deletes the rest: +what the CLI leaves behind (synced source folders, Runtime Stores whose App is gone, Lakebase +databases and roles) and whatever a crashed test or a failed teardown left. 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. +It finds resources by the run's name prefix (Apps, source folders, schemas) and by the identities +recorded in ``/apps.jsonl`` (Runtime Stores, Lakebase databases and roles, which outlive +their App), so it touches nothing outside the run. """ from __future__ import annotations @@ -18,19 +16,16 @@ import dataclasses import json import pathlib -import re from collections.abc import Callable, Collection -from typing import Any from databricks.sdk.errors import NotFound from harness.agentbricks_cli import AgentbricksCli from harness.logs import log -from harness.run import APP_PREFIX, E2EError, project_prefix, scratch_schema_prefix +from harness.run import APP_PREFIX, project_prefix, scratch_schema_prefix from harness.workspace import AppIdentity, Workspace _APPS_RECORD = "apps.jsonl" -_MISSING = re.compile(r"NOT_FOUND|RESOURCE_DOES_NOT_EXIST|does not exist|not found", re.IGNORECASE) # Identities recorded by per-test cleanup @@ -107,49 +102,6 @@ def delete_lakebase_leftovers( ) -def delete_app_resources( - cli: AgentbricksCli, - workspace: Workspace, - name: str, - identity: AppIdentity, - *, - memory_store: str | None = None, - session_store: str | None = None, -) -> list[str]: - """Delete one project's stores and App through the CLI, and the catalog grant its test made. - - ``identity`` is the App as it was before any delete; record it first, because - ``delete_run_leftovers`` needs it for what the CLI leaves behind. Every step runs regardless of - earlier failures. Returns a description of each failure; an empty list means all succeeded. - """ - failures: list[str] = [] - - def attempt(label: str, action: Callable[..., object], *args: Any) -> None: - try: - action(*args) - except Exception as exc: - failures.append(f"{label}: {exc}") - - def run(command: str) -> None: - result = cli.run(command, timeout=1200, check=False) - output = (result.stderr or result.stdout).strip() - # A store the test declared but never deployed was never created. - if result.returncode != 0 and not _MISSING.search(output): - raise E2EError(output) - - if memory_store: - attempt(f"memory store {memory_store}", run, f"memory stores delete {memory_store} --yes") - if session_store: - attempt( - f"session store {session_store}", run, f"sessions stores delete {session_store} --yes" - ) - # Before the App goes: the grant is revoked from the App's principal. - attempt(f"catalog grant {name}", workspace.revoke_catalog_grant, name) - if identity.service_principal: - attempt(f"deployment {name}", run, f"deployments delete {name} --yes") - return failures - - # One run diff --git a/integrations/agentbricks/tests/suites/README.md b/integrations/agentbricks/tests/suites/README.md index e8dd79769..1f65d2894 100644 --- a/integrations/agentbricks/tests/suites/README.md +++ b/integrations/agentbricks/tests/suites/README.md @@ -30,8 +30,7 @@ The run id and output directory (logs, generated projects, `apps.jsonl`) are pri Every resource a run creates carries its run id, and is deleted at three points: 1. **When each test ends**, by the fixture that made it: `project` deletes the test's App and its - stores through the CLI, as a user would (`harness.cleanup.delete_app_resources`); - `scratch_schema` deletes the test's schema. + stores through the CLI, as a user would; `scratch_schema` deletes the test's schema. 2. **When the run ends**, `pytest_sessionfinish` deletes what the CLI leaves behind (synced source folders, Runtime Stores whose App is gone, Lakebase databases and roles) and whatever a crashed test or a failed teardown left (`harness.cleanup.delete_run_leftovers`). A failure there fails diff --git a/integrations/agentbricks/tests/suites/conftest.py b/integrations/agentbricks/tests/suites/conftest.py index 01cbf1843..b558ab69b 100644 --- a/integrations/agentbricks/tests/suites/conftest.py +++ b/integrations/agentbricks/tests/suites/conftest.py @@ -27,7 +27,7 @@ import pytest from harness.agentbricks_cli import AgentbricksCli -from harness.cleanup import delete_app_resources, delete_run_leftovers, record_app_identity +from harness.cleanup import delete_run_leftovers, record_app_identity from harness.logs import log, log_command from harness.project import Project from harness.run import RUN_ID_PATTERN, E2EError, RunConfig, run_output_dir @@ -38,7 +38,7 @@ make_target_workspace, require_workspace, ) -from harness.workspace import AppIdentity, Workspace +from harness.workspace import Workspace PACKAGE_DIR = pathlib.Path(__file__).resolve().parents[2] _RUN_KEY = pytest.StashKey[RunConfig]() @@ -262,28 +262,30 @@ def _delete_project( # Whatever deploy got as far as creating; read before anything is deleted, because the # principal and branch come from the App and its Runtime Store. identity = workspace.app_identity(name) - deployed = identity != AppIdentity() - if deployed: + if identity.service_principal or identity.branch: # Leftover cleanup deletes what the CLI leaves behind from this record. record_app_identity(cli.run_config.output, name, identity) if failed and identity.service_principal: # The App is deleted next, so a failing test must capture its logs while it exists. workspace.app_logs(name, cli.logs_dir / f"deploy-runtime-{name}.log") # Deploy creates the stores agent.toml declares; memory stores are addressed by resource name. - memory_store = project.declared_store("memory") + commands = [] + if memory_store := project.declared_store("memory"): + commands.append(f"memory stores delete memory-stores/{memory_store} --yes") + if session_store := project.declared_store("session"): + commands.append(f"sessions stores delete {session_store} --yes") + if identity.service_principal: + commands.append(f"deployments delete {name} --yes") try: - failures = delete_app_resources( - cli, - workspace, - name, - identity, - memory_store=f"memory-stores/{memory_store}" if memory_store else None, - session_store=project.declared_store("session"), - ) + # The grant the test made, revoked while the App's principal still exists. + workspace.revoke_catalog_grant(name) except Exception as exc: - failures = [str(exc)] - for failure in failures: - log(f"cleanup warning | {name} | {failure}") + log(f"cleanup warning | {name} | catalog grant: {exc}") + # Failures only warn: the end-of-run cleanup deletes whatever is left. + for command in commands: + result = cli.run(command, timeout=1200, check=False) + if result.returncode != 0: + log(f"cleanup warning | {name} | {(result.stderr or result.stdout).strip()}") # Deleting what the whole run left behind From 8d46b1cc9fd441d609a71ca1300640df09e03be9 Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Thu, 15 Oct 2026 21:06:14 +0000 Subject: [PATCH 26/45] [agentbricks] E2E suites: direct authoring appends tools to init's agent.toml Co-authored-by: Isaac --- .../agentbricks/tests/harness/project.py | 23 ++++++++----------- integrations/agentbricks/tests/harness/run.py | 1 - 2 files changed, 9 insertions(+), 15 deletions(-) diff --git a/integrations/agentbricks/tests/harness/project.py b/integrations/agentbricks/tests/harness/project.py index 2d109267a..42108b725 100644 --- a/integrations/agentbricks/tests/harness/project.py +++ b/integrations/agentbricks/tests/harness/project.py @@ -8,7 +8,6 @@ import dataclasses import pathlib -import re import shutil import uuid from collections.abc import Mapping, Sequence @@ -18,11 +17,7 @@ import tomlkit from harness.logs import log -from harness.run import APP_PREFIX, FRAMEWORKS, RunConfig, project_prefix - -_DIRECT_HEADER = ( - 'schema_version = 1\n\n[agent]\nframework = "{framework}"\nserver = "agentbricks"\n' -) +from harness.run import APP_PREFIX, RunConfig, project_prefix class ToolSpec(Protocol): @@ -83,16 +78,16 @@ def write_files(self, files: Mapping[str, str]) -> None: self.write_file(relative_path, content) def write_manifest(self, tool_specs: Sequence[ToolSpec]) -> None: - """Direct authoring: write agent.toml with exactly these tools plus a tracing binding.""" - sections: list[str] = [_DIRECT_HEADER.format(framework=FRAMEWORKS[0])] - sections.extend(spec.toml for spec in tool_specs if spec.toml) - # Follows default_experiment_name's shape so direct authoring carries the same tracing - # binding `agentbricks init` gives CLI authoring. - slug = re.sub(r"[^a-z0-9-]+", "-", self.path.name.lower()).strip("-") or "agent" - sections.append(f'[tracing]\nexperiment_name = "/Shared/agentbricks_traces/{slug}"\n') + """Direct authoring: append these tools to the agent.toml `agentbricks init` wrote. + + Everything else, such as the stores and the tracing binding, stays as init wrote it, the + same as for CLI authoring. + """ target = self.path / "agent.toml" + tools = "".join(f"\n{spec.toml}" for spec in tool_specs if spec.toml) log(f"# write {target}: direct authoring; no agentbricks tools command") - target.write_text("\n".join(sections), encoding="utf-8") + with target.open("a", encoding="utf-8") as manifest: + manifest.write(tools) for spec in tool_specs: self.write_files(spec.files) diff --git a/integrations/agentbricks/tests/harness/run.py b/integrations/agentbricks/tests/harness/run.py index 725423f7c..8ef8e8ab0 100644 --- a/integrations/agentbricks/tests/harness/run.py +++ b/integrations/agentbricks/tests/harness/run.py @@ -12,7 +12,6 @@ import re import tempfile -FRAMEWORKS = ("langgraph",) AUTHORING_PATHS = ("cli", "direct") E2E_MODEL = "system.ai.gpt-5-2" TOOL_RESOURCE_PREFIX = "agentbricks-tool-" From 5556730de600582c287f3c064bc3312f79a8d9ff Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Thu, 15 Oct 2026 21:08:13 +0000 Subject: [PATCH 27/45] [agentbricks] E2E suites: drop the unused --bridge-sha mode Co-authored-by: Isaac --- .../agentbricks/tests/harness/project.py | 47 ++++--------------- integrations/agentbricks/tests/harness/run.py | 1 - .../agentbricks/tests/suites/conftest.py | 8 +--- 3 files changed, 9 insertions(+), 47 deletions(-) diff --git a/integrations/agentbricks/tests/harness/project.py b/integrations/agentbricks/tests/harness/project.py index 42108b725..29b85f46b 100644 --- a/integrations/agentbricks/tests/harness/project.py +++ b/integrations/agentbricks/tests/harness/project.py @@ -17,7 +17,7 @@ import tomlkit from harness.logs import log -from harness.run import APP_PREFIX, RunConfig, project_prefix +from harness.run import APP_PREFIX, E2EError, RunConfig, project_prefix class ToolSpec(Protocol): @@ -31,9 +31,8 @@ class ToolSpec(Protocol): class Project: path: pathlib.Path app_name: str - # The runtime under test: the wheel built from this checkout, or a bridge commit. - wheel: pathlib.Path | None = None - bridge_sha: str | None = None + # 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: @@ -41,15 +40,14 @@ def allocate(cls, run_config: RunConfig) -> Project: 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) - return cls(path, f"{APP_PREFIX}{name}", run_config.wheel, run_config.bridge_sha) + 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.""" - if self.bridge_sha: - _pin_bridge_sources(self.path, self.bridge_sha) - else: - assert self.wheel is not None - _pin_project_wheel(self.path, self.wheel) + _pin_project_wheel(self.path, self.wheel) def is_pinned(self) -> bool: pyproject = self.path / "pyproject.toml" @@ -108,32 +106,3 @@ def _pin_project_wheel(project: pathlib.Path, wheel: pathlib.Path) -> None: f"pin databricks-agentbricks runtime to {vendored_wheel.relative_to(project)}" ) pyproject.write_text(tomlkit.dumps(document), encoding="utf-8") - - -def _pin_bridge_sources(project: pathlib.Path, bridge_sha: str) -> None: - pyproject = project / "pyproject.toml" - document = tomlkit.parse(pyproject.read_text(encoding="utf-8")) - dependencies = document["project"]["dependencies"] - if not any(str(dependency).startswith("databricks-langchain") for dependency in dependencies): - dependencies.append("databricks-langchain>=0.17.0") - if "tool" not in document: - document["tool"] = tomlkit.table() - if "uv" not in document["tool"]: - document["tool"]["uv"] = tomlkit.table() - if "sources" not in document["tool"]["uv"]: - document["tool"]["uv"]["sources"] = tomlkit.table() - for package, subdirectory in ( - ("databricks-agentbricks", "integrations/agentbricks"), - ("databricks-langchain", "integrations/langchain"), - ): - source = tomlkit.inline_table() - source.update( - { - "git": "https://github.com/databricks/databricks-ai-bridge.git", - "rev": bridge_sha, - "subdirectory": subdirectory, - } - ) - document["tool"]["uv"]["sources"][package] = source - log(f"# write {pyproject}: pin Agent Bricks and LangChain to {bridge_sha}") - 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 index 8ef8e8ab0..67da3131d 100644 --- a/integrations/agentbricks/tests/harness/run.py +++ b/integrations/agentbricks/tests/harness/run.py @@ -54,5 +54,4 @@ class RunConfig: # Operator override; None lets the workspace pick the first running warehouse. warehouse_id: str | None preprovisioned_app_catalog_access: bool - bridge_sha: str | None wheel: pathlib.Path | None diff --git a/integrations/agentbricks/tests/suites/conftest.py b/integrations/agentbricks/tests/suites/conftest.py index b558ab69b..54ce37306 100644 --- a/integrations/agentbricks/tests/suites/conftest.py +++ b/integrations/agentbricks/tests/suites/conftest.py @@ -19,7 +19,6 @@ import dataclasses import os import pathlib -import re import subprocess import uuid from collections.abc import Iterator @@ -67,7 +66,6 @@ def pytest_addoption(parser: pytest.Parser) -> None: default="main", help="Existing catalog for temporary test schemas (default: main).", ) - group.addoption("--bridge-sha", help="Immutable bridge commit for generated App dependencies.") group.addoption( "--wheel", type=pathlib.Path, @@ -104,9 +102,6 @@ def pytest_addoption(parser: pytest.Parser) -> None: def pytest_configure(config: pytest.Config) -> None: - sha = config.getoption("bridge_sha") - if sha is not None and re.fullmatch(r"[0-9a-f]{40}", sha) is None: - raise pytest.UsageError("--bridge-sha must be a 40-character lowercase Git commit SHA.") 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.") @@ -149,7 +144,6 @@ def _base_run_config(config: pytest.Config) -> RunConfig: output=output, warehouse_id=option("warehouse_id"), preprovisioned_app_catalog_access=bool(option("preprovisioned_app_catalog_access")), - bridge_sha=option("bridge_sha"), wheel=option("wheel"), ) config.stash[_RUN_KEY] = run @@ -169,7 +163,7 @@ 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 or run.bridge_sha: + 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") From 410b6e16fbef8e1381859ec98fcf55c4858078e9 Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Thu, 15 Oct 2026 21:15:23 +0000 Subject: [PATCH 28/45] [agentbricks] E2E suites: stop relying on the CLI's tool-resource name prefix Co-authored-by: Isaac --- integrations/agentbricks/tests/harness/run.py | 1 - .../agentbricks/tests/harness/workspace.py | 30 +++++-------------- .../tests/suites/tools/test_grants.py | 11 ++----- .../tests/suites/tracing/test_tracing.py | 10 +++---- 4 files changed, 13 insertions(+), 39 deletions(-) diff --git a/integrations/agentbricks/tests/harness/run.py b/integrations/agentbricks/tests/harness/run.py index 67da3131d..a0fe29b06 100644 --- a/integrations/agentbricks/tests/harness/run.py +++ b/integrations/agentbricks/tests/harness/run.py @@ -14,7 +14,6 @@ AUTHORING_PATHS = ("cli", "direct") E2E_MODEL = "system.ai.gpt-5-2" -TOOL_RESOURCE_PREFIX = "agentbricks-tool-" APP_PREFIX = "agent-bricks-" diff --git a/integrations/agentbricks/tests/harness/workspace.py b/integrations/agentbricks/tests/harness/workspace.py index df9844fb5..c6dd357b8 100644 --- a/integrations/agentbricks/tests/harness/workspace.py +++ b/integrations/agentbricks/tests/harness/workspace.py @@ -25,7 +25,7 @@ from databricks.sdk.service.sql import ExecuteStatementRequestOnWaitTimeout, State, StatementState from harness.logs import log -from harness.run import TOOL_RESOURCE_PREFIX, E2EError, now +from harness.run import E2EError, now GrantTuple = tuple[str, str, str, str] @@ -250,9 +250,9 @@ def transitive_state(self, deployment: dict[str, Any], function: str) -> dict[st ) if any( resource.get("uc_securable", {}).get("securable_full_name") == function - for resource in tool_resources(deployment) + for resource in app_resources(deployment) ): - raise E2EError("The transitive function appeared in Agent Bricks-owned Apps resources.") + raise E2EError("The transitive function appeared among the App's resources.") return { "service_principal_client_id": principal, "direct_privileges": direct, @@ -468,30 +468,14 @@ def granted(deployment: dict[str, Any], grant: GrantTuple) -> bool: def granted_tuples(deployment: dict[str, Any]) -> set[GrantTuple]: - return app_resource_tuples(tool_resources(deployment)) + return app_resource_tuples(app_resources(deployment)) -def _split_resources(app: dict[str, Any]) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: - resources = app.get("resources") or [] +def app_resources(deployment: dict[str, Any]) -> list[dict[str, Any]]: + resources = deployment.get("resources") or [] if not isinstance(resources, list): raise E2EError(f"App resources are not a list: {resources}") - - def name(resource: dict[str, Any]) -> str: - return str(resource.get("name", "")) - - owned = [r for r in resources if name(r).startswith(TOOL_RESOURCE_PREFIX)] - other = [r for r in resources if not name(r).startswith(TOOL_RESOURCE_PREFIX)] - return sorted(owned, key=name), sorted(other, key=name) - - -def tool_resources(app: dict[str, Any]) -> list[dict[str, Any]]: - """The App resources Agent Bricks owns, i.e. those it names with ``TOOL_RESOURCE_PREFIX``.""" - return _split_resources(app)[0] - - -def non_tool_resources(app: dict[str, Any]) -> list[dict[str, Any]]: - """The App resources Agent Bricks does not own, e.g. the tracing experiment.""" - return _split_resources(app)[1] + return resources def app_resource_tuples(resources: Sequence[dict[str, Any]]) -> set[GrantTuple]: diff --git a/integrations/agentbricks/tests/suites/tools/test_grants.py b/integrations/agentbricks/tests/suites/tools/test_grants.py index ddfc2f6ed..99c7fb2ee 100644 --- a/integrations/agentbricks/tests/suites/tools/test_grants.py +++ b/integrations/agentbricks/tests/suites/tools/test_grants.py @@ -1,4 +1,4 @@ -"""Deploy grants that belong to no single tool (deploy-only): nothing transitive, nothing lost.""" +"""Deploy grants only what a tool binds: a function the bound UC function calls needs a manual grant.""" from __future__ import annotations @@ -6,8 +6,7 @@ from harness.agentbricks_cli import AgentbricksCli from harness.project import Project -from harness.run import TOOL_RESOURCE_PREFIX -from harness.workspace import Workspace, non_tool_resources +from harness.workspace import Workspace from tool_definitions import assert_declared, uc_function_tool from uc_objects import UcFunction @@ -21,7 +20,6 @@ def test_grants( subtests, ) -> None: tool = uc_function_tool(uc_function) - # `init` and direct authoring both bind tracing, a resource Agent Bricks does not own. agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) project.pin_runtime() app = project.app_name @@ -41,8 +39,3 @@ def test_grants( granted = workspace_client.grant_transitive(deployment, uc_function.nested_function) assert "EXECUTE" in granted["direct_privileges"] assert "EXECUTE" in granted["effective_privileges"] - - with subtests.test(check="non-tool-resource-preserved"): - unrelated = non_tool_resources(deployment) - assert unrelated, "No non-tool App resource; the tracing experiment should be attached" - assert not any(r["name"].startswith(TOOL_RESOURCE_PREFIX) for r in unrelated) diff --git a/integrations/agentbricks/tests/suites/tracing/test_tracing.py b/integrations/agentbricks/tests/suites/tracing/test_tracing.py index 55f707c1b..8ce92b0da 100644 --- a/integrations/agentbricks/tests/suites/tracing/test_tracing.py +++ b/integrations/agentbricks/tests/suites/tracing/test_tracing.py @@ -10,8 +10,7 @@ import pytest from harness.agentbricks_cli import AgentbricksCli from harness.project import Project -from harness.run import TOOL_RESOURCE_PREFIX -from harness.workspace import Workspace, non_tool_resources +from harness.workspace import Workspace, app_resources def test_tracing_experiment_bound( @@ -26,10 +25,9 @@ def test_tracing_experiment_bound( agentbricks_cli.run(f"--output json deployments get {project.app_name}").stdout ) - # Deploy attaches the experiment as an App resource that Agent Bricks does not own. - experiments = [r for r in non_tool_resources(deployment) if "experiment" in r] - assert experiments, "No experiment App resource among the non-tool resources" - assert not any(r["name"].startswith(TOOL_RESOURCE_PREFIX) for r in experiments) + # Deploy attaches the experiment as an App resource. + experiments = [r for r in app_resources(deployment) if "experiment" in r] + assert experiments, "No experiment App resource" expected_id = workspace_client.experiment_id(experiment_name) assert any(str(r["experiment"].get("experiment_id")) == expected_id for r in experiments), ( f"{experiment_name} ({expected_id}) not in {experiments}" From ae8f16b5cc5e6e8db6e328b216e6e7aff3bd0b5b Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Thu, 15 Oct 2026 21:34:10 +0000 Subject: [PATCH 29/45] [agentbricks] E2E suites: simplify the sweep to the run prefix plus App roles Co-authored-by: Isaac --- .../agentbricks/tests/harness/cleanup.py | 158 ++++++++---------- .../agentbricks/tests/harness/workspace.py | 43 +---- .../agentbricks/tests/suites/README.md | 8 +- 3 files changed, 72 insertions(+), 137 deletions(-) diff --git a/integrations/agentbricks/tests/harness/cleanup.py b/integrations/agentbricks/tests/harness/cleanup.py index 9acfd6328..bde2768aa 100644 --- a/integrations/agentbricks/tests/harness/cleanup.py +++ b/integrations/agentbricks/tests/harness/cleanup.py @@ -1,14 +1,12 @@ """Deleting what the E2E suites create: per App when a test ends, and per run afterwards. The ``project`` fixture deletes each test's App and stores through the CLI when the test ends, after -recording the App's identity with ``record_app_identity``. ``delete_run_leftovers`` deletes the rest: -what the CLI leaves behind (synced source folders, Runtime Stores whose App is gone, Lakebase -databases and roles) and whatever a crashed test or a failed teardown left. It is safe to repeat (a -missing resource counts as deleted), runs once after every live run, and runs on its own with +recording the App's identity with ``record_app_identity``. ``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 recorded +identities 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. -It finds resources by the run's name prefix (Apps, source folders, schemas) and by the identities -recorded in ``/apps.jsonl`` (Runtime Stores, Lakebase databases and roles, which outlive -their App), so it touches nothing outside the run. """ from __future__ import annotations @@ -16,9 +14,8 @@ import dataclasses import json import pathlib -from collections.abc import Callable, Collection - -from databricks.sdk.errors import NotFound +from collections.abc import Iterator +from typing import Any from harness.agentbricks_cli import AgentbricksCli from harness.logs import log @@ -50,58 +47,13 @@ def recorded_app_identities(output: pathlib.Path) -> dict[str, AppIdentity]: for line in path.read_text(encoding="utf-8").splitlines(): try: data = json.loads(line) - identity = AppIdentity( - data.get("service_principal"), data.get("branch"), data.get("database_id") - ) + 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 App - - -def _delete_each( - kind: str, find: Callable[[], list[str]], delete: Callable[[str], None] -) -> list[str]: - """Delete everything ``find`` returns, continuing past failures; returns one line per failure.""" - try: - names = find() - except Exception as exc: - return [f"{kind} lookup: {exc}"] - failures = [] - for name in names: - try: - delete(name) - except Exception as exc: - failures.append(f"{kind} {name}: {exc}") - return failures - - -def delete_lakebase_leftovers( - workspace: Workspace, - branch: str, - *, - principals: Collection[str], - database_ids: Collection[str] = (), -) -> list[str]: - """Delete runtime databases and then the principals' roles on one branch. - - A role that owns a database cannot be deleted, so databases go first. Only ``runtime-`` - databases are touched, and both kinds are matched by principal UUID, never by name pattern. - """ - return _delete_each( - "runtime database", - lambda: workspace.runtime_databases(branch, principals, database_ids), - workspace.delete_database, - ) + _delete_each( - "Lakebase role", - lambda: workspace.roles_of(branch, principals), - workspace.delete_role, - ) - - # One run @@ -111,10 +63,10 @@ def delete_run_leftovers( """Delete what a whole run left behind; returns a description of each failure.""" failures: list[str] = [] steps = ( - # First: the rest outlives the App, and some of it can't be deleted while the App exists. + ("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)), - ("orphaned Runtime Stores", lambda: _delete_orphaned_runtime_stores(workspace, output)), - ("leaked Lakebase databases and roles", lambda: _delete_leaked_lakebase(workspace, output)), + ("leaked Lakebase roles", lambda: _delete_leaked_roles(workspace, output)), ("leaked source folders", lambda: _delete_leaked_deployment_sources(workspace, run_id)), ("leaked schemas", lambda: _delete_leaked_schemas(workspace, catalog, run_id)), ) @@ -130,6 +82,42 @@ 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]: @@ -159,42 +147,28 @@ def _delete_leaked_apps( return failures -def _delete_orphaned_runtime_stores(workspace: Workspace, output: pathlib.Path) -> list[str]: - """Delete the Runtime Stores of this run's Apps that outlived them. - - `deployments delete` refuses once the App is gone, because it checks the store's owner against - the App's principal. - """ - failures: list[str] = [] - for name in recorded_app_identities(output): - try: - workspace.client.api_client.do("DELETE", f"/api/2.0/agents/runtime-stores/{name}") - except NotFound: - pass - except Exception as exc: - failures.append(f"runtime store {name}: {exc}") - return failures - - -def _delete_leaked_lakebase(workspace: Workspace, output: pathlib.Path) -> list[str]: - """Delete runtime databases and roles of this run's Apps, whether or not the App still exists.""" +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}): - failures += delete_lakebase_leftovers( - workspace, - 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 - ], - database_ids=[ - item.database_id for item in recorded if item.branch == branch and item.database_id - ], - ) + # 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 diff --git a/integrations/agentbricks/tests/harness/workspace.py b/integrations/agentbricks/tests/harness/workspace.py index c6dd357b8..a92492685 100644 --- a/integrations/agentbricks/tests/harness/workspace.py +++ b/integrations/agentbricks/tests/harness/workspace.py @@ -30,23 +30,20 @@ GrantTuple = tuple[str, str, str, str] _BRANCH = re.compile(r"projects/[^/]+/branches/[^/]+") -_RUNTIME_DATABASE_PREFIX = "runtime-" @dataclasses.dataclass(frozen=True) class AppIdentity: - """What cleanup needs to find an App's Lakebase leftovers once the App and its Runtime Store are gone.""" + """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 - database_id: 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, - self.database_id or other.database_id, ) @@ -342,7 +339,6 @@ def app_identity(self, app_name: str) -> AppIdentity: """The App's principal and Lakebase location as far as they still exist; never raises.""" principal: str | None = None branch: str | None = None - database_id: str | None = None try: principal = self.app(app_name).get("service_principal_client_id") except NotFound: @@ -354,40 +350,11 @@ def app_identity(self, app_name: str) -> AppIdentity: candidate = lakebase.get("branch") if isinstance(candidate, str) and _BRANCH.fullmatch(candidate): branch = candidate - database_id = lakebase.get("database_id") except NotFound: pass except Exception as exc: log(f"cleanup warning | Runtime Store lookup for {app_name} | {exc}") - return AppIdentity(principal, branch, database_id) - - def runtime_databases( - self, branch: str, principals: Collection[str], database_ids: Collection[str] - ) -> list[str]: - """Resource names of the branch's ``runtime-`` databases owned by these principals. - - Database ids embed a truncated App name plus a UUID, so ownership is read from the owner - role instead, which is named after the principal. ``database_ids`` are matched as given. - """ - wanted = [principal for principal in principals if principal] - names = [] - for database in self.client.postgres.list_databases(parent=branch): - database_id = (database.name or "").rsplit("/", 1)[-1] - owner = ( - (database.status.role if database.status else None) - or (database.spec.role if database.spec else None) - or "" - ) - if ( - database.name - and database_id.startswith(_RUNTIME_DATABASE_PREFIX) - and ( - database_id in database_ids - or any(owner.rsplit("/", 1)[-1].endswith(item) for item in wanted) - ) - ): - names.append(database.name) - return names + return AppIdentity(principal, branch) def roles_of(self, branch: str, principals: Collection[str]) -> list[str]: """Resource names of the branch's roles for these principals. @@ -401,12 +368,6 @@ def roles_of(self, branch: str, principals: Collection[str]) -> list[str]: if role.name and any(role.name.rsplit("/", 1)[-1].endswith(item) for item in wanted) ] - def delete_database(self, name: str) -> None: - try: - _wait(self.client.postgres.delete_database(name=name)) - except NotFound: - pass - def delete_role(self, name: str) -> None: try: _wait(self.client.postgres.delete_role(name=name)) diff --git a/integrations/agentbricks/tests/suites/README.md b/integrations/agentbricks/tests/suites/README.md index 1f65d2894..c2689d91a 100644 --- a/integrations/agentbricks/tests/suites/README.md +++ b/integrations/agentbricks/tests/suites/README.md @@ -31,10 +31,10 @@ Every resource a run creates carries its run id, and is deleted at three points: 1. **When each test ends**, by the fixture that made it: `project` deletes the test's App and its stores through the CLI, as a user would; `scratch_schema` deletes the test's schema. -2. **When the run ends**, `pytest_sessionfinish` deletes what the CLI leaves behind (synced source - folders, Runtime Stores whose App is gone, Lakebase databases and roles) and whatever a crashed - test or a failed teardown left (`harness.cleanup.delete_run_leftovers`). A failure there fails - the run. +2. **When the run ends**, `pytest_sessionfinish` deletes everything carrying the run's name prefix + that is still there (stores, Apps, synced source folders, schemas) and the Apps' Lakebase roles, + which `deployments delete` leaves 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 From 35e59557b1ea299e22dc0ba0b98d1a1f9a4577da Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Thu, 15 Oct 2026 21:40:08 +0000 Subject: [PATCH 30/45] [agentbricks] E2E suites: tests register their own cleanup commands Co-authored-by: Isaac --- .../agentbricks/tests/harness/cleanup.py | 60 +++++++++++++++---- .../agentbricks/tests/suites/README.md | 12 ++-- .../agentbricks/tests/suites/conftest.py | 49 ++++++--------- .../suites/memory_store/test_memory_store.py | 9 ++- .../session_store/test_session_store.py | 9 ++- .../tests/suites/tools/test_deploy.py | 7 +++ .../tests/suites/tools/test_grants.py | 7 +++ .../tests/suites/tools/test_selection.py | 7 +++ .../tests/suites/tracing/test_tracing.py | 10 +++- 9 files changed, 123 insertions(+), 47 deletions(-) diff --git a/integrations/agentbricks/tests/harness/cleanup.py b/integrations/agentbricks/tests/harness/cleanup.py index bde2768aa..8cdfbcb89 100644 --- a/integrations/agentbricks/tests/harness/cleanup.py +++ b/integrations/agentbricks/tests/harness/cleanup.py @@ -1,11 +1,11 @@ -"""Deleting what the E2E suites create: per App when a test ends, and per run afterwards. - -The ``project`` fixture deletes each test's App and stores through the CLI when the test ends, after -recording the App's identity with ``record_app_identity``. ``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 recorded -identities 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 +"""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. """ @@ -14,15 +14,55 @@ import dataclasses import json import pathlib -from collections.abc import Iterator +import re +from collections.abc import Callable, Iterator from typing import Any from harness.agentbricks_cli import AgentbricksCli from harness.logs import log -from harness.run import APP_PREFIX, project_prefix, scratch_schema_prefix +from harness.run import APP_PREFIX, E2EError, project_prefix, scratch_schema_prefix from harness.workspace import AppIdentity, Workspace _APPS_RECORD = "apps.jsonl" +_MISSING = re.compile(r"NOT_FOUND|RESOURCE_DOES_NOT_EXIST|does not exist|not found", re.IGNORECASE) + + +# 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 call(self, function: Callable[..., object], *args: object) -> None: + """Call ``function(*args)`` when the test ends, for what no CLI command deletes.""" + self._steps.append((getattr(function, "__name__", "cleanup"), lambda: function(*args))) + + 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) + message = (result.stderr or result.stdout).strip() + if result.returncode != 0 and not _MISSING.search(message): + raise E2EError(message) # Identities recorded by per-test cleanup diff --git a/integrations/agentbricks/tests/suites/README.md b/integrations/agentbricks/tests/suites/README.md index c2689d91a..36ceadd30 100644 --- a/integrations/agentbricks/tests/suites/README.md +++ b/integrations/agentbricks/tests/suites/README.md @@ -29,8 +29,10 @@ The run id and output directory (logs, generated projects, `apps.jsonl`) are pri Every resource a run creates carries its run id, and is deleted at three points: -1. **When each test ends**, by the fixture that made it: `project` deletes the test's App and its - stores through the CLI, as a user would; `scratch_schema` deletes the test's schema. +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, schemas) and the Apps' Lakebase roles, which `deployments delete` leaves behind (`harness.cleanup.delete_run_leftovers`). A failure @@ -67,5 +69,7 @@ rather than inside it: it is plain Python the suites import as `harness` (via `p `requires` names the fixtures it needs; tools whose fixtures skip are left out and logged. - **A suite:** add a folder with `test_*.py` files. 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. Add a `conftest.py` only for resources the CLI - under test does not create itself; create them per test and delete them in the fixture's teardown. + 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. Add a `conftest.py` only for resources the CLI under test does not create + itself; create them per test and delete them in the fixture's teardown. diff --git a/integrations/agentbricks/tests/suites/conftest.py b/integrations/agentbricks/tests/suites/conftest.py index 54ce37306..b7396d0b1 100644 --- a/integrations/agentbricks/tests/suites/conftest.py +++ b/integrations/agentbricks/tests/suites/conftest.py @@ -26,7 +26,7 @@ import pytest from harness.agentbricks_cli import AgentbricksCli -from harness.cleanup import delete_run_leftovers, record_app_identity +from harness.cleanup import Cleanup, delete_run_leftovers, record_app_identity from harness.logs import log, log_command from harness.project import Project from harness.run import RUN_ID_PATTERN, E2EError, RunConfig, run_output_dir @@ -225,23 +225,35 @@ def agentbricks_cli( 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`; deleted when it ends.""" + """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: - # Teardown deletes whatever App has this name, so it must not belong to anyone else. + # The test deletes whatever App has this name, so it must not belong to anyone else. workspace.assert_app_absent(project.app_name) yield project failed = request.node.nodeid in _failed_tests if workspace is not None: - _delete_project(agentbricks_cli, workspace, project, failed=failed) + _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( @@ -249,37 +261,14 @@ def project( ) -def _delete_project( - cli: AgentbricksCli, workspace: Workspace, project: Project, *, failed: bool -) -> None: - name = project.app_name - # Whatever deploy got as far as creating; read before anything is deleted, because the - # principal and branch come from the App and its Runtime Store. +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 = workspace.app_identity(name) if identity.service_principal or identity.branch: - # Leftover cleanup deletes what the CLI leaves behind from this record. + # 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: - # The App is deleted next, so a failing test must capture its logs while it exists. workspace.app_logs(name, cli.logs_dir / f"deploy-runtime-{name}.log") - # Deploy creates the stores agent.toml declares; memory stores are addressed by resource name. - commands = [] - if memory_store := project.declared_store("memory"): - commands.append(f"memory stores delete memory-stores/{memory_store} --yes") - if session_store := project.declared_store("session"): - commands.append(f"sessions stores delete {session_store} --yes") - if identity.service_principal: - commands.append(f"deployments delete {name} --yes") - try: - # The grant the test made, revoked while the App's principal still exists. - workspace.revoke_catalog_grant(name) - except Exception as exc: - log(f"cleanup warning | {name} | catalog grant: {exc}") - # Failures only warn: the end-of-run cleanup deletes whatever is left. - for command in commands: - result = cli.run(command, timeout=1200, check=False) - if result.returncode != 0: - log(f"cleanup warning | {name} | {(result.stderr or result.stdout).strip()}") # Deleting what the whole run left behind diff --git a/integrations/agentbricks/tests/suites/memory_store/test_memory_store.py b/integrations/agentbricks/tests/suites/memory_store/test_memory_store.py index 6827cc694..419ed2646 100644 --- a/integrations/agentbricks/tests/suites/memory_store/test_memory_store.py +++ b/integrations/agentbricks/tests/suites/memory_store/test_memory_store.py @@ -9,14 +9,21 @@ 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_memory_store_created(agentbricks_cli: AgentbricksCli, project: Project) -> None: +def test_memory_store_created( + agentbricks_cli: AgentbricksCli, project: Project, cleanup: Cleanup +) -> None: agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) project.pin_runtime() + # Registered before deploy: a deploy that fails partway can still leave them behind. + cleanup(f"memory stores delete memory-stores/{project.declared_store('memory')} --yes") + cleanup(f"sessions stores delete {project.declared_store('session')} --yes") + cleanup(f"deployments delete {project.app_name} --yes") declared = project.declared_store("memory") assert declared, "agentbricks init declared no memory store" agentbricks_cli.run(f"deploy {project.app_name} --source {project.path}", timeout=2400) diff --git a/integrations/agentbricks/tests/suites/session_store/test_session_store.py b/integrations/agentbricks/tests/suites/session_store/test_session_store.py index c3981e232..7959ee4d7 100644 --- a/integrations/agentbricks/tests/suites/session_store/test_session_store.py +++ b/integrations/agentbricks/tests/suites/session_store/test_session_store.py @@ -10,14 +10,21 @@ 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_session_store_created(agentbricks_cli: AgentbricksCli, project: Project) -> None: +def test_session_store_created( + agentbricks_cli: AgentbricksCli, project: Project, cleanup: Cleanup +) -> None: agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) project.pin_runtime() + # Registered before deploy: a deploy that fails partway can still leave them behind. + cleanup(f"memory stores delete memory-stores/{project.declared_store('memory')} --yes") + cleanup(f"sessions stores delete {project.declared_store('session')} --yes") + cleanup(f"deployments delete {project.app_name} --yes") declared = project.declared_store("session") assert declared, "agentbricks init declared no session store" agentbricks_cli.run(f"deploy {project.app_name} --source {project.path}", timeout=2400) diff --git a/integrations/agentbricks/tests/suites/tools/test_deploy.py b/integrations/agentbricks/tests/suites/tools/test_deploy.py index 860af2d51..d498e27f6 100644 --- a/integrations/agentbricks/tests/suites/tools/test_deploy.py +++ b/integrations/agentbricks/tests/suites/tools/test_deploy.py @@ -14,6 +14,7 @@ from typing import Any from harness.agentbricks_cli import AgentbricksCli +from harness.cleanup import Cleanup from harness.logs import log from harness.project import Project from harness.workspace import Workspace, granted, granted_tuples @@ -23,6 +24,7 @@ def test_deploy( agentbricks_cli: AgentbricksCli, project: Project, + cleanup: Cleanup, workspace_client: Workspace, tools: tuple[Tool, ...], authoring: str, @@ -31,6 +33,10 @@ def test_deploy( agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) project.pin_runtime() app = project.app_name + # Registered before deploy: a deploy that fails partway can still leave them behind. + cleanup(f"memory stores delete memory-stores/{project.declared_store('memory')} --yes") + cleanup(f"sessions stores delete {project.declared_store('session')} --yes") + cleanup(f"deployments delete {app} --yes") def assert_granted(tool: Tool, deployment: dict[str, Any]) -> None: if tool.grant is not None: @@ -47,6 +53,7 @@ def prepare_and_check(tool: Tool, deployment: dict[str, Any]) -> None: # The nested function is deliberately outside Agent Bricks' grants; the tool only # works once this manual grant is in place. workspace_client.grant_transitive(deployment, tool.transitive_function) + cleanup.call(workspace_client.revoke_catalog_grant, app) output = agentbricks_cli.run( f"--output json endpoint invoke {app} --path /api/invocations " f"--json {invocation_json(tool.prompt)}", diff --git a/integrations/agentbricks/tests/suites/tools/test_grants.py b/integrations/agentbricks/tests/suites/tools/test_grants.py index 99c7fb2ee..10b7b7ec8 100644 --- a/integrations/agentbricks/tests/suites/tools/test_grants.py +++ b/integrations/agentbricks/tests/suites/tools/test_grants.py @@ -5,6 +5,7 @@ import json from harness.agentbricks_cli import AgentbricksCli +from harness.cleanup import Cleanup from harness.project import Project from harness.workspace import Workspace from tool_definitions import assert_declared, uc_function_tool @@ -14,6 +15,7 @@ def test_grants( agentbricks_cli: AgentbricksCli, project: Project, + cleanup: Cleanup, workspace_client: Workspace, uc_function: UcFunction, authoring: str, @@ -23,6 +25,10 @@ def test_grants( agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) project.pin_runtime() app = project.app_name + # Registered before deploy: a deploy that fails partway can still leave them behind. + cleanup(f"memory stores delete memory-stores/{project.declared_store('memory')} --yes") + cleanup(f"sessions stores delete {project.declared_store('session')} --yes") + cleanup(f"deployments delete {app} --yes") if authoring == "cli": agentbricks_cli.run(f"tools add {tool.add_command} --source {project.path}") assert_declared(project, tool) @@ -37,5 +43,6 @@ def test_grants( with subtests.test(check="manual-grant-is-direct-and-effective"): granted = workspace_client.grant_transitive(deployment, uc_function.nested_function) + cleanup.call(workspace_client.revoke_catalog_grant, app) assert "EXECUTE" in granted["direct_privileges"] assert "EXECUTE" in granted["effective_privileges"] diff --git a/integrations/agentbricks/tests/suites/tools/test_selection.py b/integrations/agentbricks/tests/suites/tools/test_selection.py index 9533a049b..02d049b8e 100644 --- a/integrations/agentbricks/tests/suites/tools/test_selection.py +++ b/integrations/agentbricks/tests/suites/tools/test_selection.py @@ -9,6 +9,7 @@ import json from harness.agentbricks_cli import AgentbricksCli +from harness.cleanup import Cleanup from harness.project import Project from harness.workspace import Workspace from tool_definitions import ( @@ -31,6 +32,7 @@ def test_selects_uc_function_not_python_marker( agentbricks_cli: AgentbricksCli, project: Project, + cleanup: Cleanup, workspace_client: Workspace, tools: tuple[Tool, ...], uc_function: UcFunction, @@ -39,6 +41,10 @@ def test_selects_uc_function_not_python_marker( agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) project.pin_runtime() app = project.app_name + # Registered before deploy: a deploy that fails partway can still leave them behind. + cleanup(f"memory stores delete memory-stores/{project.declared_store('memory')} --yes") + cleanup(f"sessions stores delete {project.declared_store('session')} --yes") + cleanup(f"deployments delete {app} --yes") if authoring == "cli": for tool in tools: if tool.add_command: @@ -51,6 +57,7 @@ def test_selects_uc_function_not_python_marker( deployment = json.loads(agentbricks_cli.run(f"--output json deployments get {app}").stdout) # The nested function is deliberately outside Agent Bricks' grants. workspace_client.grant_transitive(deployment, uc_function.nested_function) + cleanup.call(workspace_client.revoke_catalog_grant, app) serialized = agentbricks_cli.run( f"--output json endpoint invoke {app} --path /api/invocations " f"--json {invocation_json(PROMPT)}", diff --git a/integrations/agentbricks/tests/suites/tracing/test_tracing.py b/integrations/agentbricks/tests/suites/tracing/test_tracing.py index 8ce92b0da..e7cc16fd5 100644 --- a/integrations/agentbricks/tests/suites/tracing/test_tracing.py +++ b/integrations/agentbricks/tests/suites/tracing/test_tracing.py @@ -9,15 +9,23 @@ import pytest from harness.agentbricks_cli import AgentbricksCli +from harness.cleanup import Cleanup from harness.project import Project from harness.workspace import Workspace, app_resources def test_tracing_experiment_bound( - agentbricks_cli: AgentbricksCli, project: Project, workspace_client: Workspace + agentbricks_cli: AgentbricksCli, + project: Project, + cleanup: Cleanup, + workspace_client: Workspace, ) -> None: agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) project.pin_runtime() + # Registered before deploy: a deploy that fails partway can still leave them behind. + cleanup(f"memory stores delete memory-stores/{project.declared_store('memory')} --yes") + cleanup(f"sessions stores delete {project.declared_store('session')} --yes") + cleanup(f"deployments delete {project.app_name} --yes") experiment_name = project.manifest().get("tracing", {}).get("experiment_name") agentbricks_cli.run(f"deploy {project.app_name} --source {project.path}", timeout=2400) assert experiment_name, "agent.toml binds no tracing experiment" From 92562c52671d8f0a177cb7622b9ae4d05c398e3b Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Thu, 15 Oct 2026 21:52:59 +0000 Subject: [PATCH 31/45] [agentbricks] E2E suites: check grants by invoking the tools, not by inspecting them Co-authored-by: Isaac --- .../agentbricks/tests/harness/cleanup.py | 4 - integrations/agentbricks/tests/harness/run.py | 1 - .../agentbricks/tests/harness/target.py | 1 - .../agentbricks/tests/harness/workspace.py | 135 +----------------- .../agentbricks/tests/suites/conftest.py | 6 - .../tests/suites/tools/conftest.py | 2 +- .../tests/suites/tools/test_deploy.py | 41 +++--- .../tests/suites/tools/test_grants.py | 48 ------- .../tests/suites/tools/test_selection.py | 8 +- .../tests/suites/tools/tool_definitions.py | 4 - .../tests/suites/tools/uc_objects.py | 15 +- 11 files changed, 26 insertions(+), 239 deletions(-) delete mode 100644 integrations/agentbricks/tests/suites/tools/test_grants.py diff --git a/integrations/agentbricks/tests/harness/cleanup.py b/integrations/agentbricks/tests/harness/cleanup.py index 8cdfbcb89..be32956a7 100644 --- a/integrations/agentbricks/tests/harness/cleanup.py +++ b/integrations/agentbricks/tests/harness/cleanup.py @@ -46,10 +46,6 @@ def __call__(self, command: str) -> None: """Run ``agentbricks `` when the test ends.""" self._steps.append((command, lambda: self._run(command))) - def call(self, function: Callable[..., object], *args: object) -> None: - """Call ``function(*args)`` when the test ends, for what no CLI command deletes.""" - self._steps.append((getattr(function, "__name__", "cleanup"), lambda: function(*args))) - def run_all(self) -> None: while self._steps: label, step = self._steps.pop() diff --git a/integrations/agentbricks/tests/harness/run.py b/integrations/agentbricks/tests/harness/run.py index a0fe29b06..6eb5b1a85 100644 --- a/integrations/agentbricks/tests/harness/run.py +++ b/integrations/agentbricks/tests/harness/run.py @@ -52,5 +52,4 @@ class RunConfig: output: pathlib.Path # Operator override; None lets the workspace pick the first running warehouse. warehouse_id: str | None - preprovisioned_app_catalog_access: bool wheel: pathlib.Path | None diff --git a/integrations/agentbricks/tests/harness/target.py b/integrations/agentbricks/tests/harness/target.py index ea892082c..7c4ce33b6 100644 --- a/integrations/agentbricks/tests/harness/target.py +++ b/integrations/agentbricks/tests/harness/target.py @@ -70,7 +70,6 @@ def client(self) -> Workspace: self._client = Workspace( self.profile, warehouse_id=self._run_config.warehouse_id, - preprovisioned_app_catalog_access=self._run_config.preprovisioned_app_catalog_access, ) return self._client diff --git a/integrations/agentbricks/tests/harness/workspace.py b/integrations/agentbricks/tests/harness/workspace.py index a92492685..805dc595a 100644 --- a/integrations/agentbricks/tests/harness/workspace.py +++ b/integrations/agentbricks/tests/harness/workspace.py @@ -16,12 +16,12 @@ import re import subprocess import time -from collections.abc import Collection, Sequence +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 PermissionsChange, Privilege, SecurableType, VolumeType +from databricks.sdk.service.catalog import SecurableType, VolumeType from databricks.sdk.service.sql import ExecuteStatementRequestOnWaitTimeout, State, StatementState from harness.logs import log @@ -53,16 +53,12 @@ def __init__( profile: str | None, *, warehouse_id: str | None = None, - preprovisioned_app_catalog_access: bool = False, ): self.profile = profile self.warehouse_id = warehouse_id - self.preprovisioned_app_catalog_access = preprovisioned_app_catalog_access self.client = WorkspaceClient(profile=profile) self._warehouse_started = False self._user_name: str | None = None - # App name -> (catalog, principal) of each USE CATALOG this suite granted, for cleanup. - self._catalog_grants: dict[str, tuple[str, str]] = {} # Existing resources the suite depends on @@ -231,82 +227,6 @@ def experiment_id(self, name: str) -> str: raise E2EError(f"Experiment {name!r} does not exist.") return str(experiment.experiment_id) - def transitive_state(self, deployment: dict[str, Any], function: str) -> dict[str, Any]: - """The App principal's privileges on a nested function before any manual grant.""" - principal = _principal(deployment) - direct = direct_privileges(self.client, SecurableType.FUNCTION, function, principal) - if "EXECUTE" in direct: - raise E2EError( - f"Agent Bricks granted the transitive function directly: {function} -> {direct}" - ) - effective = effective_privileges(self.client, SecurableType.FUNCTION, function, principal) - if "EXECUTE" in effective: - raise E2EError( - "The transitive function was already effective before the manual grant: " - f"{function} -> {effective}" - ) - if any( - resource.get("uc_securable", {}).get("securable_full_name") == function - for resource in app_resources(deployment) - ): - raise E2EError("The transitive function appeared among the App's resources.") - return { - "service_principal_client_id": principal, - "direct_privileges": direct, - "effective_privileges": effective, - } - - def revoke_catalog_grant(self, app_name: str) -> None: - """Revoke the USE CATALOG ``grant_transitive`` gave the App's principal, if it gave one.""" - if app_name not in self._catalog_grants: - return - catalog, principal = self._catalog_grants[app_name] - self.client.grants.update( - SecurableType.CATALOG.value, - catalog, - changes=[PermissionsChange(principal=principal, remove=[Privilege.USE_CATALOG])], - ) - del self._catalog_grants[app_name] - - def grant_transitive(self, deployment: dict[str, Any], function: str) -> dict[str, Any]: - """Grant the App principal EXECUTE on a nested function and confirm it is direct and effective.""" - app_name = str(deployment.get("name")) - principal = _principal(deployment) - catalog, schema, function_name = function.split(".") - quoted_principal = f"`{principal.replace('`', '``')}`" - statements = [] - if not self.preprovisioned_app_catalog_access: - # Recorded first so a partial failure is still revoked. The schema-level grants go - # with the test's schema, but the catalog is shared and outlives the run. - self._catalog_grants[app_name] = (catalog, principal) - statements.append(f"GRANT USE CATALOG ON CATALOG `{catalog}` TO {quoted_principal}") - statements.extend( - ( - f"GRANT USE SCHEMA ON SCHEMA `{catalog}`.`{schema}` TO {quoted_principal}", - f"GRANT EXECUTE ON FUNCTION `{catalog}`.`{schema}`.`{function_name}` " - f"TO {quoted_principal}", - ) - ) - for statement in statements: - self.sql(statement) - direct = direct_privileges(self.client, SecurableType.FUNCTION, function, principal) - effective = effective_privileges(self.client, SecurableType.FUNCTION, function, principal) - if "EXECUTE" not in effective: - raise E2EError( - "Manual EXECUTE grant did not become effective for the transitive function: " - f"{function} -> {effective}" - ) - if "EXECUTE" not in direct: - raise E2EError( - "Manual EXECUTE grant was not persisted directly for the transitive function: " - f"{function} -> {direct}" - ) - return { - "granted_at": now().isoformat(), - "direct_privileges": direct, - "effective_privileges": effective, - } - # Workarounds for the end-of-run cleanup: what `agentbricks deployments delete` leaves behind. def wait_until_app_gone(self, name: str, timeout: float = 1200) -> None: @@ -423,15 +343,6 @@ def _principal(deployment: dict[str, Any]) -> str: return str(principal) -def granted(deployment: dict[str, Any], grant: GrantTuple) -> bool: - """Whether deploy attached ``grant``, as (kind, name, type, permission), to the App.""" - return grant in granted_tuples(deployment) - - -def granted_tuples(deployment: dict[str, Any]) -> set[GrantTuple]: - return app_resource_tuples(app_resources(deployment)) - - def app_resources(deployment: dict[str, Any]) -> list[dict[str, Any]]: resources = deployment.get("resources") or [] if not isinstance(resources, list): @@ -439,27 +350,6 @@ def app_resources(deployment: dict[str, Any]) -> list[dict[str, Any]]: return resources -def app_resource_tuples(resources: Sequence[dict[str, Any]]) -> set[GrantTuple]: - actual = set() - for resource in resources: - if "uc_securable" in resource: - value = resource["uc_securable"] - actual.add( - ( - "uc_securable", - value.get("securable_full_name"), - value.get("securable_type"), - value.get("permission"), - ) - ) - elif "genie_space" in resource: - value = resource["genie_space"] - actual.add( - ("genie_space", value.get("space_id"), "GENIE_SPACE", value.get("permission")) - ) - return actual - - def effective_privileges( client: WorkspaceClient, securable_type: SecurableType, full_name: str, principal: str ) -> list[str]: @@ -484,24 +374,3 @@ def effective_privileges( page_token = response.next_page_token if not page_token: return sorted(privileges) - - -def direct_privileges( - client: WorkspaceClient, securable_type: SecurableType, full_name: str, principal: str -) -> list[str]: - privileges: set[str] = set() - page_token: str | None = None - while True: - response = client.grants.get( - securable_type.value, - full_name, - max_results=0, - principal=principal, - **({"page_token": page_token} if page_token else {}), - ) - for assignment in response.privilege_assignments or (): - if assignment.principal == principal: - privileges.update(privilege.value for privilege in assignment.privileges or ()) - page_token = response.next_page_token - if not page_token: - return sorted(privileges) diff --git a/integrations/agentbricks/tests/suites/conftest.py b/integrations/agentbricks/tests/suites/conftest.py index b7396d0b1..1b45bafc3 100644 --- a/integrations/agentbricks/tests/suites/conftest.py +++ b/integrations/agentbricks/tests/suites/conftest.py @@ -79,11 +79,6 @@ def pytest_addoption(parser: pytest.Parser) -> None: "--genie-space-id", help="Existing Genie space for the Genie tool; without one that tool is skipped.", ) - group.addoption( - "--preprovisioned-app-catalog-access", - action="store_true", - help="Skip per-App USE CATALOG grants because catalog access is pre-provisioned.", - ) group.addoption( "--run-id", help="Names every resource of the run, so --cleanup-only can delete what a killed run left " @@ -143,7 +138,6 @@ def _base_run_config(config: pytest.Config) -> RunConfig: run_id=run_id, output=output, warehouse_id=option("warehouse_id"), - preprovisioned_app_catalog_access=bool(option("preprovisioned_app_catalog_access")), wheel=option("wheel"), ) config.stash[_RUN_KEY] = run diff --git a/integrations/agentbricks/tests/suites/tools/conftest.py b/integrations/agentbricks/tests/suites/tools/conftest.py index 518c91f05..ab26b5dc6 100644 --- a/integrations/agentbricks/tests/suites/tools/conftest.py +++ b/integrations/agentbricks/tests/suites/tools/conftest.py @@ -45,7 +45,7 @@ def scratch_schema( @pytest.fixture def uc_function(scratch_schema: str, target_workspace: TargetWorkspace) -> UcFunction: - """The marker functions and volume in ``scratch_schema``.""" + """The marker function and volume in ``scratch_schema``.""" return create_uc_function(require_workspace(target_workspace, "a UC function"), scratch_schema) diff --git a/integrations/agentbricks/tests/suites/tools/test_deploy.py b/integrations/agentbricks/tests/suites/tools/test_deploy.py index d498e27f6..2b77747d0 100644 --- a/integrations/agentbricks/tests/suites/tools/test_deploy.py +++ b/integrations/agentbricks/tests/suites/tools/test_deploy.py @@ -1,9 +1,9 @@ -"""Each tool works on a deployed App, and its UC grant is applied, per authoring path. +"""Each tool works on a deployed App, per authoring path: invoking it returns its marker. -CLI authoring is incremental: add one tool, redeploy the same App, verify it, then add the next. -After every redeploy the grants of all tools added so far must still hold, which is the idempotency -check. Direct authoring writes every tool into agent.toml and deploys once. A subtest per tool -reports each tool independently and the sequence continues past a failure. +CLI authoring is incremental: add one tool, redeploy the same App, invoke it, then add the next. +After the last redeploy every tool is invoked again, so a redeploy that broke an earlier tool, e.g. +by dropping its grant, fails. Direct authoring writes every tool into agent.toml and deploys once. +A subtest per tool reports each tool independently and the sequence continues past a failure. `agentbricks deploy` returns once the App is running its new source, so the next step can use it. """ @@ -17,7 +17,7 @@ from harness.cleanup import Cleanup from harness.logs import log from harness.project import Project -from harness.workspace import Workspace, granted, granted_tuples +from harness.workspace import Workspace from tool_definitions import Tool, assert_declared, invocation_json @@ -38,22 +38,12 @@ def test_deploy( cleanup(f"sessions stores delete {project.declared_store('session')} --yes") cleanup(f"deployments delete {app} --yes") - def assert_granted(tool: Tool, deployment: dict[str, Any]) -> None: - if tool.grant is not None: - assert granted(deployment, tool.grant), ( - f"{tool.name}: {tool.grant} not granted to {app}: " - f"{sorted(granted_tuples(deployment))}" - ) - - def prepare_and_check(tool: Tool, deployment: dict[str, Any]) -> None: + def invoke_and_check(tool: Tool, deployment: dict[str, Any]) -> None: if tool.grant is not None: + # Deploy adds the grant before returning, but UC can apply it a little later, and a + # tool called in between is denied inside a successful response. waited = workspace_client.wait_until_effective(deployment, tool.grant) log(f"{tool.name}: {tool.grant[3]} effective for {app} after {waited:.0f}s") - if tool.transitive_function: - # The nested function is deliberately outside Agent Bricks' grants; the tool only - # works once this manual grant is in place. - workspace_client.grant_transitive(deployment, tool.transitive_function) - cleanup.call(workspace_client.revoke_catalog_grant, app) output = agentbricks_cli.run( f"--output json endpoint invoke {app} --path /api/invocations " f"--json {invocation_json(tool.prompt)}", @@ -63,7 +53,7 @@ def prepare_and_check(tool: Tool, deployment: dict[str, Any]) -> None: tool.check(output) if authoring == "cli": - for index, tool in enumerate(tools): + for tool in tools: with subtests.test(tool=tool.name): if tool.add_command: agentbricks_cli.run(f"tools add {tool.add_command} --source {project.path}") @@ -73,14 +63,15 @@ def prepare_and_check(tool: Tool, deployment: dict[str, Any]) -> None: deployment = json.loads( agentbricks_cli.run(f"--output json deployments get {app}").stdout ) - prepare_and_check(tool, deployment) - for added in tools[: index + 1]: - assert_granted(added, deployment) + invoke_and_check(tool, deployment) + deployment = json.loads(agentbricks_cli.run(f"--output json deployments get {app}").stdout) + for tool in tools: + with subtests.test(tool=tool.name, after="last-deploy"): + invoke_and_check(tool, deployment) elif authoring == "direct": project.write_manifest(tools) agentbricks_cli.run(f"deploy {app} --source {project.path}", timeout=2400) deployment = json.loads(agentbricks_cli.run(f"--output json deployments get {app}").stdout) for tool in tools: with subtests.test(tool=tool.name): - prepare_and_check(tool, deployment) - assert_granted(tool, deployment) + invoke_and_check(tool, deployment) diff --git a/integrations/agentbricks/tests/suites/tools/test_grants.py b/integrations/agentbricks/tests/suites/tools/test_grants.py deleted file mode 100644 index 10b7b7ec8..000000000 --- a/integrations/agentbricks/tests/suites/tools/test_grants.py +++ /dev/null @@ -1,48 +0,0 @@ -"""Deploy grants only what a tool binds: a function the bound UC function calls needs a manual grant.""" - -from __future__ import annotations - -import json - -from harness.agentbricks_cli import AgentbricksCli -from harness.cleanup import Cleanup -from harness.project import Project -from harness.workspace import Workspace -from tool_definitions import assert_declared, uc_function_tool -from uc_objects import UcFunction - - -def test_grants( - agentbricks_cli: AgentbricksCli, - project: Project, - cleanup: Cleanup, - workspace_client: Workspace, - uc_function: UcFunction, - authoring: str, - subtests, -) -> None: - tool = uc_function_tool(uc_function) - agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) - project.pin_runtime() - app = project.app_name - # Registered before deploy: a deploy that fails partway can still leave them behind. - cleanup(f"memory stores delete memory-stores/{project.declared_store('memory')} --yes") - cleanup(f"sessions stores delete {project.declared_store('session')} --yes") - cleanup(f"deployments delete {app} --yes") - if authoring == "cli": - agentbricks_cli.run(f"tools add {tool.add_command} --source {project.path}") - assert_declared(project, tool) - else: - project.write_manifest([tool]) - agentbricks_cli.run(f"deploy {app} --source {project.path}", timeout=2400) - deployment = json.loads(agentbricks_cli.run(f"--output json deployments get {app}").stdout) - - with subtests.test(check="transitive-function-not-auto-granted"): - # Raises if the hidden nested function is already granted or listed as an App resource. - workspace_client.transitive_state(deployment, uc_function.nested_function) - - with subtests.test(check="manual-grant-is-direct-and-effective"): - granted = workspace_client.grant_transitive(deployment, uc_function.nested_function) - cleanup.call(workspace_client.revoke_catalog_grant, app) - assert "EXECUTE" in granted["direct_privileges"] - assert "EXECUTE" in granted["effective_privileges"] diff --git a/integrations/agentbricks/tests/suites/tools/test_selection.py b/integrations/agentbricks/tests/suites/tools/test_selection.py index 02d049b8e..1057dc63a 100644 --- a/integrations/agentbricks/tests/suites/tools/test_selection.py +++ b/integrations/agentbricks/tests/suites/tools/test_selection.py @@ -19,6 +19,7 @@ assert_declared, excerpt, invocation_json, + uc_function_tool, ) from uc_objects import UcFunction @@ -55,9 +56,10 @@ def test_selects_uc_function_not_python_marker( project.write_manifest(tools) agentbricks_cli.run(f"deploy {app} --source {project.path}", timeout=2400) deployment = json.loads(agentbricks_cli.run(f"--output json deployments get {app}").stdout) - # The nested function is deliberately outside Agent Bricks' grants. - workspace_client.grant_transitive(deployment, uc_function.nested_function) - cleanup.call(workspace_client.revoke_catalog_grant, app) + # UC can apply deploy's grant a little after deploy returns; a call in between is denied. + grant = uc_function_tool(uc_function).grant + assert grant is not None + workspace_client.wait_until_effective(deployment, grant) serialized = agentbricks_cli.run( f"--output json endpoint invoke {app} --path /api/invocations " f"--json {invocation_json(PROMPT)}", diff --git a/integrations/agentbricks/tests/suites/tools/tool_definitions.py b/integrations/agentbricks/tests/suites/tools/tool_definitions.py index 64398ccb9..9ff23ba49 100644 --- a/integrations/agentbricks/tests/suites/tools/tool_definitions.py +++ b/integrations/agentbricks/tests/suites/tools/tool_definitions.py @@ -44,9 +44,6 @@ class Tool: app_auth: bool # The App resource a deploy grants for it; None when it needs no UC grant. grant: GrantTuple | None = None - # A function the tool calls that Agent Bricks deliberately does not grant; it works only after - # a manual grant. - transitive_function: str | None = None @dataclasses.dataclass(frozen=True) @@ -168,7 +165,6 @@ def uc_function_tool(uc_function: UcFunction) -> Tool: tool_id="agentbricks_uc_marker", app_auth=False, grant=("uc_securable", uc_function.function, "FUNCTION", "EXECUTE"), - transitive_function=uc_function.nested_function, ) diff --git a/integrations/agentbricks/tests/suites/tools/uc_objects.py b/integrations/agentbricks/tests/suites/tools/uc_objects.py index 2d732ef7d..21c40deb8 100644 --- a/integrations/agentbricks/tests/suites/tools/uc_objects.py +++ b/integrations/agentbricks/tests/suites/tools/uc_objects.py @@ -8,7 +8,6 @@ from harness.run import E2EError from harness.workspace import Workspace -NESTED_FUNCTION = "ab_nested_marker" FUNCTION = "uc_marker" VOLUME = "ab_volume" # The MCP server exposes a function as catalog__schema__function, capped at 64 characters. @@ -20,15 +19,13 @@ class UcFunction: """The scratch UC objects the tools use.""" function: str - # Called by ``function`` and deliberately never declared in agent.toml. - nested_function: str volume: str volume_file_path: str volume_marker: str def create_uc_function(workspace: Workspace, scratch_schema: str) -> UcFunction: - """Create the marker functions and volume in ``scratch_schema``, whose deletion removes them.""" + """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: @@ -36,18 +33,11 @@ def create_uc_function(workspace: Workspace, scratch_schema: str) -> UcFunction: f"The UC function's MCP tool name would exceed {MAX_TOOL_NAME_LENGTH} characters: " f"{exposed_tool_name!r}. Use a shorter catalog." ) - # The nested function first: the outer one resolves it when created. - workspace.sql( - f"CREATE FUNCTION `{catalog}`.`{schema}`.`{NESTED_FUNCTION}`" - "(value STRING) RETURNS STRING " - "COMMENT 'Transitive Agent Bricks E2E marker; never declared in agent.toml' " - "RETURN concat('AGENTBRICKS_UC_OK:', value)" - ) workspace.sql( f"CREATE FUNCTION `{catalog}`.`{schema}`.`{FUNCTION}`" "(value STRING) RETURNS STRING " "COMMENT 'Deterministic Agent Bricks E2E marker tool' " - f"RETURN `{catalog}`.`{schema}`.`{NESTED_FUNCTION}`(value)" + "RETURN concat('AGENTBRICKS_UC_OK:', value)" ) workspace.create_volume(catalog, schema, VOLUME) marker = f"AGENTBRICKS_VOLUME_{uuid.uuid4().hex}" @@ -55,7 +45,6 @@ def create_uc_function(workspace: Workspace, scratch_schema: str) -> UcFunction: workspace.upload(file_path, marker.encode()) return UcFunction( function=f"{scratch_schema}.{FUNCTION}", - nested_function=f"{scratch_schema}.{NESTED_FUNCTION}", volume=f"{scratch_schema}.{VOLUME}", volume_file_path=file_path, volume_marker=marker, From 844cbeef11cb1e217c209e738bb05904607a15dd Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Thu, 15 Oct 2026 22:10:02 +0000 Subject: [PATCH 32/45] [agentbricks] E2E suites: CLI-authored deploy adds every tool, then deploys once Co-authored-by: Isaac --- .../tests/suites/tools/test_deploy.py | 35 +++++++------------ .../tests/suites/tools/test_selection.py | 5 +-- 2 files changed, 15 insertions(+), 25 deletions(-) diff --git a/integrations/agentbricks/tests/suites/tools/test_deploy.py b/integrations/agentbricks/tests/suites/tools/test_deploy.py index 2b77747d0..5aef992bb 100644 --- a/integrations/agentbricks/tests/suites/tools/test_deploy.py +++ b/integrations/agentbricks/tests/suites/tools/test_deploy.py @@ -1,9 +1,8 @@ """Each tool works on a deployed App, per authoring path: invoking it returns its marker. -CLI authoring is incremental: add one tool, redeploy the same App, invoke it, then add the next. -After the last redeploy every tool is invoked again, so a redeploy that broke an earlier tool, e.g. -by dropping its grant, fails. Direct authoring writes every tool into agent.toml and deploys once. -A subtest per tool reports each tool independently and the sequence continues past a failure. +CLI authoring adds every tool with `agentbricks tools add`; direct authoring writes them into +agent.toml. Either way the App is deployed once, and a subtest per tool invokes it, so each tool is +reported independently. `agentbricks deploy` returns once the App is running its new source, so the next step can use it. """ @@ -54,24 +53,14 @@ def invoke_and_check(tool: Tool, deployment: dict[str, Any]) -> None: if authoring == "cli": for tool in tools: - with subtests.test(tool=tool.name): - if tool.add_command: - agentbricks_cli.run(f"tools add {tool.add_command} --source {project.path}") - project.write_files(tool.files) - assert_declared(project, tool) - agentbricks_cli.run(f"deploy {app} --source {project.path}", timeout=2400) - deployment = json.loads( - agentbricks_cli.run(f"--output json deployments get {app}").stdout - ) - invoke_and_check(tool, deployment) - deployment = json.loads(agentbricks_cli.run(f"--output json deployments get {app}").stdout) - for tool in tools: - with subtests.test(tool=tool.name, after="last-deploy"): - invoke_and_check(tool, deployment) + if tool.add_command: + agentbricks_cli.run(f"tools add {tool.add_command} --source {project.path}") + project.write_files(tool.files) + assert_declared(project, tool) elif authoring == "direct": project.write_manifest(tools) - agentbricks_cli.run(f"deploy {app} --source {project.path}", timeout=2400) - deployment = json.loads(agentbricks_cli.run(f"--output json deployments get {app}").stdout) - for tool in tools: - with subtests.test(tool=tool.name): - invoke_and_check(tool, deployment) + agentbricks_cli.run(f"deploy {app} --source {project.path}", timeout=2400) + deployment = json.loads(agentbricks_cli.run(f"--output json deployments get {app}").stdout) + for tool in tools: + with subtests.test(tool=tool.name): + invoke_and_check(tool, deployment) diff --git a/integrations/agentbricks/tests/suites/tools/test_selection.py b/integrations/agentbricks/tests/suites/tools/test_selection.py index 1057dc63a..7b4e168b4 100644 --- a/integrations/agentbricks/tests/suites/tools/test_selection.py +++ b/integrations/agentbricks/tests/suites/tools/test_selection.py @@ -1,7 +1,8 @@ """MULTI-TOOL SELECTION check (deploy-only): with every tool bound, the agent picks the one asked for. -The incremental tests cannot catch a regression where binding several tools makes the model, or -the tool surface, route to the wrong one, so this is the one test that binds them all at once. +`test_deploy` names the tool in each prompt; this prompt names only what the tool does, so it +catches a regression where binding several tools makes the model, or the tool surface, route to the +wrong one. """ from __future__ import annotations From 58b991a0385a836520982144fcbe0a54c186eae3 Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Thu, 15 Oct 2026 22:12:32 +0000 Subject: [PATCH 33/45] [agentbricks] E2E suites: CLI-authored dev adds every tool, then starts dev once Co-authored-by: Isaac --- .../tests/suites/tools/test_dev.py | 39 ++++++++----------- 1 file changed, 16 insertions(+), 23 deletions(-) diff --git a/integrations/agentbricks/tests/suites/tools/test_dev.py b/integrations/agentbricks/tests/suites/tools/test_dev.py index 38e9679fc..4526a76fd 100644 --- a/integrations/agentbricks/tests/suites/tools/test_dev.py +++ b/integrations/agentbricks/tests/suites/tools/test_dev.py @@ -1,9 +1,9 @@ """Each tool works in `agentbricks dev`, per authoring path. -CLI authoring is incremental: add one tool, restart dev, verify it, then add the next on the same -project. Direct authoring writes every tool into agent.toml and runs dev once. A subtest per tool -reports each tool independently and the sequence continues past a failure. Every check calls the -dev server with `agentbricks endpoint invoke`, the command the CLI tells users to run. +CLI authoring adds every tool with `agentbricks tools add`; direct authoring writes them into +agent.toml. Either way dev runs once, and a subtest per tool invokes it, so each tool is reported +independently. Every check calls the dev server with `agentbricks endpoint invoke`, the command the +CLI tells users to run. 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. @@ -41,24 +41,17 @@ def check(tool: Tool, port: int) -> None: if authoring == "cli": for tool in tools: - with subtests.test(tool=tool.name): - if tool.add_command: - agentbricks_cli.run(f"tools add {tool.add_command} --source {project.path}") - project.write_files(tool.files) - assert_declared(project, tool) - 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) - check(tool, port) + if tool.add_command: + agentbricks_cli.run(f"tools add {tool.add_command} --source {project.path}") + project.write_files(tool.files) + assert_declared(project, tool) elif authoring == "direct": project.write_manifest(tools) - 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): - check(tool, port) + 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): + check(tool, port) From 4cd8e084d0ada8613e926820caffa4e904a6563e Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Thu, 15 Oct 2026 22:21:25 +0000 Subject: [PATCH 34/45] [agentbricks] E2E suites: read the App principal through the CLI; name the UC tool by purpose Co-authored-by: Isaac --- .../agentbricks/tests/harness/cleanup.py | 17 ++++++++++++- .../agentbricks/tests/harness/workspace.py | 24 +++++-------------- .../agentbricks/tests/suites/conftest.py | 4 ++-- .../tests/suites/tools/tool_definitions.py | 4 ++-- 4 files changed, 26 insertions(+), 23 deletions(-) diff --git a/integrations/agentbricks/tests/harness/cleanup.py b/integrations/agentbricks/tests/harness/cleanup.py index be32956a7..b95418fe5 100644 --- a/integrations/agentbricks/tests/harness/cleanup.py +++ b/integrations/agentbricks/tests/harness/cleanup.py @@ -64,6 +64,21 @@ def _run(self, command: str) -> None: # 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 _MISSING.search(message): + 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" @@ -169,7 +184,7 @@ def _delete_leaked_apps( continue log(f"deleting leaked App {name}") # The principal stops resolving once the App is gone, and the later steps need it. - record_app_identity(output, name, workspace.app_identity(name)) + 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()}") diff --git a/integrations/agentbricks/tests/harness/workspace.py b/integrations/agentbricks/tests/harness/workspace.py index 805dc595a..7913e491b 100644 --- a/integrations/agentbricks/tests/harness/workspace.py +++ b/integrations/agentbricks/tests/harness/workspace.py @@ -161,9 +161,6 @@ def create_volume(self, catalog: str, schema: str, name: str) -> None: # Apps - def app(self, name: str) -> dict[str, Any]: - return self.client.apps.get(name).as_dict() - def assert_app_absent(self, name: str) -> None: try: self.client.apps.get(name) @@ -255,26 +252,17 @@ def runtime_store(self, app_name: str) -> dict[str, Any]: self.client.api_client.do("GET", f"/api/2.0/agents/runtime-stores/{app_name}"), ) - def app_identity(self, app_name: str) -> AppIdentity: - """The App's principal and Lakebase location as far as they still exist; never raises.""" - principal: str | None = None - branch: str | None = None - try: - principal = self.app(app_name).get("service_principal_client_id") - except NotFound: - pass - except Exception as exc: - log(f"cleanup warning | principal lookup for {app_name} | {exc}") + 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", {}) - candidate = lakebase.get("branch") - if isinstance(candidate, str) and _BRANCH.fullmatch(candidate): - branch = candidate except NotFound: - pass + return None except Exception as exc: log(f"cleanup warning | Runtime Store lookup for {app_name} | {exc}") - return AppIdentity(principal, branch) + 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. diff --git a/integrations/agentbricks/tests/suites/conftest.py b/integrations/agentbricks/tests/suites/conftest.py index 1b45bafc3..f8963686f 100644 --- a/integrations/agentbricks/tests/suites/conftest.py +++ b/integrations/agentbricks/tests/suites/conftest.py @@ -26,7 +26,7 @@ import pytest from harness.agentbricks_cli import AgentbricksCli -from harness.cleanup import Cleanup, delete_run_leftovers, record_app_identity +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 RUN_ID_PATTERN, E2EError, RunConfig, run_output_dir @@ -257,7 +257,7 @@ def project( 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 = workspace.app_identity(name) + 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) diff --git a/integrations/agentbricks/tests/suites/tools/tool_definitions.py b/integrations/agentbricks/tests/suites/tools/tool_definitions.py index 9ff23ba49..8b5b80262 100644 --- a/integrations/agentbricks/tests/suites/tools/tool_definitions.py +++ b/integrations/agentbricks/tests/suites/tools/tool_definitions.py @@ -152,8 +152,8 @@ def uc_function_tool(uc_function: UcFunction) -> Tool: return Tool( name="uc_function", prompt=( - f"You must call the tool named {uc_function.function.replace('.', '__')} with value " - "'matrix'. Return the called tool's exact result." + "You must call the Unity Catalog function tool that computes the Agent Bricks marker, " + "with value 'matrix'. Return the called tool's exact result." ), check=contains(UC_MARKER), add_command=f"uc-function {uc_function.function} --name agentbricks_uc_marker", From 82d87c671b6b1702ee2e9a06421fe2706af729e1 Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Thu, 15 Oct 2026 23:26:47 +0000 Subject: [PATCH 35/45] [agentbricks] E2E suites: check App existence with deployments get, not the Apps API Co-authored-by: Isaac --- .../agentbricks/tests/harness/cleanup.py | 42 +++++++++++++++---- .../agentbricks/tests/harness/workspace.py | 31 +------------- .../agentbricks/tests/suites/conftest.py | 11 ++++- 3 files changed, 45 insertions(+), 39 deletions(-) diff --git a/integrations/agentbricks/tests/harness/cleanup.py b/integrations/agentbricks/tests/harness/cleanup.py index b95418fe5..c9fbcdd75 100644 --- a/integrations/agentbricks/tests/harness/cleanup.py +++ b/integrations/agentbricks/tests/harness/cleanup.py @@ -15,12 +15,13 @@ import json import pathlib import re +import time from collections.abc import Callable, Iterator from typing import Any from harness.agentbricks_cli import AgentbricksCli from harness.logs import log -from harness.run import APP_PREFIX, E2EError, project_prefix, scratch_schema_prefix +from harness.run import APP_PREFIX, E2EError, now, project_prefix, scratch_schema_prefix from harness.workspace import AppIdentity, Workspace _APPS_RECORD = "apps.jsonl" @@ -64,6 +65,17 @@ def _run(self, command: str) -> None: # Identities recorded by per-test cleanup +def app_exists(cli: AgentbricksCli, name: str) -> bool: + """Whether `agentbricks deployments get` finds the App; raises when it can't tell.""" + result = cli.run(f"--output json deployments get {name}", check=False) + if result.returncode == 0: + return True + message = (result.stderr or result.stdout).strip() + if _MISSING.search(message): + return False + raise E2EError(f"Could not tell whether App {name} exists: {message}") + + 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 @@ -178,26 +190,42 @@ def _delete_leaked_apps( 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 app in workspace.client.apps.list(): - if app.name and app.name.startswith(prefix): - try: - workspace.wait_until_app_gone(app.name) - except Exception as exc: - failures.append(f"App {app.name}: {exc}") + 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 app_exists(cli, name): + 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()) diff --git a/integrations/agentbricks/tests/harness/workspace.py b/integrations/agentbricks/tests/harness/workspace.py index 7913e491b..2d58da0de 100644 --- a/integrations/agentbricks/tests/harness/workspace.py +++ b/integrations/agentbricks/tests/harness/workspace.py @@ -25,7 +25,7 @@ from databricks.sdk.service.sql import ExecuteStatementRequestOnWaitTimeout, State, StatementState from harness.logs import log -from harness.run import E2EError, now +from harness.run import E2EError GrantTuple = tuple[str, str, str, str] @@ -161,15 +161,6 @@ def create_volume(self, catalog: str, schema: str, name: str) -> None: # Apps - def assert_app_absent(self, name: str) -> None: - try: - self.client.apps.get(name) - except NotFound: - return - except DatabricksError as exc: - raise E2EError(f"Could not verify App {name} is absent: {exc}") from exc - raise E2EError(f"App {name} already exists; refusing to deploy over it.") - 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"] @@ -226,26 +217,6 @@ def experiment_id(self, name: str) -> str: # Workarounds for the end-of-run cleanup: what `agentbricks deployments delete` leaves behind. - def wait_until_app_gone(self, name: str, timeout: float = 1200) -> None: - """Block until the App is gone; `apps delete` returns while it is still deleting.""" - started = time.monotonic() - next_tick = 0.0 - while True: - try: - self.client.apps.get(name) - except NotFound: - log(f"tick {now():%H:%M} | delete-{name} | absent") - return - except DatabricksError as exc: - raise E2EError(f"Could not verify deletion of App {name!r}: {exc}") from exc - 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 runtime_store(self, app_name: str) -> dict[str, Any]: return cast( dict[str, Any], diff --git a/integrations/agentbricks/tests/suites/conftest.py b/integrations/agentbricks/tests/suites/conftest.py index f8963686f..3007f18d1 100644 --- a/integrations/agentbricks/tests/suites/conftest.py +++ b/integrations/agentbricks/tests/suites/conftest.py @@ -26,7 +26,13 @@ import pytest from harness.agentbricks_cli import AgentbricksCli -from harness.cleanup import Cleanup, app_identity, delete_run_leftovers, record_app_identity +from harness.cleanup import ( + Cleanup, + app_exists, + app_identity, + delete_run_leftovers, + record_app_identity, +) from harness.logs import log, log_command from harness.project import Project from harness.run import RUN_ID_PATTERN, E2EError, RunConfig, run_output_dir @@ -243,7 +249,8 @@ def project( 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. - workspace.assert_app_absent(project.app_name) + if app_exists(agentbricks_cli, project.app_name): + raise E2EError(f"App {project.app_name} already exists; refusing to deploy over it.") yield project failed = request.node.nodeid in _failed_tests if workspace is not None: From cadc1b91723042de684b6435a75377df9e7b8155 Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Thu, 15 Oct 2026 23:28:27 +0000 Subject: [PATCH 36/45] [agentbricks] E2E suites: run deployments get directly instead of through app_exists Co-authored-by: Isaac --- .../tests/harness/agentbricks_cli.py | 8 +++++++ .../agentbricks/tests/harness/cleanup.py | 24 ++++--------------- .../agentbricks/tests/suites/conftest.py | 10 +++++--- 3 files changed, 20 insertions(+), 22 deletions(-) diff --git a/integrations/agentbricks/tests/harness/agentbricks_cli.py b/integrations/agentbricks/tests/harness/agentbricks_cli.py index 69e9ecf04..64cb30078 100644 --- a/integrations/agentbricks/tests/harness/agentbricks_cli.py +++ b/integrations/agentbricks/tests/harness/agentbricks_cli.py @@ -31,6 +31,9 @@ # `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 @@ -221,6 +224,11 @@ def terminate_process_group(process: subprocess.Popen[str], *, grace: float = 20 pass +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)) + + def free_port() -> int: with socket.socket() as probe: probe.bind(("127.0.0.1", 0)) diff --git a/integrations/agentbricks/tests/harness/cleanup.py b/integrations/agentbricks/tests/harness/cleanup.py index c9fbcdd75..61e1c53d9 100644 --- a/integrations/agentbricks/tests/harness/cleanup.py +++ b/integrations/agentbricks/tests/harness/cleanup.py @@ -14,18 +14,16 @@ import dataclasses import json import pathlib -import re import time from collections.abc import Callable, Iterator from typing import Any -from harness.agentbricks_cli import AgentbricksCli +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" -_MISSING = re.compile(r"NOT_FOUND|RESOURCE_DOES_NOT_EXIST|does not exist|not found", re.IGNORECASE) # One test @@ -57,25 +55,13 @@ def run_all(self) -> None: def _run(self, command: str) -> None: result = self._cli.run(command, timeout=1200, check=False) - message = (result.stderr or result.stdout).strip() - if result.returncode != 0 and not _MISSING.search(message): - raise E2EError(message) + 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_exists(cli: AgentbricksCli, name: str) -> bool: - """Whether `agentbricks deployments get` finds the App; raises when it can't tell.""" - result = cli.run(f"--output json deployments get {name}", check=False) - if result.returncode == 0: - return True - message = (result.stderr or result.stdout).strip() - if _MISSING.search(message): - return False - raise E2EError(f"Could not tell whether App {name} exists: {message}") - - 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 @@ -86,7 +72,7 @@ def app_identity(cli: AgentbricksCli, workspace: Workspace, name: str) -> AppIde principal = json.loads(result.stdout).get("service_principal_client_id") except (ValueError, AttributeError): log(f"cleanup warning | principal lookup for {name} | {message}") - elif not _MISSING.search(message): + elif not not_found(result): log(f"cleanup warning | principal lookup for {name} | {message}") return AppIdentity(principal, workspace.lakebase_branch(name)) @@ -216,7 +202,7 @@ def _wait_until_app_gone(cli: AgentbricksCli, name: str, timeout: float = 1200) """Block until the App is gone; `deployments delete` returns while it is still deleting.""" started = time.monotonic() next_tick = 0.0 - while app_exists(cli, name): + 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.") diff --git a/integrations/agentbricks/tests/suites/conftest.py b/integrations/agentbricks/tests/suites/conftest.py index 3007f18d1..729943b41 100644 --- a/integrations/agentbricks/tests/suites/conftest.py +++ b/integrations/agentbricks/tests/suites/conftest.py @@ -25,10 +25,9 @@ from typing import Any import pytest -from harness.agentbricks_cli import AgentbricksCli +from harness.agentbricks_cli import AgentbricksCli, not_found from harness.cleanup import ( Cleanup, - app_exists, app_identity, delete_run_leftovers, record_app_identity, @@ -249,8 +248,13 @@ def project( 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. - if app_exists(agentbricks_cli, project.app_name): + 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: From 0160964406cb16b75021f454fa3170f4160f71d7 Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Fri, 16 Oct 2026 00:05:46 +0000 Subject: [PATCH 37/45] [agentbricks] E2E suites: one happy path per command group, in flat files Co-authored-by: Isaac --- .../agentbricks/tests/harness/cleanup.py | 22 ++++- .../agentbricks/tests/harness/project.py | 25 +----- integrations/agentbricks/tests/harness/run.py | 1 - .../agentbricks/tests/harness/workspace.py | 16 ++++ .../agentbricks/tests/suites/README.md | 22 ++--- .../agentbricks/tests/suites/conftest.py | 83 +++++++++++++++++-- .../suites/memory_store/test_memory_store.py | 40 --------- .../session_store/test_session_store.py | 38 --------- .../agentbricks/tests/suites/test_deploy.py | 59 +++++++++++++ .../agentbricks/tests/suites/test_dev.py | 42 ++++++++++ .../agentbricks/tests/suites/test_memory.py | 29 +++++++ .../agentbricks/tests/suites/test_sessions.py | 27 ++++++ .../agentbricks/tests/suites/test_tools.py | 23 +++++ .../agentbricks/tests/suites/test_tracing.py | 45 ++++++++++ .../suites/{tools => }/tool_definitions.py | 80 +++++++++--------- .../tests/suites/tools/conftest.py | 77 ----------------- .../tests/suites/tools/test_deploy.py | 66 --------------- .../tests/suites/tools/test_dev.py | 57 ------------- .../tests/suites/tools/test_selection.py | 74 ----------------- .../tests/suites/tools/test_tools_add.py | 34 -------- .../tests/suites/tracing/test_tracing.py | 47 ----------- .../tests/suites/{tools => }/uc_objects.py | 0 22 files changed, 390 insertions(+), 517 deletions(-) delete mode 100644 integrations/agentbricks/tests/suites/memory_store/test_memory_store.py delete mode 100644 integrations/agentbricks/tests/suites/session_store/test_session_store.py create mode 100644 integrations/agentbricks/tests/suites/test_deploy.py create mode 100644 integrations/agentbricks/tests/suites/test_dev.py create mode 100644 integrations/agentbricks/tests/suites/test_memory.py create mode 100644 integrations/agentbricks/tests/suites/test_sessions.py create mode 100644 integrations/agentbricks/tests/suites/test_tools.py create mode 100644 integrations/agentbricks/tests/suites/test_tracing.py rename integrations/agentbricks/tests/suites/{tools => }/tool_definitions.py (75%) delete mode 100644 integrations/agentbricks/tests/suites/tools/conftest.py delete mode 100644 integrations/agentbricks/tests/suites/tools/test_deploy.py delete mode 100644 integrations/agentbricks/tests/suites/tools/test_dev.py delete mode 100644 integrations/agentbricks/tests/suites/tools/test_selection.py delete mode 100644 integrations/agentbricks/tests/suites/tools/test_tools_add.py delete mode 100644 integrations/agentbricks/tests/suites/tracing/test_tracing.py rename integrations/agentbricks/tests/suites/{tools => }/uc_objects.py (100%) diff --git a/integrations/agentbricks/tests/harness/cleanup.py b/integrations/agentbricks/tests/harness/cleanup.py index 61e1c53d9..8df2acb02 100644 --- a/integrations/agentbricks/tests/harness/cleanup.py +++ b/integrations/agentbricks/tests/harness/cleanup.py @@ -117,6 +117,7 @@ def delete_run_leftovers( ("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: @@ -192,13 +193,13 @@ def _delete_leaked_apps( failures: list[str] = [] for name in sorted(set(leaked) | set(recorded_app_identities(output))): try: - _wait_until_app_gone(cli, name) + 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: +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 @@ -249,6 +250,23 @@ def _delete_leaked_deployment_sources(workspace: Workspace, run_id: str) -> list 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] = [] diff --git a/integrations/agentbricks/tests/harness/project.py b/integrations/agentbricks/tests/harness/project.py index 29b85f46b..bd9626989 100644 --- a/integrations/agentbricks/tests/harness/project.py +++ b/integrations/agentbricks/tests/harness/project.py @@ -10,8 +10,8 @@ import pathlib import shutil import uuid -from collections.abc import Mapping, Sequence -from typing import Any, Protocol +from collections.abc import Mapping +from typing import Any import tomli import tomlkit @@ -20,13 +20,6 @@ from harness.run import APP_PREFIX, E2EError, RunConfig, project_prefix -class ToolSpec(Protocol): - """What ``write_manifest`` needs from a tool: its [[tools]] TOML and any user-owned files.""" - - toml: str - files: Mapping[str, str] - - @dataclasses.dataclass(frozen=True) class Project: path: pathlib.Path @@ -75,20 +68,6 @@ def write_files(self, files: Mapping[str, str]) -> None: for relative_path, content in files.items(): self.write_file(relative_path, content) - def write_manifest(self, tool_specs: Sequence[ToolSpec]) -> None: - """Direct authoring: append these tools to the agent.toml `agentbricks init` wrote. - - Everything else, such as the stores and the tracing binding, stays as init wrote it, the - same as for CLI authoring. - """ - target = self.path / "agent.toml" - tools = "".join(f"\n{spec.toml}" for spec in tool_specs if spec.toml) - log(f"# write {target}: direct authoring; no agentbricks tools command") - with target.open("a", encoding="utf-8") as manifest: - manifest.write(tools) - for spec in tool_specs: - self.write_files(spec.files) - def _pin_project_wheel(project: pathlib.Path, wheel: pathlib.Path) -> None: vendor_dir = project / "agentbricks_e2e_wheels" diff --git a/integrations/agentbricks/tests/harness/run.py b/integrations/agentbricks/tests/harness/run.py index 6eb5b1a85..75ede3909 100644 --- a/integrations/agentbricks/tests/harness/run.py +++ b/integrations/agentbricks/tests/harness/run.py @@ -12,7 +12,6 @@ import re import tempfile -AUTHORING_PATHS = ("cli", "direct") E2E_MODEL = "system.ai.gpt-5-2" APP_PREFIX = "agent-bricks-" diff --git a/integrations/agentbricks/tests/harness/workspace.py b/integrations/agentbricks/tests/harness/workspace.py index 2d58da0de..3454f5b83 100644 --- a/integrations/agentbricks/tests/harness/workspace.py +++ b/integrations/agentbricks/tests/harness/workspace.py @@ -268,6 +268,22 @@ def user_name(self) -> str: 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())) diff --git a/integrations/agentbricks/tests/suites/README.md b/integrations/agentbricks/tests/suites/README.md index 36ceadd30..5834c727a 100644 --- a/integrations/agentbricks/tests/suites/README.md +++ b/integrations/agentbricks/tests/suites/README.md @@ -1,7 +1,8 @@ # Agent Bricks E2E suites -Black-box tests of the `agentbricks` CLI: each test runs the real CLI (`init`, `tools add`, `dev`, -`deploy`, `endpoint invoke`) and checks the agent it produces. +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 @@ -12,7 +13,7 @@ uv run pytest tests/suites -n 6 --dist loadgroup \ --databricks-profile --catalog ``` -- Narrow it to a suite (`tests/suites/tools`), a file, or a test (`"...test_deploy.py::test_deploy[cli]"`). +- 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`, @@ -34,8 +35,8 @@ Every resource a run creates carries its run id, and is deleted at three points: `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, schemas) and the Apps' Lakebase roles, - which `deployments delete` leaves behind (`harness.cleanup.delete_run_leftovers`). A failure + 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: @@ -57,7 +58,8 @@ Every resource a run creates carries its run id, and is deleted at three points: | `tests/harness/workspace.py` | Databricks SDK helpers for what is not under test: prerequisites (UC objects, warehouse, Genie), UC privilege checks, 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//` | One folder per feature: `tools`, `memory_store`, `session_store`, `tracing` | +| `tests/suites/test_.py` | One file per CLI command group: `test_tools`, `test_dev`, `test_deploy` (with `deployments`), `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 @@ -65,11 +67,11 @@ rather than inside it: it is plain Python the suites import as `harness` (via `p ## Adding things -- **A tool:** add one `ToolDefinition(name, requires, build)` in `tools/tool_definitions.py`. +- **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 suite:** add a folder with `test_*.py` files. A test takes the `project` fixture, runs +- **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. Add a `conftest.py` only for resources the CLI under test does not create - itself; create them per test and delete them in the fixture's teardown. + 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 index 729943b41..cf1c0d79c 100644 --- a/integrations/agentbricks/tests/suites/conftest.py +++ b/integrations/agentbricks/tests/suites/conftest.py @@ -1,24 +1,25 @@ -"""All pytest wiring for the E2E suites: options, hooks, and the fixtures every suite shares. +"""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. One command runs a suite, from ``integrations/agentbricks``: +answers what calls what. From ``integrations/agentbricks``: - uv run pytest tests/suites/tools -n 6 --dist loadgroup \ - --databricks-profile + uv run pytest tests/suites -n 6 --dist loadgroup --databricks-profile ``--workspace fake`` runs against no workspace; whatever needs one is skipped. -Each test deletes what it created when it ends (the ``agentbricks_cli`` fixture here, and suite -fixtures such as the tools suite's ``scratch_schema``). After every live run, the controller also -deletes whatever a crashed test or a failed teardown left (``harness.cleanup``). A run killed -before that point is cleaned up afterwards with ``--cleanup-only --run-id ``. +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 @@ -34,7 +35,14 @@ ) from harness.logs import log, log_command from harness.project import Project -from harness.run import RUN_ID_PATTERN, E2EError, RunConfig, run_output_dir +from harness.run import ( + RUN_ID_PATTERN, + E2EError, + RunConfig, + now, + run_output_dir, + scratch_schema_prefix, +) from harness.target import ( TargetWorkspace, TargetWorkspaceUnavailable, @@ -43,6 +51,8 @@ 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]() @@ -208,6 +218,61 @@ def catalog(request: pytest.FixtureRequest, target_workspace: TargetWorkspace) - 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 diff --git a/integrations/agentbricks/tests/suites/memory_store/test_memory_store.py b/integrations/agentbricks/tests/suites/memory_store/test_memory_store.py deleted file mode 100644 index 419ed2646..000000000 --- a/integrations/agentbricks/tests/suites/memory_store/test_memory_store.py +++ /dev/null @@ -1,40 +0,0 @@ -"""STUB: an `init` agent (default bindings); deploy created the memory store agent.toml declares. - -Real case still to write: the agent writes a memory in one invocation and recalls it in another. -""" - -from __future__ import annotations - -import json - -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_memory_store_created( - agentbricks_cli: AgentbricksCli, project: Project, cleanup: Cleanup -) -> None: - agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) - project.pin_runtime() - # Registered before deploy: a deploy that fails partway can still leave them behind. - cleanup(f"memory stores delete memory-stores/{project.declared_store('memory')} --yes") - cleanup(f"sessions stores delete {project.declared_store('session')} --yes") - cleanup(f"deployments delete {project.app_name} --yes") - declared = project.declared_store("memory") - assert declared, "agentbricks init declared no memory store" - agentbricks_cli.run(f"deploy {project.app_name} --source {project.path}", timeout=2400) - - store = json.loads( - agentbricks_cli.run(f"--output json memory stores get memory-stores/{declared}").stdout - ) - assert store.get("display_name") == declared, store - assert store.get("name") == f"memory-stores/{declared}", store - - -@pytest.mark.xfail(reason="STUB: no memory write/recall case written yet", run=False) -def test_memory_recalled_across_invocations(agentbricks_cli: AgentbricksCli) -> None: - raise NotImplementedError diff --git a/integrations/agentbricks/tests/suites/session_store/test_session_store.py b/integrations/agentbricks/tests/suites/session_store/test_session_store.py deleted file mode 100644 index 7959ee4d7..000000000 --- a/integrations/agentbricks/tests/suites/session_store/test_session_store.py +++ /dev/null @@ -1,38 +0,0 @@ -"""STUB: an `init` agent (default bindings); deploy created the session store agent.toml declares. - -Real case still to write: conversation state persists across invocations of one session_id and is -isolated between sessions. -""" - -from __future__ import annotations - -import json - -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_session_store_created( - agentbricks_cli: AgentbricksCli, project: Project, cleanup: Cleanup -) -> None: - agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) - project.pin_runtime() - # Registered before deploy: a deploy that fails partway can still leave them behind. - cleanup(f"memory stores delete memory-stores/{project.declared_store('memory')} --yes") - cleanup(f"sessions stores delete {project.declared_store('session')} --yes") - cleanup(f"deployments delete {project.app_name} --yes") - declared = project.declared_store("session") - assert declared, "agentbricks init declared no session store" - agentbricks_cli.run(f"deploy {project.app_name} --source {project.path}", timeout=2400) - - store = json.loads(agentbricks_cli.run(f"--output json sessions stores get {declared}").stdout) - assert store.get("session_store_name") == declared, store - - -@pytest.mark.xfail(reason="STUB: no session-persistence case written yet", run=False) -def test_session_persists_across_invocations(agentbricks_cli: AgentbricksCli) -> None: - raise NotImplementedError diff --git a/integrations/agentbricks/tests/suites/test_deploy.py b/integrations/agentbricks/tests/suites/test_deploy.py new file mode 100644 index 000000000..a86a68bfd --- /dev/null +++ b/integrations/agentbricks/tests/suites/test_deploy.py @@ -0,0 +1,59 @@ +"""`agentbricks deploy` and `deployments`: every tool works on the deployed App, then delete removes it. + +`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, not_found +from harness.cleanup import Cleanup, wait_until_app_gone +from harness.logs import log +from harness.project import Project +from harness.workspace import Workspace +from tool_definitions import Tool, assert_declared, invocation_json + + +def test_deploy( + agentbricks_cli: AgentbricksCli, + project: Project, + cleanup: Cleanup, + workspace_client: Workspace, + tools: tuple[Tool, ...], + subtests, +) -> None: + agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) + project.pin_runtime() + app = project.app_name + # Registered before deploy: a deploy that fails partway can still leave them behind. + cleanup(f"memory stores delete memory-stores/{project.declared_store('memory')} --yes") + cleanup(f"sessions stores delete {project.declared_store('session')} --yes") + 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) + assert_declared(project, tool) + + agentbricks_cli.run(f"deploy {app} --source {project.path}", timeout=2400) + deployment = json.loads(agentbricks_cli.run(f"--output json deployments get {app}").stdout) + for tool in tools: + with subtests.test(tool=tool.name): + if tool.grant is not None: + # Deploy adds the grant before returning, but UC can apply it a little later, and + # a tool called in between is denied inside a successful response. + waited = workspace_client.wait_until_effective(deployment, tool.grant) + log(f"{tool.name}: {tool.grant[3]} effective for {app} after {waited:.0f}s") + output = agentbricks_cli.run( + f"--output json endpoint invoke {app} --path /api/invocations " + f"--json {invocation_json(tool.prompt)}", + timeout=360, + retry_server_errors=True, + ).stdout + tool.check(output) + + agentbricks_cli.run(f"deployments delete {app} --yes", timeout=1200) + # `deployments delete` returns while the App is still deleting. + wait_until_app_gone(agentbricks_cli, app) + assert not_found(agentbricks_cli.run(f"--output json deployments get {app}", check=False)) diff --git a/integrations/agentbricks/tests/suites/test_dev.py b/integrations/agentbricks/tests/suites/test_dev.py new file mode 100644 index 000000000..2e3344f15 --- /dev/null +++ b/integrations/agentbricks/tests/suites/test_dev.py @@ -0,0 +1,42 @@ +"""`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, assert_declared, invocation_json + +# 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) + assert_declared(project, tool) + + 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.run( + f"--output json endpoint invoke --url http://127.0.0.1:{port} --no-auth " + f"--path /api/invocations --json {invocation_json(tool.prompt)}", + timeout=360, + retry_server_errors=True, + ).stdout + 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..82eb49d6a --- /dev/null +++ b/integrations/agentbricks/tests/suites/test_memory.py @@ -0,0 +1,29 @@ +"""`agentbricks memory`: deploy creates the memory store `agentbricks init` declares.""" + +from __future__ import annotations + +import json + +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_memory_store(agentbricks_cli: AgentbricksCli, project: Project, cleanup: Cleanup) -> None: + agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) + project.pin_runtime() + store = project.declared_store("memory") + assert store, "agentbricks init declared no memory store" + # Registered before deploy: a deploy that fails partway can still leave them behind. + cleanup(f"memory stores delete memory-stores/{store} --yes") + cleanup(f"sessions stores delete {project.declared_store('session')} --yes") + cleanup(f"deployments delete {project.app_name} --yes") + agentbricks_cli.run(f"deploy {project.app_name} --source {project.path}", timeout=2400) + + found = json.loads( + agentbricks_cli.run(f"--output json memory stores get memory-stores/{store}").stdout + ) + assert found.get("name") == f"memory-stores/{store}", found diff --git a/integrations/agentbricks/tests/suites/test_sessions.py b/integrations/agentbricks/tests/suites/test_sessions.py new file mode 100644 index 000000000..ae9f2e80f --- /dev/null +++ b/integrations/agentbricks/tests/suites/test_sessions.py @@ -0,0 +1,27 @@ +"""`agentbricks sessions`: deploy creates the session store `agentbricks init` declares.""" + +from __future__ import annotations + +import json + +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_session_store(agentbricks_cli: AgentbricksCli, project: Project, cleanup: Cleanup) -> None: + agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) + project.pin_runtime() + store = project.declared_store("session") + assert store, "agentbricks init declared no session store" + # Registered before deploy: a deploy that fails partway can still leave them behind. + cleanup(f"memory stores delete memory-stores/{project.declared_store('memory')} --yes") + cleanup(f"sessions stores delete {store} --yes") + cleanup(f"deployments delete {project.app_name} --yes") + agentbricks_cli.run(f"deploy {project.app_name} --source {project.path}", timeout=2400) + + found = json.loads(agentbricks_cli.run(f"--output json sessions stores get {store}").stdout) + assert found.get("session_store_name") == store, found diff --git a/integrations/agentbricks/tests/suites/test_tools.py b/integrations/agentbricks/tests/suites/test_tools.py new file mode 100644 index 000000000..ff839a4cb --- /dev/null +++ b/integrations/agentbricks/tests/suites/test_tools.py @@ -0,0 +1,23 @@ +"""`agentbricks tools`: `tools add` declares each tool in agent.toml, and `tools remove` takes it out.""" + +from __future__ import annotations + +from harness.agentbricks_cli import AgentbricksCli +from harness.project import Project +from tool_definitions import Tool, assert_declared + + +def test_tools_add_and_remove( + agentbricks_cli: AgentbricksCli, project: Project, tools: tuple[Tool, ...] +) -> None: + agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) + project.pin_runtime() + added = [tool for tool in tools if tool.add_command and tool.tool_id] + for tool in added: + agentbricks_cli.run(f"tools add {tool.add_command} --source {project.path}") + assert_declared(project, tool) + + for tool in added: + agentbricks_cli.run(f"tools remove {tool.tool_id} --source {project.path}") + declared = [t["id"] for t in project.manifest().get("tools", [])] + assert not declared, f"tools remove left {declared} in agent.toml" diff --git a/integrations/agentbricks/tests/suites/test_tracing.py b/integrations/agentbricks/tests/suites/test_tracing.py new file mode 100644 index 000000000..23ed5f3ec --- /dev/null +++ b/integrations/agentbricks/tests/suites/test_tracing.py @@ -0,0 +1,45 @@ +"""`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 +from tool_definitions import invocation_json + + +# 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 + # Registered before deploy: a deploy that fails partway can still leave them behind. + cleanup(f"memory stores delete memory-stores/{project.declared_store('memory')} --yes") + cleanup(f"sessions stores delete {project.declared_store('session')} --yes") + cleanup(f"deployments delete {app} --yes") + agentbricks_cli.run(f"deploy {app} --source {project.path}", timeout=2400) + agentbricks_cli.run( + f"--output json endpoint invoke {app} --path /api/invocations " + f"--json {invocation_json('Reply with one word: hello.')}", + timeout=360, + retry_server_errors=True, + ) + + # 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/tools/tool_definitions.py b/integrations/agentbricks/tests/suites/tool_definitions.py similarity index 75% rename from integrations/agentbricks/tests/suites/tools/tool_definitions.py rename to integrations/agentbricks/tests/suites/tool_definitions.py index 8b5b80262..65853a529 100644 --- a/integrations/agentbricks/tests/suites/tools/tool_definitions.py +++ b/integrations/agentbricks/tests/suites/tool_definitions.py @@ -7,6 +7,7 @@ import shlex import uuid from collections.abc import Callable, Mapping +from typing import Any from harness.project import Project from harness.run import E2E_MODEL @@ -32,12 +33,11 @@ class Tool: name: str prompt: str - # Asserts the serialized response proves the tool ran. + # Asserts the `endpoint invoke --output json` output shows the tool ran and returned its result. check: Check # What follows `agentbricks tools add`; None when the tool is a user-owned file instead. add_command: str | None - # The [[tools]] entry for direct authoring, and any user-owned files the tool needs. - toml: str + # User-owned files the tool needs. files: Mapping[str, str] # The id agent.toml declares it under, and whether it must be App-authed. tool_id: str | None @@ -63,27 +63,49 @@ def excerpt(serialized: str, head: int = 1000, tail: int = 3000) -> str: return f"{serialized[:head]} ...[{omitted} chars omitted]... {serialized[-tail:]}" -def contains(marker: str) -> Check: +def tool_results(serialized: str) -> list[dict[str, Any]]: + """The tool steps of an `endpoint invoke --output json` response, as the agent reports them. + + The LangGraph agent lists every step in ``body.output.output``; a tool's own result is the item + with ``type == "tool"``, separate from the model's reply. + """ + 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) and step.get("type") == "tool"] + + +def _successful(serialized: str, matches: Callable[[dict[str, Any]], bool], what: str) -> None: + results = tool_results(serialized) + assert any(r.get("status") == "success" and matches(r) for r in results), ( + f"No successful tool result {what}: {excerpt(serialized)}" + ) + + +def returned(marker: str) -> Check: + """A tool returned ``marker``, which only the tool can produce.""" + def check(serialized: str) -> None: - assert marker in serialized, f"Missing exact marker {marker!r}: {excerpt(serialized)}" + _successful(serialized, lambda r: marker in str(r.get("content")), f"containing {marker!r}") return check def check_web_search(serialized: str) -> None: - lowered = serialized.lower() - tool_evidence = any(v in lowered for v in ("web_search", "web search", "search")) - assert tool_evidence and "https" in lowered and len(serialized) >= 80, ( - f"Missing web-search execution/result evidence: {excerpt(serialized)}" + _successful( + serialized, + lambda r: r.get("name") == "web_search" and "https://" in str(r.get("content")), + "from web_search with an https result", ) def check_genie(serialized: str) -> None: - lowered = serialized.lower() - assert ( - any(v in lowered for v in ("genie_ask", "genie", "conversation_id")) - and len(serialized) >= 80 - ), f"Missing Genie execution/result evidence: {excerpt(serialized)}" + _successful( + serialized, + lambda r: "genie" in str(r.get("name")) and bool(r.get("content")), + "from the Genie tool", + ) def sandbox_tool(uc_function: UcFunction) -> Tool: @@ -94,17 +116,10 @@ def sandbox_tool(uc_function: UcFunction) -> Tool: f"entire text file at {uc_function.volume_file_path}. Return only its exact contents; " "do not fabricate them." ), - check=contains(uc_function.volume_marker), + check=returned(uc_function.volume_marker), # The CLI's default auth: the sandbox reads the volume as the user invoking the agent, so - # deploy adds no App grant. A manifest entry without `auth` would run as the App instead. + # deploy adds no App grant. add_command=f"sandbox --scope volume:{uc_function.volume}", - toml=( - '[[tools]]\nid = "sandbox"\nauth = "user"\n' - 'source = { kind = "sandbox", service = "system.ai.sandbox" }\n' - "policy = { downscope = [\n" - f' {{ resource = "volume:{uc_function.volume}", permission = "read_only" }},\n' - "] }\n" - ), files={}, tool_id="sandbox", app_auth=False, @@ -122,10 +137,6 @@ def web_search_tool(**_resources: object) -> Tool: ), check=check_web_search, add_command="mcp system.ai.web_search --auth app", - toml=( - '[[tools]]\nid = "web_search"\nauth = "app"\n' - 'source = { kind = "mcp", service = "system.ai.web_search" }\n' - ), files={}, tool_id="web_search", app_auth=True, @@ -139,9 +150,8 @@ def python_tool() -> Tool: "You must call the matrix_marker Python tool with value 'matrix'. " "Return its exact result." ), - check=contains(PYTHON_MARKER), + check=returned(PYTHON_MARKER), add_command=None, - toml="", files={PYTHON_MARKER_FILE: PYTHON_MARKER_SOURCE}, tool_id=None, app_auth=False, @@ -155,12 +165,8 @@ def uc_function_tool(uc_function: UcFunction) -> Tool: "You must call the Unity Catalog function tool that computes the Agent Bricks marker, " "with value 'matrix'. Return the called tool's exact result." ), - check=contains(UC_MARKER), + check=returned(UC_MARKER), add_command=f"uc-function {uc_function.function} --name agentbricks_uc_marker", - toml=( - '[[tools]]\nid = "agentbricks_uc_marker"\n' - f'source = {{ kind = "uc_function", function = "{uc_function.function}" }}\n' - ), files={}, tool_id="agentbricks_uc_marker", app_auth=False, @@ -178,10 +184,6 @@ def genie_tool(genie_space: str) -> Tool: ), check=check_genie, add_command=f"genie-agent {genie_space} --name genie --auth app", - toml=( - '[[tools]]\nid = "genie"\nauth = "app"\n' - f'source = {{ kind = "genie_agent", space_id = "{genie_space}" }}\n' - ), files={}, tool_id="genie", app_auth=True, @@ -200,7 +202,7 @@ def genie_tool(genie_space: str) -> Tool: def assert_declared(project: Project, tool: Tool) -> None: - """After CLI authoring, agent.toml declares the tool under its id, App-authed if it must be.""" + """After `tools add`, agent.toml declares the tool under its id, App-authed if it must be.""" if tool.tool_id is None: return declared = {t["id"]: t for t in project.manifest().get("tools", [])} diff --git a/integrations/agentbricks/tests/suites/tools/conftest.py b/integrations/agentbricks/tests/suites/tools/conftest.py deleted file mode 100644 index ab26b5dc6..000000000 --- a/integrations/agentbricks/tests/suites/tools/conftest.py +++ /dev/null @@ -1,77 +0,0 @@ -"""Tools suite: every test runs once per authoring path, over the tools the target workspace can support. - -Each test gets its own temporary UC schema (``ab_e2e__``) holding the function and -volume its tools use, and deletes it, grants included, in its own teardown. Leftover cleanup deletes -any a crashed test left behind, by the run's prefix (see ``harness.cleanup``). -""" - -from __future__ import annotations - -import datetime as dt -import re -import uuid -from collections.abc import Iterator - -import pytest -from harness.logs import log -from harness.run import AUTHORING_PATHS, E2EError, RunConfig, now, scratch_schema_prefix -from harness.target import TargetWorkspace, require_workspace -from tool_definitions import TOOL_DEFINITIONS, Tool -from uc_objects import UcFunction, create_uc_function - - -@pytest.fixture(params=AUTHORING_PATHS) -def authoring(request: pytest.FixtureRequest) -> str: - return request.param - - -@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) diff --git a/integrations/agentbricks/tests/suites/tools/test_deploy.py b/integrations/agentbricks/tests/suites/tools/test_deploy.py deleted file mode 100644 index 5aef992bb..000000000 --- a/integrations/agentbricks/tests/suites/tools/test_deploy.py +++ /dev/null @@ -1,66 +0,0 @@ -"""Each tool works on a deployed App, per authoring path: invoking it returns its marker. - -CLI authoring adds every tool with `agentbricks tools add`; direct authoring writes them into -agent.toml. Either way the App is deployed once, and a subtest per tool invokes it, so each tool is -reported independently. - -`agentbricks deploy` returns once the App is running its new source, so the next step can use it. -""" - -from __future__ import annotations - -import json -from typing import Any - -from harness.agentbricks_cli import AgentbricksCli -from harness.cleanup import Cleanup -from harness.logs import log -from harness.project import Project -from harness.workspace import Workspace -from tool_definitions import Tool, assert_declared, invocation_json - - -def test_deploy( - agentbricks_cli: AgentbricksCli, - project: Project, - cleanup: Cleanup, - workspace_client: Workspace, - tools: tuple[Tool, ...], - authoring: str, - subtests, -) -> None: - agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) - project.pin_runtime() - app = project.app_name - # Registered before deploy: a deploy that fails partway can still leave them behind. - cleanup(f"memory stores delete memory-stores/{project.declared_store('memory')} --yes") - cleanup(f"sessions stores delete {project.declared_store('session')} --yes") - cleanup(f"deployments delete {app} --yes") - - def invoke_and_check(tool: Tool, deployment: dict[str, Any]) -> None: - if tool.grant is not None: - # Deploy adds the grant before returning, but UC can apply it a little later, and a - # tool called in between is denied inside a successful response. - waited = workspace_client.wait_until_effective(deployment, tool.grant) - log(f"{tool.name}: {tool.grant[3]} effective for {app} after {waited:.0f}s") - output = agentbricks_cli.run( - f"--output json endpoint invoke {app} --path /api/invocations " - f"--json {invocation_json(tool.prompt)}", - timeout=360, - retry_server_errors=True, - ).stdout - tool.check(output) - - if authoring == "cli": - 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) - assert_declared(project, tool) - elif authoring == "direct": - project.write_manifest(tools) - agentbricks_cli.run(f"deploy {app} --source {project.path}", timeout=2400) - deployment = json.loads(agentbricks_cli.run(f"--output json deployments get {app}").stdout) - for tool in tools: - with subtests.test(tool=tool.name): - invoke_and_check(tool, deployment) diff --git a/integrations/agentbricks/tests/suites/tools/test_dev.py b/integrations/agentbricks/tests/suites/tools/test_dev.py deleted file mode 100644 index 4526a76fd..000000000 --- a/integrations/agentbricks/tests/suites/tools/test_dev.py +++ /dev/null @@ -1,57 +0,0 @@ -"""Each tool works in `agentbricks dev`, per authoring path. - -CLI authoring adds every tool with `agentbricks tools add`; direct authoring writes them into -agent.toml. Either way dev runs once, and a subtest per tool invokes it, so each tool is reported -independently. Every check calls the dev server with `agentbricks endpoint invoke`, the command the -CLI tells users to run. - -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, assert_declared, invocation_json - -# 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, ...], - authoring: str, - subtests, -) -> None: - agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) - project.pin_runtime() - - def check(tool: Tool, port: int) -> None: - output = agentbricks_cli.run( - f"--output json endpoint invoke --url http://127.0.0.1:{port} --no-auth " - f"--path /api/invocations --json {invocation_json(tool.prompt)}", - timeout=360, - retry_server_errors=True, - ).stdout - tool.check(output) - - if authoring == "cli": - 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) - assert_declared(project, tool) - elif authoring == "direct": - project.write_manifest(tools) - 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): - check(tool, port) diff --git a/integrations/agentbricks/tests/suites/tools/test_selection.py b/integrations/agentbricks/tests/suites/tools/test_selection.py deleted file mode 100644 index 7b4e168b4..000000000 --- a/integrations/agentbricks/tests/suites/tools/test_selection.py +++ /dev/null @@ -1,74 +0,0 @@ -"""MULTI-TOOL SELECTION check (deploy-only): with every tool bound, the agent picks the one asked for. - -`test_deploy` names the tool in each prompt; this prompt names only what the tool does, so it -catches a regression where binding several tools makes the model, or the tool surface, route to the -wrong one. -""" - -from __future__ import annotations - -import json - -from harness.agentbricks_cli import AgentbricksCli -from harness.cleanup import Cleanup -from harness.project import Project -from harness.workspace import Workspace -from tool_definitions import ( - PYTHON_MARKER, - UC_MARKER, - Tool, - assert_declared, - excerpt, - invocation_json, - uc_function_tool, -) -from uc_objects import UcFunction - -# Names the Unity Catalog function without its tool id, and the Python helper only to exclude it. -PROMPT = ( - "Use the Unity Catalog function tool to compute the Agent Bricks marker for the value " - "'matrix', not the local Python helper tool. Return the called tool's exact result." -) - - -def test_selects_uc_function_not_python_marker( - agentbricks_cli: AgentbricksCli, - project: Project, - cleanup: Cleanup, - workspace_client: Workspace, - tools: tuple[Tool, ...], - uc_function: UcFunction, - authoring: str, -) -> None: - agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) - project.pin_runtime() - app = project.app_name - # Registered before deploy: a deploy that fails partway can still leave them behind. - cleanup(f"memory stores delete memory-stores/{project.declared_store('memory')} --yes") - cleanup(f"sessions stores delete {project.declared_store('session')} --yes") - cleanup(f"deployments delete {app} --yes") - if authoring == "cli": - 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) - assert_declared(project, tool) - else: - project.write_manifest(tools) - agentbricks_cli.run(f"deploy {app} --source {project.path}", timeout=2400) - deployment = json.loads(agentbricks_cli.run(f"--output json deployments get {app}").stdout) - # UC can apply deploy's grant a little after deploy returns; a call in between is denied. - grant = uc_function_tool(uc_function).grant - assert grant is not None - workspace_client.wait_until_effective(deployment, grant) - serialized = agentbricks_cli.run( - f"--output json endpoint invoke {app} --path /api/invocations " - f"--json {invocation_json(PROMPT)}", - timeout=360, - retry_server_errors=True, - ).stdout - - assert UC_MARKER in serialized, f"The UC function tool was not called: {excerpt(serialized)}" - assert PYTHON_MARKER not in serialized, ( - f"The Python marker tool was called too: {excerpt(serialized)}" - ) diff --git a/integrations/agentbricks/tests/suites/tools/test_tools_add.py b/integrations/agentbricks/tests/suites/tools/test_tools_add.py deleted file mode 100644 index f4708aa22..000000000 --- a/integrations/agentbricks/tests/suites/tools/test_tools_add.py +++ /dev/null @@ -1,34 +0,0 @@ -"""`agentbricks tools add` validates what it adds before it touches agent.toml.""" - -from __future__ import annotations - -import json - -import pytest -from harness.agentbricks_cli import AgentbricksCli -from harness.project import Project - - -@pytest.mark.usefixtures("workspace_client") -def test_rejects_unavailable_mcp(agentbricks_cli: AgentbricksCli, project: Project) -> None: - """An MCP service that does not exist is refused, and agent.toml is left untouched.""" - agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) - project.pin_runtime() - before = (project.path / "agent.toml").read_bytes() - - rejected = agentbricks_cli.run( - "--output json tools add mcp system.ai.missing_service --name broken_mcp " - f"--source {project.path}", - check=False, - ) - - assert rejected.returncode != 0, "tools add accepted an unavailable MCP service" - assert (project.path / "agent.toml").read_bytes() == before, ( - "a refused tools add changed agent.toml" - ) - code = json.loads(rejected.stderr).get("error", {}).get("code") - assert code in {"NOT_FOUND", "RESOURCE_DOES_NOT_EXIST"}, rejected.stderr - - agentbricks_cli.run(f"tools remove mcp system.ai.missing_service --source {project.path}") - declared = [tool["id"] for tool in project.manifest().get("tools", [])] - assert "broken_mcp" not in declared, "tools remove left the broken MCP binding in agent.toml" diff --git a/integrations/agentbricks/tests/suites/tracing/test_tracing.py b/integrations/agentbricks/tests/suites/tracing/test_tracing.py deleted file mode 100644 index e7cc16fd5..000000000 --- a/integrations/agentbricks/tests/suites/tracing/test_tracing.py +++ /dev/null @@ -1,47 +0,0 @@ -"""STUB: an `init` agent (default bindings); the tracing experiment is bound and attached to the App. - -Real case still to write: an invocation exports a trace to the experiment. -""" - -from __future__ import annotations - -import json - -import pytest -from harness.agentbricks_cli import AgentbricksCli -from harness.cleanup import Cleanup -from harness.project import Project -from harness.workspace import Workspace, app_resources - - -def test_tracing_experiment_bound( - agentbricks_cli: AgentbricksCli, - project: Project, - cleanup: Cleanup, - workspace_client: Workspace, -) -> None: - agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) - project.pin_runtime() - # Registered before deploy: a deploy that fails partway can still leave them behind. - cleanup(f"memory stores delete memory-stores/{project.declared_store('memory')} --yes") - cleanup(f"sessions stores delete {project.declared_store('session')} --yes") - cleanup(f"deployments delete {project.app_name} --yes") - experiment_name = project.manifest().get("tracing", {}).get("experiment_name") - agentbricks_cli.run(f"deploy {project.app_name} --source {project.path}", timeout=2400) - assert experiment_name, "agent.toml binds no tracing experiment" - deployment = json.loads( - agentbricks_cli.run(f"--output json deployments get {project.app_name}").stdout - ) - - # Deploy attaches the experiment as an App resource. - experiments = [r for r in app_resources(deployment) if "experiment" in r] - assert experiments, "No experiment App resource" - expected_id = workspace_client.experiment_id(experiment_name) - assert any(str(r["experiment"].get("experiment_id")) == expected_id for r in experiments), ( - f"{experiment_name} ({expected_id}) not in {experiments}" - ) - - -@pytest.mark.xfail(reason="STUB: no trace-export case written yet", run=False) -def test_invocation_exports_trace(agentbricks_cli: AgentbricksCli) -> None: - raise NotImplementedError diff --git a/integrations/agentbricks/tests/suites/tools/uc_objects.py b/integrations/agentbricks/tests/suites/uc_objects.py similarity index 100% rename from integrations/agentbricks/tests/suites/tools/uc_objects.py rename to integrations/agentbricks/tests/suites/uc_objects.py From 649c28fe711ae8337c87a98de85c84136b5b57f1 Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Fri, 16 Oct 2026 18:05:22 +0000 Subject: [PATCH 38/45] [agentbricks] E2E suites: move invocation_json into the harness CLI module Co-authored-by: Isaac --- .../agentbricks/tests/harness/agentbricks_cli.py | 16 +++++++++++++++- integrations/agentbricks/tests/suites/README.md | 2 +- .../agentbricks/tests/suites/test_deploy.py | 4 ++-- .../agentbricks/tests/suites/test_dev.py | 4 ++-- .../agentbricks/tests/suites/test_tracing.py | 3 +-- .../agentbricks/tests/suites/tool_definitions.py | 15 --------------- 6 files changed, 21 insertions(+), 23 deletions(-) diff --git a/integrations/agentbricks/tests/harness/agentbricks_cli.py b/integrations/agentbricks/tests/harness/agentbricks_cli.py index 64cb30078..8abff11d3 100644 --- a/integrations/agentbricks/tests/harness/agentbricks_cli.py +++ b/integrations/agentbricks/tests/harness/agentbricks_cli.py @@ -10,6 +10,7 @@ import contextlib import dataclasses +import json import os import pathlib import re @@ -22,10 +23,11 @@ import time import urllib.error import urllib.request +import uuid from collections.abc import Iterator, Mapping from harness.logs import last_lines, last_nonempty_line, log, log_command -from harness.run import E2EError, RunConfig, now +from harness.run import E2E_MODEL, 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 @@ -229,6 +231,18 @@ def not_found(result: subprocess.CompletedProcess[str]) -> bool: return result.returncode != 0 and bool(_NOT_FOUND.search(result.stderr or result.stdout)) +def invocation_json(prompt: str) -> str: + """A shell-quoted /api/invocations body for `agentbricks endpoint invoke --json`.""" + invocation_id = str(uuid.uuid4()) + body = { + "id": invocation_id, + # session_id is a top-level invocation field; the adapter rejects it nested in input. + "session_id": invocation_id, + "input": {"model": E2E_MODEL, "messages": [{"role": "user", "content": prompt}]}, + } + return shlex.quote(json.dumps(body)) + + def free_port() -> int: with socket.socket() as probe: probe.bind(("127.0.0.1", 0)) diff --git a/integrations/agentbricks/tests/suites/README.md b/integrations/agentbricks/tests/suites/README.md index 5834c727a..88e642581 100644 --- a/integrations/agentbricks/tests/suites/README.md +++ b/integrations/agentbricks/tests/suites/README.md @@ -53,7 +53,7 @@ Every resource a run creates carries its run id, and is deleted at three points: | `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 | +| `tests/harness/agentbricks_cli.py` | Runs `agentbricks` commands: `cli.run("tools add ...")`, and `cli.start("dev ...")` in the background; `invocation_json` builds the `endpoint invoke --json` body | | `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), UC privilege checks, 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 | diff --git a/integrations/agentbricks/tests/suites/test_deploy.py b/integrations/agentbricks/tests/suites/test_deploy.py index a86a68bfd..0ac4e0f1d 100644 --- a/integrations/agentbricks/tests/suites/test_deploy.py +++ b/integrations/agentbricks/tests/suites/test_deploy.py @@ -7,12 +7,12 @@ import json -from harness.agentbricks_cli import AgentbricksCli, not_found +from harness.agentbricks_cli import AgentbricksCli, invocation_json, not_found from harness.cleanup import Cleanup, wait_until_app_gone from harness.logs import log from harness.project import Project from harness.workspace import Workspace -from tool_definitions import Tool, assert_declared, invocation_json +from tool_definitions import Tool, assert_declared def test_deploy( diff --git a/integrations/agentbricks/tests/suites/test_dev.py b/integrations/agentbricks/tests/suites/test_dev.py index 2e3344f15..18c3e61ac 100644 --- a/integrations/agentbricks/tests/suites/test_dev.py +++ b/integrations/agentbricks/tests/suites/test_dev.py @@ -7,9 +7,9 @@ from __future__ import annotations import pytest -from harness.agentbricks_cli import AgentbricksCli, free_port +from harness.agentbricks_cli import AgentbricksCli, free_port, invocation_json from harness.project import Project -from tool_definitions import Tool, assert_declared, invocation_json +from tool_definitions import Tool, assert_declared # 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") diff --git a/integrations/agentbricks/tests/suites/test_tracing.py b/integrations/agentbricks/tests/suites/test_tracing.py index 23ed5f3ec..af69461be 100644 --- a/integrations/agentbricks/tests/suites/test_tracing.py +++ b/integrations/agentbricks/tests/suites/test_tracing.py @@ -9,10 +9,9 @@ import time import pytest -from harness.agentbricks_cli import AgentbricksCli +from harness.agentbricks_cli import AgentbricksCli, invocation_json from harness.cleanup import Cleanup from harness.project import Project -from tool_definitions import invocation_json # Deploy needs a live workspace. diff --git a/integrations/agentbricks/tests/suites/tool_definitions.py b/integrations/agentbricks/tests/suites/tool_definitions.py index 65853a529..3c43904ff 100644 --- a/integrations/agentbricks/tests/suites/tool_definitions.py +++ b/integrations/agentbricks/tests/suites/tool_definitions.py @@ -4,13 +4,10 @@ import dataclasses import json -import shlex -import uuid from collections.abc import Callable, Mapping from typing import Any from harness.project import Project -from harness.run import E2E_MODEL from harness.workspace import GrantTuple from uc_objects import UcFunction @@ -213,15 +210,3 @@ def assert_declared(project: Project, tool: Tool) -> None: assert declared[tool.tool_id].get("auth") == "app", ( f"CLI-authored {tool.tool_id} binding is not App-auth." ) - - -def invocation_json(prompt: str) -> str: - """A shell-quoted /api/invocations body for `agentbricks endpoint invoke --json`.""" - invocation_id = str(uuid.uuid4()) - body = { - "id": invocation_id, - # session_id is a top-level invocation field; the adapter rejects it nested in input. - "session_id": invocation_id, - "input": {"model": E2E_MODEL, "messages": [{"role": "user", "content": prompt}]}, - } - return shlex.quote(json.dumps(body)) From 29cceb3bd9a3229ca178931d3463639d08064a33 Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Fri, 16 Oct 2026 18:29:58 +0000 Subject: [PATCH 39/45] [agentbricks] E2E suites: disable CLI telemetry for test runs Co-authored-by: Isaac --- .../agentbricks/tests/harness/agentbricks_cli.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/integrations/agentbricks/tests/harness/agentbricks_cli.py b/integrations/agentbricks/tests/harness/agentbricks_cli.py index 8abff11d3..63a9f11e1 100644 --- a/integrations/agentbricks/tests/harness/agentbricks_cli.py +++ b/integrations/agentbricks/tests/harness/agentbricks_cli.py @@ -87,8 +87,14 @@ def __init__(self, target_workspace: TargetWorkspace, run_config: RunConfig, tes 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. - self._env = {**os.environ, **target_workspace.env, "AGENTBRICKS_DISABLE_AGENT_HINT": "1"} + # 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( From f33218066bc0393b1ba2ebd124bd8833cd6a1fc3 Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Fri, 16 Oct 2026 18:46:24 +0000 Subject: [PATCH 40/45] [agentbricks] E2E suites: check stores and tools through invoke; read store names from deploy Co-authored-by: Isaac --- .../tests/harness/agentbricks_cli.py | 51 +++++++++++++++-- .../agentbricks/tests/harness/project.py | 10 ---- .../agentbricks/tests/suites/test_deploy.py | 15 +++-- .../agentbricks/tests/suites/test_dev.py | 3 +- .../agentbricks/tests/suites/test_memory.py | 57 ++++++++++++++----- .../agentbricks/tests/suites/test_sessions.py | 49 +++++++++++----- .../agentbricks/tests/suites/test_tools.py | 44 ++++++++++---- .../agentbricks/tests/suites/test_tracing.py | 12 ++-- .../tests/suites/tool_definitions.py | 46 +-------------- 9 files changed, 178 insertions(+), 109 deletions(-) diff --git a/integrations/agentbricks/tests/harness/agentbricks_cli.py b/integrations/agentbricks/tests/harness/agentbricks_cli.py index 63a9f11e1..b613e3ee4 100644 --- a/integrations/agentbricks/tests/harness/agentbricks_cli.py +++ b/integrations/agentbricks/tests/harness/agentbricks_cli.py @@ -25,6 +25,7 @@ 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 E2E_MODEL, E2EError, RunConfig, now @@ -237,18 +238,60 @@ def not_found(result: subprocess.CompletedProcess[str]) -> bool: return result.returncode != 0 and bool(_NOT_FOUND.search(result.stderr or result.stdout)) -def invocation_json(prompt: str) -> str: - """A shell-quoted /api/invocations body for `agentbricks endpoint invoke --json`.""" +def invocation_json(prompt: str, *, session_id: str | None = None, actor: str | None = None) -> str: + """A shell-quoted /api/invocations body for `agentbricks endpoint invoke --json`. + + 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": E2E_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": invocation_id, - "input": {"model": E2E_MODEL, "messages": [{"role": "user", "content": prompt}]}, + "session_id": session_id or invocation_id, + "input": agent_input, } return shlex.quote(json.dumps(body)) +# 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)) diff --git a/integrations/agentbricks/tests/harness/project.py b/integrations/agentbricks/tests/harness/project.py index bd9626989..ebd42e135 100644 --- a/integrations/agentbricks/tests/harness/project.py +++ b/integrations/agentbricks/tests/harness/project.py @@ -11,7 +11,6 @@ import shutil import uuid from collections.abc import Mapping -from typing import Any import tomli import tomlkit @@ -48,15 +47,6 @@ def is_pinned(self) -> bool: sources = document.get("tool", {}).get("uv", {}).get("sources", {}) return "databricks-agentbricks" in sources - def manifest(self) -> dict[str, Any]: - return tomli.loads((self.path / "agent.toml").read_text(encoding="utf-8")) - - def declared_store(self, kind: str) -> str | None: - """The name agent.toml gives the ``memory`` or ``session`` store, if it declares one.""" - if not (self.path / "agent.toml").exists(): - return None - return self.manifest().get(f"{kind}_store", {}).get("name") - 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 diff --git a/integrations/agentbricks/tests/suites/test_deploy.py b/integrations/agentbricks/tests/suites/test_deploy.py index 0ac4e0f1d..30d599872 100644 --- a/integrations/agentbricks/tests/suites/test_deploy.py +++ b/integrations/agentbricks/tests/suites/test_deploy.py @@ -12,7 +12,7 @@ from harness.logs import log from harness.project import Project from harness.workspace import Workspace -from tool_definitions import Tool, assert_declared +from tool_definitions import Tool def test_deploy( @@ -26,17 +26,20 @@ def test_deploy( agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) project.pin_runtime() app = project.app_name - # Registered before deploy: a deploy that fails partway can still leave them behind. - cleanup(f"memory stores delete memory-stores/{project.declared_store('memory')} --yes") - cleanup(f"sessions stores delete {project.declared_store('session')} --yes") + # 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) - assert_declared(project, tool) - agentbricks_cli.run(f"deploy {app} --source {project.path}", timeout=2400) + 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") deployment = json.loads(agentbricks_cli.run(f"--output json deployments get {app}").stdout) for tool in tools: with subtests.test(tool=tool.name): diff --git a/integrations/agentbricks/tests/suites/test_dev.py b/integrations/agentbricks/tests/suites/test_dev.py index 18c3e61ac..554c4cde5 100644 --- a/integrations/agentbricks/tests/suites/test_dev.py +++ b/integrations/agentbricks/tests/suites/test_dev.py @@ -9,7 +9,7 @@ import pytest from harness.agentbricks_cli import AgentbricksCli, free_port, invocation_json from harness.project import Project -from tool_definitions import Tool, assert_declared +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") @@ -24,7 +24,6 @@ def test_dev( if tool.add_command: agentbricks_cli.run(f"tools add {tool.add_command} --source {project.path}") project.write_files(tool.files) - assert_declared(project, tool) port = free_port() with agentbricks_cli.start( diff --git a/integrations/agentbricks/tests/suites/test_memory.py b/integrations/agentbricks/tests/suites/test_memory.py index 82eb49d6a..c6192f9dc 100644 --- a/integrations/agentbricks/tests/suites/test_memory.py +++ b/integrations/agentbricks/tests/suites/test_memory.py @@ -1,29 +1,58 @@ -"""`agentbricks memory`: deploy creates the memory store `agentbricks init` declares.""" +"""`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 +from harness.agentbricks_cli import AgentbricksCli, excerpt, invocation_json, 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_store(agentbricks_cli: AgentbricksCli, project: Project, cleanup: Cleanup) -> None: +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() - store = project.declared_store("memory") - assert store, "agentbricks init declared no memory store" - # Registered before deploy: a deploy that fails partway can still leave them behind. - cleanup(f"memory stores delete memory-stores/{store} --yes") - cleanup(f"sessions stores delete {project.declared_store('session')} --yes") - cleanup(f"deployments delete {project.app_name} --yes") - agentbricks_cli.run(f"deploy {project.app_name} --source {project.path}", timeout=2400) - - found = json.loads( - agentbricks_cli.run(f"--output json memory stores get memory-stores/{store}").stdout + 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]}" + + def invoke(prompt: str) -> str: + # Each invocation is a new session; only the actor carries over. + return agentbricks_cli.run( + f"--output json endpoint invoke {app} --path /api/invocations " + f"--json {invocation_json(prompt, actor=actor)}", + timeout=360, + retry_server_errors=True, + ).stdout + + remembered = invoke( + f"Remember this fact about me in long-term memory: my codename is {codename}." ) - assert found.get("name") == f"memory-stores/{store}", found + 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 = invoke("What is my codename? Look it up in your long-term memory.") + 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 index ae9f2e80f..c5873b28e 100644 --- a/integrations/agentbricks/tests/suites/test_sessions.py +++ b/integrations/agentbricks/tests/suites/test_sessions.py @@ -1,27 +1,50 @@ -"""`agentbricks sessions`: deploy creates the session store `agentbricks init` declares.""" +"""`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 +from harness.agentbricks_cli import AgentbricksCli, excerpt, invocation_json, 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_session_store(agentbricks_cli: AgentbricksCli, project: Project, cleanup: Cleanup) -> None: +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() - store = project.declared_store("session") - assert store, "agentbricks init declared no session store" - # Registered before deploy: a deploy that fails partway can still leave them behind. - cleanup(f"memory stores delete memory-stores/{project.declared_store('memory')} --yes") - cleanup(f"sessions stores delete {store} --yes") - cleanup(f"deployments delete {project.app_name} --yes") - agentbricks_cli.run(f"deploy {project.app_name} --source {project.path}", timeout=2400) - - found = json.loads(agentbricks_cli.run(f"--output json sessions stores get {store}").stdout) - assert found.get("session_store_name") == store, found + 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]}" + + def invoke(prompt: str, session_id: str) -> str: + return agentbricks_cli.run( + f"--output json endpoint invoke {app} --path /api/invocations " + f"--json {invocation_json(prompt, session_id=session_id)}", + timeout=360, + retry_server_errors=True, + ).stdout + + invoke(f"The secret word for this conversation is {word}. Reply with just: noted.", session) + same = invoke(QUESTION, session) + assert word in reply(same), f"The session lost its history: {excerpt(same)}" + other = invoke(QUESTION, str(uuid.uuid4())) + 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 index ff839a4cb..0745cc895 100644 --- a/integrations/agentbricks/tests/suites/test_tools.py +++ b/integrations/agentbricks/tests/suites/test_tools.py @@ -1,23 +1,43 @@ -"""`agentbricks tools`: `tools add` declares each tool in agent.toml, and `tools remove` takes it out.""" +"""`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 -from harness.agentbricks_cli import AgentbricksCli +import pytest +from harness.agentbricks_cli import AgentbricksCli, free_port, invocation_json from harness.project import Project -from tool_definitions import Tool, assert_declared +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_add_and_remove( - agentbricks_cli: AgentbricksCli, project: Project, tools: tuple[Tool, ...] +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() - added = [tool for tool in tools if tool.add_command and tool.tool_id] - for tool in added: + 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}") - assert_declared(project, tool) - - for tool in added: + for tool in removed: agentbricks_cli.run(f"tools remove {tool.tool_id} --source {project.path}") - declared = [t["id"] for t in project.manifest().get("tools", [])] - assert not declared, f"tools remove left {declared} in agent.toml" + + 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.run( + f"--output json endpoint invoke --url http://127.0.0.1:{port} --no-auth " + f"--path /api/invocations --json {invocation_json(tool.prompt)}", + timeout=360, + retry_server_errors=True, + ).stdout + # The tool's own check passes only when the tool ran and returned its result. + with pytest.raises(AssertionError): + tool.check(output) diff --git a/integrations/agentbricks/tests/suites/test_tracing.py b/integrations/agentbricks/tests/suites/test_tracing.py index af69461be..1c68de518 100644 --- a/integrations/agentbricks/tests/suites/test_tracing.py +++ b/integrations/agentbricks/tests/suites/test_tracing.py @@ -20,11 +20,15 @@ def test_tracing(agentbricks_cli: AgentbricksCli, project: Project, cleanup: Cle agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600) project.pin_runtime() app = project.app_name - # Registered before deploy: a deploy that fails partway can still leave them behind. - cleanup(f"memory stores delete memory-stores/{project.declared_store('memory')} --yes") - cleanup(f"sessions stores delete {project.declared_store('session')} --yes") + # Before deploy: a deploy that fails partway can still leave the App behind. cleanup(f"deployments delete {app} --yes") - agentbricks_cli.run(f"deploy {app} --source {project.path}", timeout=2400) + 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.run( f"--output json endpoint invoke {app} --path /api/invocations " f"--json {invocation_json('Reply with one word: hello.')}", diff --git a/integrations/agentbricks/tests/suites/tool_definitions.py b/integrations/agentbricks/tests/suites/tool_definitions.py index 3c43904ff..7feaf3486 100644 --- a/integrations/agentbricks/tests/suites/tool_definitions.py +++ b/integrations/agentbricks/tests/suites/tool_definitions.py @@ -3,11 +3,10 @@ from __future__ import annotations import dataclasses -import json from collections.abc import Callable, Mapping from typing import Any -from harness.project import Project +from harness.agentbricks_cli import excerpt, tool_results from harness.workspace import GrantTuple from uc_objects import UcFunction @@ -36,9 +35,8 @@ class Tool: add_command: str | None # User-owned files the tool needs. files: Mapping[str, str] - # The id agent.toml declares it under, and whether it must be App-authed. + # The id `tools add` binds it under, which `tools remove` takes. tool_id: str | None - app_auth: bool # The App resource a deploy grants for it; None when it needs no UC grant. grant: GrantTuple | None = None @@ -52,27 +50,6 @@ class ToolDefinition: build: Callable[..., 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 tool_results(serialized: str) -> list[dict[str, Any]]: - """The tool steps of an `endpoint invoke --output json` response, as the agent reports them. - - The LangGraph agent lists every step in ``body.output.output``; a tool's own result is the item - with ``type == "tool"``, separate from the model's reply. - """ - 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) and step.get("type") == "tool"] - - def _successful(serialized: str, matches: Callable[[dict[str, Any]], bool], what: str) -> None: results = tool_results(serialized) assert any(r.get("status") == "success" and matches(r) for r in results), ( @@ -119,7 +96,6 @@ def sandbox_tool(uc_function: UcFunction) -> Tool: add_command=f"sandbox --scope volume:{uc_function.volume}", files={}, tool_id="sandbox", - app_auth=False, grant=None, ) @@ -136,7 +112,6 @@ def web_search_tool(**_resources: object) -> Tool: add_command="mcp system.ai.web_search --auth app", files={}, tool_id="web_search", - app_auth=True, ) @@ -151,7 +126,6 @@ def python_tool() -> Tool: add_command=None, files={PYTHON_MARKER_FILE: PYTHON_MARKER_SOURCE}, tool_id=None, - app_auth=False, ) @@ -166,7 +140,6 @@ def uc_function_tool(uc_function: UcFunction) -> Tool: add_command=f"uc-function {uc_function.function} --name agentbricks_uc_marker", files={}, tool_id="agentbricks_uc_marker", - app_auth=False, grant=("uc_securable", uc_function.function, "FUNCTION", "EXECUTE"), ) @@ -183,7 +156,6 @@ def genie_tool(genie_space: str) -> Tool: add_command=f"genie-agent {genie_space} --name genie --auth app", files={}, tool_id="genie", - app_auth=True, grant=("genie_space", genie_space, "GENIE_SPACE", "CAN_RUN"), ) @@ -196,17 +168,3 @@ def genie_tool(genie_space: str) -> Tool: ToolDefinition("uc_function", ("uc_function",), uc_function_tool), ToolDefinition("genie", ("genie_space",), genie_tool), ) - - -def assert_declared(project: Project, tool: Tool) -> None: - """After `tools add`, agent.toml declares the tool under its id, App-authed if it must be.""" - if tool.tool_id is None: - return - declared = {t["id"]: t for t in project.manifest().get("tools", [])} - assert tool.tool_id in declared, ( - f"tools add did not declare {tool.tool_id!r}: {sorted(declared)}" - ) - if tool.app_auth: - assert declared[tool.tool_id].get("auth") == "app", ( - f"CLI-authored {tool.tool_id} binding is not App-auth." - ) From f776c64befb05dcc12191d84b36bcbee255bb56e Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Fri, 16 Oct 2026 19:04:04 +0000 Subject: [PATCH 41/45] [agentbricks] E2E suites: add --model for workspaces that name models differently Co-authored-by: Isaac --- .../tests/harness/agentbricks_cli.py | 47 ++++++++++--------- integrations/agentbricks/tests/harness/run.py | 4 ++ .../agentbricks/tests/suites/README.md | 4 +- .../agentbricks/tests/suites/conftest.py | 7 +++ .../agentbricks/tests/suites/test_deploy.py | 4 +- .../agentbricks/tests/suites/test_dev.py | 4 +- .../agentbricks/tests/suites/test_memory.py | 4 +- .../agentbricks/tests/suites/test_sessions.py | 4 +- .../agentbricks/tests/suites/test_tools.py | 4 +- .../agentbricks/tests/suites/test_tracing.py | 4 +- 10 files changed, 50 insertions(+), 36 deletions(-) diff --git a/integrations/agentbricks/tests/harness/agentbricks_cli.py b/integrations/agentbricks/tests/harness/agentbricks_cli.py index b613e3ee4..71ffa7c0f 100644 --- a/integrations/agentbricks/tests/harness/agentbricks_cli.py +++ b/integrations/agentbricks/tests/harness/agentbricks_cli.py @@ -28,7 +28,7 @@ from typing import Any from harness.logs import last_lines, last_nonempty_line, log, log_command -from harness.run import E2E_MODEL, E2EError, RunConfig, now +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 @@ -128,6 +128,29 @@ def run( raise E2EError(f"Command failed ({result.returncode}): {shlex.join(argv)}\n{output}") return result + def invocation_json( + self, prompt: str, *, session_id: str | None = None, actor: str | None = None + ) -> str: + """A shell-quoted /api/invocations body for `agentbricks endpoint invoke --json`. + + 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 shlex.quote(json.dumps(body)) + @contextlib.contextmanager def start( self, command: str, *, env: Mapping[str, str] | None = None @@ -238,28 +261,6 @@ def not_found(result: subprocess.CompletedProcess[str]) -> bool: return result.returncode != 0 and bool(_NOT_FOUND.search(result.stderr or result.stdout)) -def invocation_json(prompt: str, *, session_id: str | None = None, actor: str | None = None) -> str: - """A shell-quoted /api/invocations body for `agentbricks endpoint invoke --json`. - - 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": E2E_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 shlex.quote(json.dumps(body)) - - # 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"``). diff --git a/integrations/agentbricks/tests/harness/run.py b/integrations/agentbricks/tests/harness/run.py index 75ede3909..f5acbf7f9 100644 --- a/integrations/agentbricks/tests/harness/run.py +++ b/integrations/agentbricks/tests/harness/run.py @@ -13,6 +13,8 @@ import tempfile E2E_MODEL = "system.ai.gpt-5-2" + + APP_PREFIX = "agent-bricks-" @@ -51,4 +53,6 @@ class RunConfig: 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/suites/README.md b/integrations/agentbricks/tests/suites/README.md index 88e642581..104da7a5b 100644 --- a/integrations/agentbricks/tests/suites/README.md +++ b/integrations/agentbricks/tests/suites/README.md @@ -20,6 +20,8 @@ uv run pytest tests/suites -n 6 --dist loadgroup \ `--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. @@ -53,7 +55,7 @@ Every resource a run creates carries its run id, and is deleted at three points: | `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; `invocation_json` builds the `endpoint invoke --json` body | +| `tests/harness/agentbricks_cli.py` | Runs `agentbricks` commands: `cli.run("tools add ...")`, and `cli.start("dev ...")` in the background; `cli.invocation_json(...)` builds the `endpoint invoke --json` body, 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), UC privilege checks, 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 | diff --git a/integrations/agentbricks/tests/suites/conftest.py b/integrations/agentbricks/tests/suites/conftest.py index cf1c0d79c..7e6406401 100644 --- a/integrations/agentbricks/tests/suites/conftest.py +++ b/integrations/agentbricks/tests/suites/conftest.py @@ -36,6 +36,7 @@ from harness.logs import log, log_command from harness.project import Project from harness.run import ( + E2E_MODEL, RUN_ID_PATTERN, E2EError, RunConfig, @@ -81,6 +82,11 @@ def pytest_addoption(parser: pytest.Parser) -> None: 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, @@ -153,6 +159,7 @@ def _base_run_config(config: pytest.Config) -> RunConfig: run_id=run_id, output=output, warehouse_id=option("warehouse_id"), + model=option("model"), wheel=option("wheel"), ) config.stash[_RUN_KEY] = run diff --git a/integrations/agentbricks/tests/suites/test_deploy.py b/integrations/agentbricks/tests/suites/test_deploy.py index 30d599872..baebd5af2 100644 --- a/integrations/agentbricks/tests/suites/test_deploy.py +++ b/integrations/agentbricks/tests/suites/test_deploy.py @@ -7,7 +7,7 @@ import json -from harness.agentbricks_cli import AgentbricksCli, invocation_json, not_found +from harness.agentbricks_cli import AgentbricksCli, not_found from harness.cleanup import Cleanup, wait_until_app_gone from harness.logs import log from harness.project import Project @@ -50,7 +50,7 @@ def test_deploy( log(f"{tool.name}: {tool.grant[3]} effective for {app} after {waited:.0f}s") output = agentbricks_cli.run( f"--output json endpoint invoke {app} --path /api/invocations " - f"--json {invocation_json(tool.prompt)}", + f"--json {agentbricks_cli.invocation_json(tool.prompt)}", timeout=360, retry_server_errors=True, ).stdout diff --git a/integrations/agentbricks/tests/suites/test_dev.py b/integrations/agentbricks/tests/suites/test_dev.py index 554c4cde5..32a0efb07 100644 --- a/integrations/agentbricks/tests/suites/test_dev.py +++ b/integrations/agentbricks/tests/suites/test_dev.py @@ -7,7 +7,7 @@ from __future__ import annotations import pytest -from harness.agentbricks_cli import AgentbricksCli, free_port, invocation_json +from harness.agentbricks_cli import AgentbricksCli, free_port from harness.project import Project from tool_definitions import Tool @@ -34,7 +34,7 @@ def test_dev( with subtests.test(tool=tool.name): output = agentbricks_cli.run( f"--output json endpoint invoke --url http://127.0.0.1:{port} --no-auth " - f"--path /api/invocations --json {invocation_json(tool.prompt)}", + f"--path /api/invocations --json {agentbricks_cli.invocation_json(tool.prompt)}", timeout=360, retry_server_errors=True, ).stdout diff --git a/integrations/agentbricks/tests/suites/test_memory.py b/integrations/agentbricks/tests/suites/test_memory.py index c6192f9dc..92cdd2613 100644 --- a/integrations/agentbricks/tests/suites/test_memory.py +++ b/integrations/agentbricks/tests/suites/test_memory.py @@ -10,7 +10,7 @@ import uuid import pytest -from harness.agentbricks_cli import AgentbricksCli, excerpt, invocation_json, tool_results +from harness.agentbricks_cli import AgentbricksCli, excerpt, tool_results from harness.cleanup import Cleanup from harness.project import Project @@ -38,7 +38,7 @@ def invoke(prompt: str) -> str: # Each invocation is a new session; only the actor carries over. return agentbricks_cli.run( f"--output json endpoint invoke {app} --path /api/invocations " - f"--json {invocation_json(prompt, actor=actor)}", + f"--json {agentbricks_cli.invocation_json(prompt, actor=actor)}", timeout=360, retry_server_errors=True, ).stdout diff --git a/integrations/agentbricks/tests/suites/test_sessions.py b/integrations/agentbricks/tests/suites/test_sessions.py index c5873b28e..12d594feb 100644 --- a/integrations/agentbricks/tests/suites/test_sessions.py +++ b/integrations/agentbricks/tests/suites/test_sessions.py @@ -9,7 +9,7 @@ import uuid import pytest -from harness.agentbricks_cli import AgentbricksCli, excerpt, invocation_json, reply +from harness.agentbricks_cli import AgentbricksCli, excerpt, reply from harness.cleanup import Cleanup from harness.project import Project @@ -38,7 +38,7 @@ def test_sessions(agentbricks_cli: AgentbricksCli, project: Project, cleanup: Cl def invoke(prompt: str, session_id: str) -> str: return agentbricks_cli.run( f"--output json endpoint invoke {app} --path /api/invocations " - f"--json {invocation_json(prompt, session_id=session_id)}", + f"--json {agentbricks_cli.invocation_json(prompt, session_id=session_id)}", timeout=360, retry_server_errors=True, ).stdout diff --git a/integrations/agentbricks/tests/suites/test_tools.py b/integrations/agentbricks/tests/suites/test_tools.py index 0745cc895..5823f86fa 100644 --- a/integrations/agentbricks/tests/suites/test_tools.py +++ b/integrations/agentbricks/tests/suites/test_tools.py @@ -6,7 +6,7 @@ from __future__ import annotations import pytest -from harness.agentbricks_cli import AgentbricksCli, free_port, invocation_json +from harness.agentbricks_cli import AgentbricksCli, free_port from harness.project import Project from tool_definitions import Tool @@ -34,7 +34,7 @@ def test_tools_remove( with subtests.test(tool=tool.name): output = agentbricks_cli.run( f"--output json endpoint invoke --url http://127.0.0.1:{port} --no-auth " - f"--path /api/invocations --json {invocation_json(tool.prompt)}", + f"--path /api/invocations --json {agentbricks_cli.invocation_json(tool.prompt)}", timeout=360, retry_server_errors=True, ).stdout diff --git a/integrations/agentbricks/tests/suites/test_tracing.py b/integrations/agentbricks/tests/suites/test_tracing.py index 1c68de518..e388f8f86 100644 --- a/integrations/agentbricks/tests/suites/test_tracing.py +++ b/integrations/agentbricks/tests/suites/test_tracing.py @@ -9,7 +9,7 @@ import time import pytest -from harness.agentbricks_cli import AgentbricksCli, invocation_json +from harness.agentbricks_cli import AgentbricksCli from harness.cleanup import Cleanup from harness.project import Project @@ -31,7 +31,7 @@ def test_tracing(agentbricks_cli: AgentbricksCli, project: Project, cleanup: Cle cleanup(f"sessions stores delete {deployed['session_store']} --yes") agentbricks_cli.run( f"--output json endpoint invoke {app} --path /api/invocations " - f"--json {invocation_json('Reply with one word: hello.')}", + f"--json {agentbricks_cli.invocation_json('Reply with one word: hello.')}", timeout=360, retry_server_errors=True, ) From 0c49ea2177786a8e50a5a08f944702b3e815dd24 Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Fri, 16 Oct 2026 20:56:18 +0000 Subject: [PATCH 42/45] [agentbricks] E2E suites: no grant checks; test_deploy leaves deletion to its registered cleanup Co-authored-by: Isaac --- .../agentbricks/tests/harness/target.py | 2 +- .../agentbricks/tests/harness/workspace.py | 91 +------------------ .../agentbricks/tests/suites/README.md | 4 +- .../agentbricks/tests/suites/test_deploy.py | 20 +--- .../tests/suites/tool_definitions.py | 10 +- 5 files changed, 12 insertions(+), 115 deletions(-) diff --git a/integrations/agentbricks/tests/harness/target.py b/integrations/agentbricks/tests/harness/target.py index 7c4ce33b6..857b345ed 100644 --- a/integrations/agentbricks/tests/harness/target.py +++ b/integrations/agentbricks/tests/harness/target.py @@ -43,7 +43,7 @@ def cli_args(self) -> Sequence[str]: @property def client(self) -> Workspace | None: - """SDK access for grants, cleanup and UC setup; None for the fake workspace.""" + """SDK access for cleanup and UC setup; None for the fake workspace.""" ... diff --git a/integrations/agentbricks/tests/harness/workspace.py b/integrations/agentbricks/tests/harness/workspace.py index 3454f5b83..35c1c7f3e 100644 --- a/integrations/agentbricks/tests/harness/workspace.py +++ b/integrations/agentbricks/tests/harness/workspace.py @@ -1,10 +1,9 @@ -"""Databricks SDK helpers for what is not under test: the prerequisites the suites need, and checks -of results the CLI has no command to show. +"""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, UC privilege checks, 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``. +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 @@ -21,14 +20,12 @@ from databricks.sdk import WorkspaceClient from databricks.sdk.errors import DatabricksError, NotFound -from databricks.sdk.service.catalog import SecurableType, VolumeType +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 -GrantTuple = tuple[str, str, str, str] - _BRANCH = re.compile(r"projects/[^/]+/branches/[^/]+") @@ -180,41 +177,6 @@ def app_logs(self, name: str, log_path: pathlib.Path) -> pathlib.Path | None: log(f"App runtime logs captured: {log_path}") return log_path - # Grants - - def wait_until_effective( - self, deployment: dict[str, Any], grant: GrantTuple, timeout: float = 120 - ) -> float: - """Wait until a UC ``grant`` deploy added is effective for the App; returns the seconds waited. - - ``deployment`` is ``agentbricks deployments get --output json``. Deploy reconciles the App's - UC resources before returning, but the privilege can become effective for its principal a - little later; a tool called in between is denied. - """ - kind, full_name, securable_type, privilege = grant - if kind != "uc_securable": - return 0.0 - principal = _principal(deployment) - started = time.monotonic() - while True: - held = effective_privileges( - self.client, SecurableType(securable_type), full_name, principal - ) - if privilege in held or "ALL_PRIVILEGES" in held: - return time.monotonic() - started - if time.monotonic() - started > timeout: - raise E2EError( - f"{privilege} on {full_name} not effective for {deployment.get('name')} " - f"after {timeout:.0f}s: {held}" - ) - time.sleep(5) - - def experiment_id(self, name: str) -> str: - experiment = self.client.experiments.get_by_name(name).experiment - if experiment is None or not experiment.experiment_id: - raise E2EError(f"Experiment {name!r} does not exist.") - return str(experiment.experiment_id) - # Workarounds for the end-of-run cleanup: what `agentbricks deployments delete` leaves behind. def runtime_store(self, app_name: str) -> dict[str, Any]: @@ -306,46 +268,3 @@ def _wait(operation: Any) -> None: wait = getattr(operation, "wait", None) if callable(wait): wait() - - -# Reading `agentbricks deployments get --output json` - - -def _principal(deployment: dict[str, Any]) -> str: - principal = deployment.get("service_principal_client_id") - if not principal: - raise E2EError(f"Deployment has no service_principal_client_id: {deployment}") - return str(principal) - - -def app_resources(deployment: dict[str, Any]) -> list[dict[str, Any]]: - resources = deployment.get("resources") or [] - if not isinstance(resources, list): - raise E2EError(f"App resources are not a list: {resources}") - return resources - - -def effective_privileges( - client: WorkspaceClient, securable_type: SecurableType, full_name: str, principal: str -) -> list[str]: - privileges: set[str] = set() - page_token: str | None = None - while True: - response = client.grants.get_effective( - securable_type.value, - full_name, - max_results=0, - principal=principal, - **({"page_token": page_token} if page_token else {}), - ) - for assignment in response.privilege_assignments or (): - if assignment.principal != principal: - continue - privileges.update( - privilege.privilege.value - for privilege in assignment.privileges or () - if privilege.privilege is not None - ) - page_token = response.next_page_token - if not page_token: - return sorted(privileges) diff --git a/integrations/agentbricks/tests/suites/README.md b/integrations/agentbricks/tests/suites/README.md index 104da7a5b..5f19e0009 100644 --- a/integrations/agentbricks/tests/suites/README.md +++ b/integrations/agentbricks/tests/suites/README.md @@ -57,10 +57,10 @@ Every resource a run creates carries its run id, and is deleted at three points: | `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.invocation_json(...)` builds the `endpoint invoke --json` body, 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), UC privilege checks, and cleanup no CLI command does yet | +| `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` (with `deployments`), `test_memory`, `test_sessions`, `test_tracing` | +| `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/` diff --git a/integrations/agentbricks/tests/suites/test_deploy.py b/integrations/agentbricks/tests/suites/test_deploy.py index baebd5af2..a07c888b5 100644 --- a/integrations/agentbricks/tests/suites/test_deploy.py +++ b/integrations/agentbricks/tests/suites/test_deploy.py @@ -1,4 +1,4 @@ -"""`agentbricks deploy` and `deployments`: every tool works on the deployed App, then delete removes it. +"""`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. """ @@ -7,11 +7,9 @@ import json -from harness.agentbricks_cli import AgentbricksCli, not_found -from harness.cleanup import Cleanup, wait_until_app_gone -from harness.logs import log +from harness.agentbricks_cli import AgentbricksCli +from harness.cleanup import Cleanup from harness.project import Project -from harness.workspace import Workspace from tool_definitions import Tool @@ -19,7 +17,6 @@ def test_deploy( agentbricks_cli: AgentbricksCli, project: Project, cleanup: Cleanup, - workspace_client: Workspace, tools: tuple[Tool, ...], subtests, ) -> None: @@ -40,14 +37,8 @@ def test_deploy( ) cleanup(f"memory stores delete memory-stores/{deployed['memory_store']} --yes") cleanup(f"sessions stores delete {deployed['session_store']} --yes") - deployment = json.loads(agentbricks_cli.run(f"--output json deployments get {app}").stdout) for tool in tools: with subtests.test(tool=tool.name): - if tool.grant is not None: - # Deploy adds the grant before returning, but UC can apply it a little later, and - # a tool called in between is denied inside a successful response. - waited = workspace_client.wait_until_effective(deployment, tool.grant) - log(f"{tool.name}: {tool.grant[3]} effective for {app} after {waited:.0f}s") output = agentbricks_cli.run( f"--output json endpoint invoke {app} --path /api/invocations " f"--json {agentbricks_cli.invocation_json(tool.prompt)}", @@ -55,8 +46,3 @@ def test_deploy( retry_server_errors=True, ).stdout tool.check(output) - - agentbricks_cli.run(f"deployments delete {app} --yes", timeout=1200) - # `deployments delete` returns while the App is still deleting. - wait_until_app_gone(agentbricks_cli, app) - assert not_found(agentbricks_cli.run(f"--output json deployments get {app}", check=False)) diff --git a/integrations/agentbricks/tests/suites/tool_definitions.py b/integrations/agentbricks/tests/suites/tool_definitions.py index 7feaf3486..c170661bd 100644 --- a/integrations/agentbricks/tests/suites/tool_definitions.py +++ b/integrations/agentbricks/tests/suites/tool_definitions.py @@ -1,4 +1,4 @@ -"""The tools under test: how each is added, prompted, verified, and what a deploy grants for it.""" +"""The tools under test: how each is added, prompted and verified.""" from __future__ import annotations @@ -7,7 +7,6 @@ from typing import Any from harness.agentbricks_cli import excerpt, tool_results -from harness.workspace import GrantTuple from uc_objects import UcFunction PYTHON_MARKER_FILE = "agent/tools/matrix_marker.py" @@ -37,8 +36,6 @@ class Tool: files: Mapping[str, str] # The id `tools add` binds it under, which `tools remove` takes. tool_id: str | None - # The App resource a deploy grants for it; None when it needs no UC grant. - grant: GrantTuple | None = None @dataclasses.dataclass(frozen=True) @@ -91,12 +88,9 @@ def sandbox_tool(uc_function: UcFunction) -> Tool: "do not fabricate them." ), check=returned(uc_function.volume_marker), - # The CLI's default auth: the sandbox reads the volume as the user invoking the agent, so - # deploy adds no App grant. add_command=f"sandbox --scope volume:{uc_function.volume}", files={}, tool_id="sandbox", - grant=None, ) @@ -140,7 +134,6 @@ def uc_function_tool(uc_function: UcFunction) -> Tool: add_command=f"uc-function {uc_function.function} --name agentbricks_uc_marker", files={}, tool_id="agentbricks_uc_marker", - grant=("uc_securable", uc_function.function, "FUNCTION", "EXECUTE"), ) @@ -156,7 +149,6 @@ def genie_tool(genie_space: str) -> Tool: add_command=f"genie-agent {genie_space} --name genie --auth app", files={}, tool_id="genie", - grant=("genie_space", genie_space, "GENIE_SPACE", "CAN_RUN"), ) From 950e0627ae0c972aefc3da1d61610f8ae7986426 Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Fri, 16 Oct 2026 21:10:49 +0000 Subject: [PATCH 43/45] [agentbricks] E2E suites: one AgentbricksCli.invoke for endpoint invoke Co-authored-by: Isaac --- .../tests/harness/agentbricks_cli.py | 18 +++++++++++------ .../agentbricks/tests/suites/README.md | 2 +- .../agentbricks/tests/suites/test_deploy.py | 7 +------ .../agentbricks/tests/suites/test_dev.py | 9 +++------ .../agentbricks/tests/suites/test_memory.py | 20 ++++++++----------- .../agentbricks/tests/suites/test_sessions.py | 18 +++++++---------- .../agentbricks/tests/suites/test_tools.py | 9 +++------ .../agentbricks/tests/suites/test_tracing.py | 7 +------ 8 files changed, 36 insertions(+), 54 deletions(-) diff --git a/integrations/agentbricks/tests/harness/agentbricks_cli.py b/integrations/agentbricks/tests/harness/agentbricks_cli.py index 71ffa7c0f..fc3f5e21d 100644 --- a/integrations/agentbricks/tests/harness/agentbricks_cli.py +++ b/integrations/agentbricks/tests/harness/agentbricks_cli.py @@ -128,13 +128,14 @@ def run( raise E2EError(f"Command failed ({result.returncode}): {shlex.join(argv)}\n{output}") return result - def invocation_json( - self, prompt: str, *, session_id: str | None = None, actor: str | None = None + def invoke( + self, target: str, prompt: str, *, session_id: str | None = None, actor: str | None = None ) -> str: - """A shell-quoted /api/invocations body for `agentbricks endpoint invoke --json`. + """``agentbricks --output json endpoint invoke `` with ``prompt``; returns its output. - Each call is a new session unless ``session_id`` is given. The agent keys long-term memory by - ``actor``, which defaults to the session. + ``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] = { @@ -149,7 +150,12 @@ def invocation_json( "session_id": session_id or invocation_id, "input": agent_input, } - return shlex.quote(json.dumps(body)) + 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( diff --git a/integrations/agentbricks/tests/suites/README.md b/integrations/agentbricks/tests/suites/README.md index 5f19e0009..b66a0d1bb 100644 --- a/integrations/agentbricks/tests/suites/README.md +++ b/integrations/agentbricks/tests/suites/README.md @@ -55,7 +55,7 @@ Every resource a run creates carries its run id, and is deleted at three points: | `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.invocation_json(...)` builds the `endpoint invoke --json` body, with the `--model` the run asks for | +| `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 | diff --git a/integrations/agentbricks/tests/suites/test_deploy.py b/integrations/agentbricks/tests/suites/test_deploy.py index a07c888b5..e9fc72d7c 100644 --- a/integrations/agentbricks/tests/suites/test_deploy.py +++ b/integrations/agentbricks/tests/suites/test_deploy.py @@ -39,10 +39,5 @@ def test_deploy( cleanup(f"sessions stores delete {deployed['session_store']} --yes") for tool in tools: with subtests.test(tool=tool.name): - output = agentbricks_cli.run( - f"--output json endpoint invoke {app} --path /api/invocations " - f"--json {agentbricks_cli.invocation_json(tool.prompt)}", - timeout=360, - retry_server_errors=True, - ).stdout + 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 index 32a0efb07..e507f625b 100644 --- a/integrations/agentbricks/tests/suites/test_dev.py +++ b/integrations/agentbricks/tests/suites/test_dev.py @@ -32,10 +32,7 @@ def test_dev( server.wait_until_serving(port) for tool in tools: with subtests.test(tool=tool.name): - output = agentbricks_cli.run( - f"--output json endpoint invoke --url http://127.0.0.1:{port} --no-auth " - f"--path /api/invocations --json {agentbricks_cli.invocation_json(tool.prompt)}", - timeout=360, - retry_server_errors=True, - ).stdout + 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 index 92cdd2613..131668da0 100644 --- a/integrations/agentbricks/tests/suites/test_memory.py +++ b/integrations/agentbricks/tests/suites/test_memory.py @@ -34,24 +34,20 @@ def test_memory(agentbricks_cli: AgentbricksCli, project: Project, cleanup: Clea actor = f"e2e-{uuid.uuid4().hex[:8]}" codename = f"AGENTBRICKS_MEMORY_{uuid.uuid4().hex[:12]}" - def invoke(prompt: str) -> str: - # Each invocation is a new session; only the actor carries over. - return agentbricks_cli.run( - f"--output json endpoint invoke {app} --path /api/invocations " - f"--json {agentbricks_cli.invocation_json(prompt, actor=actor)}", - timeout=360, - retry_server_errors=True, - ).stdout - - remembered = invoke( - f"Remember this fact about me in long-term memory: my codename is {codename}." + # 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 = invoke("What is my codename? Look it up in your long-term memory.") + 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) diff --git a/integrations/agentbricks/tests/suites/test_sessions.py b/integrations/agentbricks/tests/suites/test_sessions.py index 12d594feb..aee7b208c 100644 --- a/integrations/agentbricks/tests/suites/test_sessions.py +++ b/integrations/agentbricks/tests/suites/test_sessions.py @@ -35,16 +35,12 @@ def test_sessions(agentbricks_cli: AgentbricksCli, project: Project, cleanup: Cl session = str(uuid.uuid4()) word = f"AGENTBRICKS_SESSION_{uuid.uuid4().hex[:12]}" - def invoke(prompt: str, session_id: str) -> str: - return agentbricks_cli.run( - f"--output json endpoint invoke {app} --path /api/invocations " - f"--json {agentbricks_cli.invocation_json(prompt, session_id=session_id)}", - timeout=360, - retry_server_errors=True, - ).stdout - - invoke(f"The secret word for this conversation is {word}. Reply with just: noted.", session) - same = invoke(QUESTION, session) + 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 = invoke(QUESTION, str(uuid.uuid4())) + 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 index 5823f86fa..98f492b4e 100644 --- a/integrations/agentbricks/tests/suites/test_tools.py +++ b/integrations/agentbricks/tests/suites/test_tools.py @@ -32,12 +32,9 @@ def test_tools_remove( server.wait_until_serving(port) for tool in removed: with subtests.test(tool=tool.name): - output = agentbricks_cli.run( - f"--output json endpoint invoke --url http://127.0.0.1:{port} --no-auth " - f"--path /api/invocations --json {agentbricks_cli.invocation_json(tool.prompt)}", - timeout=360, - retry_server_errors=True, - ).stdout + output = agentbricks_cli.invoke( + f"--url http://127.0.0.1:{port} --no-auth", tool.prompt + ) # The tool's own check passes only when the tool ran and returned its result. with pytest.raises(AssertionError): tool.check(output) diff --git a/integrations/agentbricks/tests/suites/test_tracing.py b/integrations/agentbricks/tests/suites/test_tracing.py index e388f8f86..ecf0314bf 100644 --- a/integrations/agentbricks/tests/suites/test_tracing.py +++ b/integrations/agentbricks/tests/suites/test_tracing.py @@ -29,12 +29,7 @@ def test_tracing(agentbricks_cli: AgentbricksCli, project: Project, cleanup: Cle ) cleanup(f"memory stores delete memory-stores/{deployed['memory_store']} --yes") cleanup(f"sessions stores delete {deployed['session_store']} --yes") - agentbricks_cli.run( - f"--output json endpoint invoke {app} --path /api/invocations " - f"--json {agentbricks_cli.invocation_json('Reply with one word: hello.')}", - timeout=360, - retry_server_errors=True, - ) + 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 From 3c086fecb238857aff853bc6fcdefcf3ccd6c1ba Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Fri, 16 Oct 2026 21:16:25 +0000 Subject: [PATCH 44/45] [agentbricks] E2E suites: Tool.check_not_called for removed tools Co-authored-by: Isaac --- .../agentbricks/tests/suites/test_tools.py | 4 +- .../tests/suites/tool_definitions.py | 62 +++++++++---------- 2 files changed, 29 insertions(+), 37 deletions(-) diff --git a/integrations/agentbricks/tests/suites/test_tools.py b/integrations/agentbricks/tests/suites/test_tools.py index 98f492b4e..12c9bc74f 100644 --- a/integrations/agentbricks/tests/suites/test_tools.py +++ b/integrations/agentbricks/tests/suites/test_tools.py @@ -35,6 +35,4 @@ def test_tools_remove( output = agentbricks_cli.invoke( f"--url http://127.0.0.1:{port} --no-auth", tool.prompt ) - # The tool's own check passes only when the tool ran and returned its result. - with pytest.raises(AssertionError): - tool.check(output) + tool.check_not_called(output) diff --git a/integrations/agentbricks/tests/suites/tool_definitions.py b/integrations/agentbricks/tests/suites/tool_definitions.py index c170661bd..020f81c46 100644 --- a/integrations/agentbricks/tests/suites/tool_definitions.py +++ b/integrations/agentbricks/tests/suites/tool_definitions.py @@ -21,15 +21,15 @@ PYTHON_MARKER = "AGENTBRICKS_PYTHON_OK" -Check = Callable[[str], None] +# 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 - # Asserts the `endpoint invoke --output json` output shows the tool ran and returned its result. - check: Check + 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. @@ -37,6 +37,19 @@ class Tool: # 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: @@ -47,36 +60,17 @@ class ToolDefinition: build: Callable[..., Tool] -def _successful(serialized: str, matches: Callable[[dict[str, Any]], bool], what: str) -> None: - results = tool_results(serialized) - assert any(r.get("status") == "success" and matches(r) for r in results), ( - f"No successful tool result {what}: {excerpt(serialized)}" - ) - - -def returned(marker: str) -> Check: - """A tool returned ``marker``, which only the tool can produce.""" - - def check(serialized: str) -> None: - _successful(serialized, lambda r: marker in str(r.get("content")), f"containing {marker!r}") +def returned(marker: str) -> Called: + """The step returned ``marker``, which only the tool can produce.""" + return lambda r: marker in str(r.get("content")) - return check +def web_search_called(r: dict[str, Any]) -> bool: + return r.get("name") == "web_search" and "https://" in str(r.get("content")) -def check_web_search(serialized: str) -> None: - _successful( - serialized, - lambda r: r.get("name") == "web_search" and "https://" in str(r.get("content")), - "from web_search with an https result", - ) - -def check_genie(serialized: str) -> None: - _successful( - serialized, - lambda r: "genie" in str(r.get("name")) and bool(r.get("content")), - "from the Genie tool", - ) +def genie_called(r: dict[str, Any]) -> bool: + return "genie" in str(r.get("name")) and bool(r.get("content")) def sandbox_tool(uc_function: UcFunction) -> Tool: @@ -87,7 +81,7 @@ def sandbox_tool(uc_function: UcFunction) -> Tool: f"entire text file at {uc_function.volume_file_path}. Return only its exact contents; " "do not fabricate them." ), - check=returned(uc_function.volume_marker), + called=returned(uc_function.volume_marker), add_command=f"sandbox --scope volume:{uc_function.volume}", files={}, tool_id="sandbox", @@ -102,7 +96,7 @@ def web_search_tool(**_resources: object) -> Tool: "Search official Databricks documentation for Model Context Protocol, then " "return the title and https URL of one result. Do not answer from memory." ), - check=check_web_search, + called=web_search_called, add_command="mcp system.ai.web_search --auth app", files={}, tool_id="web_search", @@ -116,7 +110,7 @@ def python_tool() -> Tool: "You must call the matrix_marker Python tool with value 'matrix'. " "Return its exact result." ), - check=returned(PYTHON_MARKER), + called=returned(PYTHON_MARKER), add_command=None, files={PYTHON_MARKER_FILE: PYTHON_MARKER_SOURCE}, tool_id=None, @@ -130,7 +124,7 @@ def uc_function_tool(uc_function: UcFunction) -> Tool: "You must call the Unity Catalog function tool that computes the Agent Bricks marker, " "with value 'matrix'. Return the called tool's exact result." ), - check=returned(UC_MARKER), + called=returned(UC_MARKER), add_command=f"uc-function {uc_function.function} --name agentbricks_uc_marker", files={}, tool_id="agentbricks_uc_marker", @@ -145,7 +139,7 @@ def genie_tool(genie_space: str) -> Tool: "configured Genie space. Return a one-sentence summary based only on the tool " "response." ), - check=check_genie, + called=genie_called, add_command=f"genie-agent {genie_space} --name genie --auth app", files={}, tool_id="genie", From e5b2658586d0e89f372d4d597f2cbd0ede83322f Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Fri, 16 Oct 2026 21:36:49 +0000 Subject: [PATCH 45/45] [agentbricks] E2E suites: stop timed-out commands before teardown; require a completed Genie answer Co-authored-by: Isaac --- .../agentbricks/tests/harness/agentbricks_cli.py | 11 +++++++---- .../agentbricks/tests/suites/tool_definitions.py | 15 ++++++++++++++- 2 files changed, 21 insertions(+), 5 deletions(-) diff --git a/integrations/agentbricks/tests/harness/agentbricks_cli.py b/integrations/agentbricks/tests/harness/agentbricks_cli.py index fc3f5e21d..a46b4a38b 100644 --- a/integrations/agentbricks/tests/harness/agentbricks_cli.py +++ b/integrations/agentbricks/tests/harness/agentbricks_cli.py @@ -222,7 +222,7 @@ def _run_once( while process.poll() is None: elapsed = time.monotonic() - started if elapsed >= timeout: - terminate_process_group(process, grace=0) + 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) @@ -247,12 +247,14 @@ def _agentbricks_bin() -> str: 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: - return - if grace <= 0: - return + pass try: process.wait(timeout=grace) except subprocess.TimeoutExpired: @@ -260,6 +262,7 @@ def terminate_process_group(process: subprocess.Popen[str], *, grace: float = 20 os.killpg(process.pid, signal.SIGKILL) except ProcessLookupError: pass + process.wait() def not_found(result: subprocess.CompletedProcess[str]) -> bool: diff --git a/integrations/agentbricks/tests/suites/tool_definitions.py b/integrations/agentbricks/tests/suites/tool_definitions.py index 020f81c46..7c4e20ee1 100644 --- a/integrations/agentbricks/tests/suites/tool_definitions.py +++ b/integrations/agentbricks/tests/suites/tool_definitions.py @@ -3,6 +3,7 @@ from __future__ import annotations import dataclasses +import json from collections.abc import Callable, Mapping from typing import Any @@ -70,7 +71,19 @@ def web_search_called(r: dict[str, Any]) -> bool: def genie_called(r: dict[str, Any]) -> bool: - return "genie" in str(r.get("name")) and bool(r.get("content")) + # 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: