[elainewang-db]

Summary

Adds pytest E2E suites for the Agent Bricks CLI under integrations/agentbricks/tests/, to replace tests/e2e/tool_matrix.py. They sit at the top of the test pyramid: one happy-path scenario per CLI command group, each running the real CLI and checking the agent it produces. Every step under test is the command a developer would type:

agentbricks_cli.run(f"init --framework langgraph {project.path}", timeout=600)
project.pin_runtime()   # use the wheel built from this checkout, not the PyPI release
cleanup(f"deployments delete {app} --yes")   # runs when the test ends, pass or fail
agentbricks_cli.run(f"tools add mcp system.ai.web_search --auth app --source {project.path}")
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")
agentbricks_cli.invoke(app, prompt)   # --output json endpoint invoke, with the run's --model

• Tests only use the CLI and check only what endpoint invoke returns. They don't read agent.toml or depend on CLI internals. The SDK is used only for what isn't under test (UC objects, warehouse, Genie) and for cleanup the CLI can't do yet.
• Plain pytest. Setup and teardown are fixtures; no custom runner or config file.
• Each test deletes what it creates by registering the CLI delete command with a cleanup fixture (like t.Cleanup in databricks/cli). An end-of-run sweep deletes anything left.
• Ready for a fake workspace through a --workspace live|fake seam; fake is a stub for now.

The old script stays until the nightly runner switches over. tests/suites/README.md covers running, cleanup, layout, and adding a tool or command group.

uv run pytest tests/suites -n 6 --dist loadgroup \
    --databricks-profile <profile> --catalog <catalog> \
    [--model <model>] [--warehouse-id <id>] [--genie-space-id <id>]

Tests (tests/suites/, one file per command group)

File	The scenario
test_tools.py	tools add then tools remove each tool kind → dev → asking for each tool, the agent no longer has it
test_dev.py	init → tools add every tool → dev → endpoint invoke each tool
test_deploy.py	the same with deploy
test_memory.py	deploy → remember a random codename under one actor → in a new session, the recall tool returns it
test_sessions.py	deploy → tell the agent a random word → the same session answers with it, a different session doesn't
test_tracing.py	deploy → endpoint invoke → tracing list shows a trace in the project's experiment

The tools are sandbox, MCP (web search), Python, UC function and Genie (skipped without --genie-space-id), defined in tool_definitions.py. A tool passes only if the response contains its own successful tool step with the expected result: the random marker in the sandbox's volume file, the Python or UC function's marker, a web_search result with an https link, or a Genie answer with status COMPLETED, since Genie reports a failed answer inside a successful tool step. The model writing about the tool isn't enough. The tests check no grants: a tool that answers shows deploy gave the agent the access it needs. --model sets the model every invocation asks for, since workspaces name models differently (system.ai.gpt-5-2 on production, system.ai.databricks-gpt-5-2 on staging). endpoint invoke is retried only on a 5xx (3 attempts, 15s apart), which a new App can return while its session store comes up.

Harness (tests/harness/)

File	What it is
agentbricks_cli.py	Runs agentbricks commands (run, and start for dev), logging each to logs/<test>/
project.py	The test's run-scoped path and App name, its files, and pin_runtime()
cleanup.py	Cleanup, the apps.jsonl record, and the end-of-run sweep
workspace.py	SDK helpers for what isn't under test, plus the cleanup workarounds
target.py, run.py, logs.py	The live/fake seam, run id and name prefixes, log output

Cleanup

1. When a test ends, its registered commands run, newest first; they are the only per-test deletes. Deploying tests register deployments delete before deploying, so a deploy that fails partway is covered, and the store deletes right after, from the names deploy --output json reports ([agentbricks] Report the deployed memory and session stores in deploy --output json #760):

   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")

   Something already gone counts as deleted; other failures only warn. Before they run, the project fixture records the App's principal and Lakebase branch in apps.jsonl, saves a failed test's App logs, and fails a test that never called pin_runtime().

2. When the run ends, the sweep deletes everything carrying the run's name prefix (stores, Apps, synced source folders, tracing experiments, schemas), plus the Apps' Lakebase roles from apps.jsonl. A failure here fails the run.

3. After a killed run, pytest tests/suites --cleanup-only --run-id <id> ... runs the sweep alone, for an always() CI step.

Three sweep steps are workarounds for what nothing deletes yet: the App's synced source folder, its Lakebase role (agents-<principal>), and the tracing experiment deploy creates under /Shared/agentbricks_traces/.

Testing

• On the head (e5b2658): ruff and ty pass; 6 tests collected; --workspace fake skips all 6.
• Live, integration workspace, run in04: all 6 tests pass in 5 min 52 s, on the head's code (e5b2658), with [agentbricks] Keep the databricks CLI's own output out of --output json stdout #762 applied before it merged: test_deploy (sandbox, MCP, Python, UC function), test_memory, test_sessions, test_tracing, test_dev and test_tools_remove. The Genie tool was skipped (no --genie-space-id). A second --cleanup-only --run-id in04 found nothing left. Without [agentbricks] Keep the databricks CLI's own output out of --output json stdout #762, every deploy test failed to parse deploy --output json (run in01).
• Genie, live (run gn01), with --genie-space-id for the workspace's "Databricks AI Bridge Integration Test" space: test_dev and test_deploy pass with all 5 tools in 5 min 39 s. Both runs' genie_ask returned status COMPLETED with 2 attachments. The deployed App's first call returned 502 twice while it started; the 5xx retry covered it.
• test_tools_remove, live (run dv13): after tools remove, each prompt's response had no tool steps, and the agent said it doesn't have the tool.
• Live, integration workspace, run dv12: test_dev (4 tools, without --genie-space-id) and test_tools pass in 2 min 38s, with no cleanup warnings. In test_tools, none of the removed tools' prompts got a tool step, and the agent replied that the tool isn't available.
• The end-of-run sweep is checked live: in ml01 it deleted the 3 memory stores, 3 session stores (listing across pages) and 3 experiments the failed deploys left; on an older run it deleted that run's 8 experiments.
• pytest tests/unit_tests passes on main with [agentbricks] Report the deployed memory and session stores in deploy --output json #760's telemetry fix; this PR only touches tests/suites and tests/harness.

Issues found while building this, fixed here or in the CLI PRs it depends on: VIRTUAL_ENV leaking into apps run-local (#714); a stale app.agentbricksdev.yaml after SIGTERM (#726); the coding-agent hint corrupting --output json errors (#718, the suite sets AGENTBRICKS_DISABLE_AGENT_HINT=1); CLI usage telemetry counting test runs as product usage (the suite sets AGENTBRICKS_DISABLE_TELEMETRY=1); SDK auth picking the wrong profile when two share a host; an inherited PORT overriding the dev port; the 64-character MCP tool-name limit for UC functions (use a catalog of about 38 characters or fewer); and deploy requiring auth on every MCP, sandbox and Genie entry once any tool uses auth = "user".

Follow-ups

• Lakebase role limit (ES-2291356): the integration workspace's shared branch (databricks-internal-custom-agents/production) filled its 700 roles, so deploys failed; after the cleanup on the ticket it has 98. Most of it is likely per-database databricks_reader_*/databricks_writer_* roles Lakebase doesn't drop when a database is deleted; the API protects them. LKB-12168, an in-progress Lakebase epic, removes these roles.
• DeleteRuntimeStore leaves the App's role (agents-<principal>); the server creates it, so the server should drop it.
• deployments delete has no E2E test of its own yet; registered cleanup and the sweep use it, so a broken delete fails the run in the sweep. It should wait until the App is gone, remove the synced source folder (a CLI PR is in progress), and treat an App that doesn't exist as already deleted (today it errors: "Could not resolve the app's service principal for Runtime Store cleanup"). It also leaves the tracing experiment and the stores deploy created; whether it should remove those is a question for the Agent Bricks team, since an agent.toml can share them.
• Tracing beyond the happy path: tracing unbind, then tracing bind to a new experiment, each followed by deploy and invoke.
• Sandbox as an App: with --auth app, the sandbox was denied reading its volume in 5 of 7 deploys on tilefood, even with READ_VOLUME, USE_SCHEMA and USE_CATALOG effective for the App's principal. The suite covers only user-auth sandbox until that's fixed.
• Auth defaults disagree: tools add without --auth means user, but an agent.toml entry without auth runs as the App.
• SessionStoreClient.get_session doesn't retry a truncated response, so a first message to a new deploy can return a 500.
• Suite: a framework fixture (langgraph, openai), with a per-framework reader for tool steps; the fake workspace, so the suites can run in PR CI; then switch the nightly runner to these suites and remove tests/e2e/tool_matrix.py.

This pull request and its description were written by Isaac.

[jamesbxwu]

thanks for addressing the comments! new structure is much cleaner. I found a few examples where we are still leaking internal details.

From a test organization PoV, since this is for the CLI, can we have one test file per command group? (e.g. test_init.py test_dev.py, test_deploy, test_tools.py etc). It makes it easier to map from our commands to test files.

[jamesbxwu]

instead of this, can we use agentbricks deployment get to instead which stays at the CLI interface boundary?

[jamesbxwu]

can this be replaced with agentbricks deployments get?

[jamesbxwu]

oh interesting, should deployments delete actually ensures it only returns once it's fully deleted?

[elainewang-db]

yes, I have a follow up pr for deployments delete improvements that I will add this to and drop from the test harness in that PR

[jamesbxwu]

this test is not an e2e testing test to me. IMO it should be the following and we don't need to verify/care about app resources. (note we can break up my example into smaller test cases if needed)

1. init
2. dev
3. tracing get --experiment-name DEFAULT_EXP_FROM_INIT
4. deploy + invoke
5. tracing get
6. unbind
7. deploy + invoke
8. tracing get (verify no tracing)
9. bind --experiment NEW_EXPERIMENT
10. deploy + invoke
11. tracing get --experiment-name NEW_EXPERIMENT

[elainewang-db]

yep, definitely need to improve test coverage across all our resources, will create more in-depth follow ups for tracing as well as session store + memory store

[jamesbxwu]

should we rename to test_tool_selection.py?

[jamesbxwu]

why have different folders for memory_store, session_store, tools and tracing? does it make sense to flatten all into suites? (I noticed deploy/dev tests are under tools folder)

[elainewang-db]

they are the deploy/dev tests for the tools tests

[jamesbxwu]

thanks for addressing the comments, the remaining ones are minor.

Maybe I missed it but did we get a completed live e2e run? if so, can we include it in the PR?

[jamesbxwu]

why do we delete the session/memory store after deploy?

[elainewang-db]

cleanup registers the cli delete commands so that they will run after the test regardless of whether the test succeeds or fails. but they need to be registered after deploy bc thats where we're returning the created session and memory store names

[jamesbxwu]

i see this setup and cleanup being repeated across multiple tests, does it make sense to refactor?

[elainewang-db]

I think as long as the tests stay small and clean, its nice to see exactly what cli commands are run per test in each file, but we can move to a shared test helper if it's too much repeated code, wdyt? I initially had this in a helper but decided to unwrap it so all the cli commands can be read easily in the tests

[jamesbxwu]

[P2] Wait for timed-out commands to exit before teardown

Could we wait for the process group to stop before raising here? With grace=0, terminate_process_group returns immediately after SIGTERM and skips the SIGKILL fallback. A command or child process that delays or ignores termination can continue provisioning while pytest tears resources down, potentially leaving resources after the sweep. A bounded graceful wait, followed by SIGKILL if needed and reaping the process, would make the timeout safe for teardown.

[jamesbxwu]

[P2] Check the Genie payload's completion status

Could we validate the status inside the Genie response as well? GenieAgent.ask and poll return terminal failures (FAILED, CANCELLED, QUERY_RESULT_EXPIRED) and timeout results as ordinary payloads. The outer ToolMessage can therefore have status == "success" and nonempty content even when Genie produced no successful answer, so this assertion passes. Please require a completed ask/poll response with answer content, and handle an in-progress response through polling if that is part of this scenario.