Run dedicated CUJs through installed agent suite

scripts/run_integration.py

@@ -57,6 +57,12 @@
     "codex": "test_ug_codex_headless.py::test_ug_codex_headless_prompt_argument",
     "opencode": "test_ug_opencode_headless.py::test_ug_opencode_headless_prompt_argument",
 }
+CUJ_TEST_ENVIRONMENT = (
+    "DATABRICKS_BEARER",
+    "UG_CUJ_SP_CLIENT_ID",
+    "UG_CUJ_SP_CLIENT_SECRET",
+    "UG_CUJ1_WORKSPACE",
+)
 
 
 def installer_environment(
@@ -97,6 +103,17 @@ def redact_secrets(value: str, secrets: Iterable[str]) -> str:
     return value
 
 
+def suite_test_environment(
+    suite: str,
+    runtime_environment: Mapping[str, str],
+    source_environment: Mapping[str, str],
+) -> dict[str, str]:
+    environment = dict(runtime_environment)
+    if suite == "e2e-cuj":
+        environment.update({key: source_environment.get(key, "") for key in CUJ_TEST_ENVIRONMENT})
+    return environment
+
+
 def managed_policy_paths(
     agents: Iterable[str],
     environment: Mapping[str, str] | None = None,
@@ -279,6 +296,12 @@ def arguments(
     platform_name = os.name if platform_name is None else platform_name
     environment = os.environ if environment is None else environment
     parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--suite",
+        choices=["integration", "e2e-cuj"],
+        default="integration",
+        help="CUJ tests declare their own workspace and use shared service-principal credentials.",
+    )
     source = parser.add_mutually_exclusive_group()
     source.add_argument("--ug-version", default="checkout", help="Exact ug release, or checkout.")
     source.add_argument(
@@ -354,7 +377,7 @@ def arguments(
     )
     parser.add_argument("--npm-registry", default="https://registry.npmjs.org")
     parser.add_argument("--profile", help="Explicit Databricks profile to mint the live bearer.")
-    parser.add_argument("--workspace", default=environment.get("UCODE_TEST_WORKSPACE"))
+    parser.add_argument("--workspace")
     parser.add_argument(
         "--second-workspace",
         default=environment.get("UCODE_TEST_SECOND_WORKSPACE"),
@@ -373,6 +396,14 @@ def arguments(
         "pytest_args", nargs=argparse.REMAINDER, help="After --, pass pytest filters."
     )
     args = parser.parse_args(argv)
+    if args.suite == "integration" and args.workspace is None:
+        args.workspace = environment.get("UCODE_TEST_WORKSPACE")
+    if args.suite == "e2e-cuj" and (args.workspace or args.profile):
+        parser.error(
+            "CUJ tests declare their workspace and auth; do not use --workspace/--profile."
+        )
+    if args.suite == "e2e-cuj" and (args.installation_only or args.headless_only):
+        parser.error("--installation-only and --headless-only belong to the integration suite.")
     if platform_name != "posix" and not (args.installation_only or args.headless_only):
         parser.error(
             "Live agent/TUI integration requires POSIX PTY, managed-settings, and signal "
@@ -387,7 +418,7 @@ def arguments(
     filters.add_argument("--maxfail", type=int)
     extra = args.pytest_args[1:] if args.pytest_args[:1] == ["--"] else args.pytest_args
     selected = filters.parse_args(extra)
-    marker = selected.m or "live"
+    marker = selected.m or ("cuj" if args.suite == "e2e-cuj" else "live")
     if args.installation_only:
         marker = f"installation and ({selected.m})" if selected.m else "installation"
     args.pytest_args = []
@@ -409,7 +440,7 @@ def arguments(
     for dependency in args.dependency:
         if not re.fullmatch(r"[A-Za-z0-9_.-]+==[A-Za-z0-9_.!+-]+", dependency):
             parser.error("--dependency requires an exact PACKAGE==VERSION constraint.")
-    if not args.installation_only:
+    if not args.installation_only and args.suite == "integration":
         if not args.workspace or not args.workspace.startswith("https://"):
             parser.error(
                 "Set UCODE_TEST_WORKSPACE to the existing e2e workspace, or use --workspace."
@@ -523,6 +554,7 @@ def terminate(signum, frame):
     target_bearers: dict[str, str] = {}
     client_secrets = (
         os.environ.get("DATABRICKS_CLIENT_SECRET", ""),
+        os.environ.get("UG_CUJ_SP_CLIENT_SECRET", ""),
         os.environ.get("UG_MPS_DEFAULTS_CLIENT_SECRET", ""),
         os.environ.get("UG_PARENT_SCHEMA_DEFAULTS_CLIENT_SECRET", ""),
     )
@@ -591,6 +623,7 @@ def run(command, *, cwd=output, env=base_env, timeout=600) -> str:
             "warehouse_id": args.warehouse_id,
         },
         "platform": platform.platform(),
+        "suite": args.suite,
         "installation_only": args.installation_only,
         "headless_only": args.headless_only,
     }
@@ -797,7 +830,7 @@ def run(command, *, cwd=output, env=base_env, timeout=600) -> str:
                 raise RuntimeError(f"Expected {agent} {expected}, got {version!r}")
             report["agents"][agent] = version
 
-        if args.profile and not args.installation_only:
+        if args.profile and not args.installation_only and args.suite == "integration":
             # Never select a local profile implicitly. Do not persist auth output.
             with managed_process(
                 [
@@ -825,13 +858,18 @@ def run(command, *, cwd=output, env=base_env, timeout=600) -> str:
             if not bearer:
                 raise RuntimeError("Selected profile returned no access token.")
 
-        if not bearer and not args.profile and not args.installation_only:
+        if (
+            not bearer
+            and not args.profile
+            and not args.installation_only
+            and args.suite == "integration"
+        ):
             client_id = os.environ.get("DATABRICKS_CLIENT_ID", "").strip()
             client_secret = os.environ.get("DATABRICKS_CLIENT_SECRET", "").strip()
             if client_id and client_secret:
                 bearer = mint_m2m_token(args.workspace, client_id, client_secret)
 
-        if not args.installation_only:
+        if not args.installation_only and args.suite == "integration":
             for bearer_env, target_workspace, client_id, secret_env in MANAGED_DEFAULTS_TARGETS:
                 secret = os.environ.get(secret_env, "").strip()
                 if args.workspace.rstrip("/") == target_workspace:
@@ -840,6 +878,8 @@ def run(command, *, cwd=output, env=base_env, timeout=600) -> str:
                     target_bearers[bearer_env] = mint_m2m_token(target_workspace, client_id, secret)
 
         test_dependencies = ["pytest==9.0.3"]
+        if args.suite == "e2e-cuj":
+            test_dependencies.append("databricks-sdk==0.135.0")
         if os.name == "posix":
             test_dependencies.extend(["pexpect==4.9.0", "pyte==0.8.2"])
         run(
@@ -865,33 +905,46 @@ def run(command, *, cwd=output, env=base_env, timeout=600) -> str:
                 "UG_INTEGRATION_AGENTS": ",".join(agents),
                 "UG_INTEGRATION_CLAUDE_PROVIDER": args.claude_provider,
                 "UG_INTEGRATION_CLAUDE_RELAYED_PROVIDER": args.claude_relayed_provider,
-                "UG_INTEGRATION_CLAUDE_OAUTH_TOKEN": oauth_token,
                 "UG_INTEGRATION_CLAUDE_PROVIDER_MODEL": args.claude_provider_model,
                 "UG_INTEGRATION_CLAUDE_BEDROCK_ALLOW_ALL_PROVIDER": args.claude_bedrock_allow_all_provider,
                 "UG_INTEGRATION_CLAUDE_BEDROCK_ALLOW_ALL_MODEL": args.claude_bedrock_allow_all_model,
                 "UG_INTEGRATION_CODEX_PROVIDER": args.codex_provider,
                 "UG_INTEGRATION_CODEX_PROVIDER_MODEL": args.codex_provider_model,
+                "UG_INTEGRATION_WAREHOUSE_ID": args.warehouse_id or "",
                 "UG_INTEGRATION_PARENT_SCHEMA": args.parent_schema,
                 "UG_INTEGRATION_CLAUDE_PARENT_MODEL": args.claude_parent_model,
                 "UG_INTEGRATION_CODEX_PARENT_MODEL": args.codex_parent_model,
-                "UCODE_TEST_WORKSPACE": args.workspace or "",
-                "DATABRICKS_BEARER": bearer,
-                "UG_MPS_DEFAULTS_BEARER": target_bearers.get("UG_MPS_DEFAULTS_BEARER", ""),
-                "UG_PARENT_SCHEMA_DEFAULTS_BEARER": target_bearers.get(
-                    "UG_PARENT_SCHEMA_DEFAULTS_BEARER", ""
-                ),
-                "UCODE_TEST_SECOND_WORKSPACE": args.second_workspace or "",
-                "DATABRICKS_SECOND_BEARER": second_bearer,
-                "UG_INTEGRATION_WAREHOUSE_ID": args.warehouse_id or "",
             }
         )
+        if args.suite == "integration":
+            runtime_env.update(
+                {
+                    "UCODE_TEST_WORKSPACE": args.workspace or "",
+                    "DATABRICKS_BEARER": bearer,
+                    "UG_MPS_DEFAULTS_BEARER": target_bearers.get("UG_MPS_DEFAULTS_BEARER", ""),
+                    "UG_PARENT_SCHEMA_DEFAULTS_BEARER": target_bearers.get(
+                        "UG_PARENT_SCHEMA_DEFAULTS_BEARER", ""
+                    ),
+                    "UCODE_TEST_SECOND_WORKSPACE": args.second_workspace or "",
+                    "DATABRICKS_SECOND_BEARER": second_bearer,
+                    "UG_INTEGRATION_CLAUDE_OAUTH_TOKEN": oauth_token,
+                }
+            )
         for agent in agents:
             runtime_env[f"UG_INTEGRATION_{agent.upper()}_MODEL"] = (
                 getattr(args, f"{agent}_model") or ""
             )
-        suite = ROOT / "tests/integration"
+        suite = ROOT / "tests" / args.suite.replace("-", "_")
         suite_hash = hashlib.sha256()
-        for path in [Path(__file__), *sorted(suite.rglob("*.py")), suite / "pytest.ini"]:
+        sources = [Path(__file__), *sorted(suite.rglob("*.py")), suite / "pytest.ini"]
+        if args.suite == "e2e-cuj":
+            sources.extend(
+                [
+                    ROOT / "tests/integration/conftest.py",
+                    *sorted((ROOT / "tests/integration/utils").rglob("*.py")),
+                ]
+            )
+        for path in sources:
             suite_hash.update(str(path.relative_to(ROOT)).encode() + b"\0" + path.read_bytes())
         report["suite_sha256"] = suite_hash.hexdigest()
         extra = args.pytest_args
@@ -905,6 +958,7 @@ def run(command, *, cwd=output, env=base_env, timeout=600) -> str:
             installation_only=args.installation_only,
             headless_only=args.headless_only,
         )
+        test_env = suite_test_environment(args.suite, runtime_env, os.environ)
         with managed_process(
             [
                 test_python,
@@ -920,7 +974,7 @@ def run(command, *, cwd=output, env=base_env, timeout=600) -> str:
                 f"--junitxml={output / 'junit.xml'}",
                 *extra,
             ],
-            env=runtime_env,
+            env=test_env,
             cwd=output,
             stdin=subprocess.DEVNULL,
             interrupt=True,

tests/e2e_cuj/AGENTS.md

@@ -3,8 +3,10 @@
 Each CUJ owns a separate workspace. Subclass `BaseCujTest` from `base.py` and set
 `WORKSPACE_URL`. Setup provides `self.workspace`, a Databricks SDK client using
 `UG_CUJ_SP_CLIENT_ID` and `UG_CUJ_SP_CLIENT_SECRET` with OAuth M2M authentication.
+Local runs may instead use an explicitly supplied `DATABRICKS_BEARER`; partial SP credentials fail.
 Do not share the workspace between concurrent runs. No stubbed configuration.
 
-Run from the repository root:
-`uv run pytest --confcutdir=tests/e2e_cuj tests/e2e_cuj`.
-This keeps the parent suite's mocked fixtures out of CUJs. There are no tests yet.
+Run from the repository root with `scripts/run_integration.py --suite e2e-cuj` to install
+ug and pinned agent versions and keep the parent suite's mocked fixtures out of CUJs.
+Set `UG_CUJ_SP_CLIENT_ID` and `UG_CUJ_SP_CLIENT_SECRET`; local runs can instead supply
+`DATABRICKS_BEARER`. CUJ1 also accepts `UG_CUJ1_WORKSPACE` to verify its assigned URL.

tests/e2e_cuj/base.py

@@ -1,4 +1,4 @@
-"""Base class for CUJs with a dedicated workspace and service-principal auth."""
+"""Base class for CUJs with a dedicated workspace."""
 
 import os
 from typing import ClassVar
@@ -18,12 +18,20 @@ def setup_workspace(self):
 
         client_id = os.environ.get("UG_CUJ_SP_CLIENT_ID", "")
         client_secret = os.environ.get("UG_CUJ_SP_CLIENT_SECRET", "")
-        if not client_id or not client_secret:
-            pytest.fail("Set UG_CUJ_SP_CLIENT_ID and UG_CUJ_SP_CLIENT_SECRET.", pytrace=False)
-
-        self.workspace = WorkspaceClient(
-            host=self.WORKSPACE_URL,
-            client_id=client_id,
-            client_secret=client_secret,
-            auth_type="oauth-m2m",
-        )
+        if bool(client_id) != bool(client_secret):
+            pytest.fail("Set both UG_CUJ_SP_CLIENT_ID and UG_CUJ_SP_CLIENT_SECRET.", pytrace=False)
+        if client_id:
+            self.workspace = WorkspaceClient(
+                host=self.WORKSPACE_URL,
+                client_id=client_id,
+                client_secret=client_secret,
+                auth_type="oauth-m2m",
+            )
+        elif bearer := os.environ.get("DATABRICKS_BEARER", "").strip():
+            self.workspace = WorkspaceClient(host=self.WORKSPACE_URL, token=bearer)
+        else:
+            pytest.fail(
+                "Set UG_CUJ_SP_CLIENT_ID and UG_CUJ_SP_CLIENT_SECRET, "
+                "or DATABRICKS_BEARER for a local run.",
+                pytrace=False,
+            )

tests/e2e_cuj/conftest.py

@@ -0,0 +1,36 @@
+"""Fixtures for dedicated-workspace CUJs."""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import pytest
+
+_INTEGRATION = Path(__file__).parents[1] / "integration"
+_ROOT = _INTEGRATION.parents[1]
+_SUITE = Path(__file__).parent
+for _path in (_ROOT, _INTEGRATION, _SUITE):
+    if str(_path) not in sys.path:
+        sys.path.insert(0, str(_path))
+
+from base import BaseCujTest  # noqa: E402
+
+from tests.integration.conftest import installed_binary as installed_binary  # noqa: E402,F401
+from tests.integration.conftest import session as _integration_session  # noqa: E402
+
+session = _integration_session
+
+
+@pytest.fixture
+def live_session(request, session):
+    test = request.instance
+    if not isinstance(test, BaseCujTest):
+        pytest.fail("live_session requires a BaseCujTest instance.", pytrace=False)
+
+    authorization = test.workspace.config.authenticate().get("Authorization", "")
+    bearer = authorization.removeprefix("Bearer ").strip()
+    if not bearer:
+        pytest.fail("Workspace authentication returned no bearer token.", pytrace=False)
+    session.env["DATABRICKS_BEARER"] = bearer
+    return session

tests/e2e_cuj/pytest.ini

@@ -0,0 +1,5 @@
+[pytest]
+testpaths = .
+addopts = --strict-markers --tb=short
+markers =
+    cuj: dedicated-workspace user journey

tests/integration/utils/terminal.py

@@ -232,6 +232,13 @@ def boot(self, timeout=120):
             # These are interactions with ordinary UI choices, not pre-written
             # onboarding state. Only the test's disposable project is trusted.
             dialogs = [
+                (
+                    "external-imports",
+                    "Allow external CLAUDE.md file imports?" in text
+                    and "No, disable external imports" in text
+                    and "Yes, allow external imports" in text,
+                    "\r",
+                ),
                 (
                     "theme",
                     "Choose the text style" in text and "Dark mode" in text,

tests/test_e2e_cuj_base.py

@@ -0,0 +1,31 @@
+"""Authentication boundaries for dedicated CUJ workspaces."""
+
+import pytest
+
+from tests.e2e_cuj import base
+from tests.e2e_cuj.base import BaseCujTest
+
+
+class _Cuj(BaseCujTest):
+    WORKSPACE_URL = "https://example.cloud.databricks.com"
+
+
+def test_local_run_uses_explicit_bearer(monkeypatch):
+    monkeypatch.delenv("UG_CUJ_SP_CLIENT_ID", raising=False)
+    monkeypatch.delenv("UG_CUJ_SP_CLIENT_SECRET", raising=False)
+    monkeypatch.setenv("DATABRICKS_BEARER", "local-test-token")
+    created = []
+    monkeypatch.setattr(base, "WorkspaceClient", lambda **kwargs: created.append(kwargs))
+
+    BaseCujTest.setup_workspace.__wrapped__(_Cuj())
+
+    assert created == [{"host": _Cuj.WORKSPACE_URL, "token": "local-test-token"}]
+
+
+def test_partial_service_principal_credentials_fail_even_with_local_bearer(monkeypatch):
+    monkeypatch.setenv("UG_CUJ_SP_CLIENT_ID", "incomplete-client-id")
+    monkeypatch.delenv("UG_CUJ_SP_CLIENT_SECRET", raising=False)
+    monkeypatch.setenv("DATABRICKS_BEARER", "local-test-token")
+
+    with pytest.raises(pytest.fail.Exception, match="Set both UG_CUJ_SP_CLIENT_ID"):
+        BaseCujTest.setup_workspace.__wrapped__(_Cuj())