diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index 5efbaedd4..ac50d32aa 100644 --- a/.github/workflows/integration.yml +++ b/.github/workflows/integration.yml @@ -431,6 +431,8 @@ jobs: UV_INDEX_URL: https://pypi.org/simple UG_CUJ_SP_CLIENT_ID: ${{ secrets.UG_CUJ_SP_CLIENT_ID }} UG_CUJ_SP_CLIENT_SECRET: ${{ secrets.UG_CUJ_SP_CLIENT_SECRET }} + UG_BUDGET_CUJ_SP_CLIENT_ID: ${{ secrets.UG_BUDGET_CUJ_SP_CLIENT_ID }} + UG_BUDGET_CUJ_SP_CLIENT_SECRET: ${{ secrets.UG_BUDGET_CUJ_SP_CLIENT_SECRET }} steps: - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 - uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020 # v4 diff --git a/tests/README.md b/tests/README.md index c4450b54d..dcfc7497f 100644 --- a/tests/README.md +++ b/tests/README.md @@ -7,7 +7,7 @@ mocks, monkeypatching, fake binaries/services, or fabricated ug state. | Category | Location | What it proves | | --- | --- | --- | | Unit/component | Existing `test_*.py` files | Individual behavior; dependencies may be mocked | -| Dedicated-workspace CUJ | `e2e_cuj/test_cuj_*.py` | Real user journeys against read-only, preconfigured workspaces | +| Dedicated-workspace CUJ | `e2e_cuj/test_*.py` | Real user journeys against read-only, preconfigured workspaces | | Existing e2e | `test_e2e*.py` | Real workspace behavior with some patched setup/internal calls | | Integration CUJs | `integration/test_*.py` | Public configure, TUI, script, command, protocol, and lifecycle journeys | | Installation | `integration/test_installation.py` | Fresh installed package, CLI, and local helpers without credentials on Linux and advisory native Windows | @@ -181,6 +181,7 @@ integration utilities; only CUJ-specific evidence correlation stays in a test fi | `test_ug_configure_claude_cleans_stale_skills_mcp_on_workspace_switch` | Configure the first workspace, register its skills MCP, switch to a second real workspace, and use Claude | Old registration removed from Claude and the new workspace state; old workspace bucket preserved; repeat configure stays clean; real file task completes on the second workspace | | `test_ug_configure_claude_rejects_invalid_credentials`, `test_ug_configure_codex_rejects_invalid_credentials` | Configure with a rejected bearer against the real workspace | Authentication failure; no successful saved setup | | `test_ug_configure_managed_claude`, `test_ug_configure_managed_codex` | Configure against a workspace that publishes a managed CodingAgentConfig | No agent selector; each agent's generated config exposes exactly the admin's static model_services; real gateway prompt on launch. The Codex case also checks the shared catalog pointer, restart guidance, and a fresh bare app-server's visible model list | +| `e2e_cuj/test_ug_budget_defaults.py` | Launch bare `ug` with separate low-spend and above-tier principals against the fixed 1% tier (`spending_percentage=0.01`); explicitly launch `ug claude` above the tier | Bare launches check Claude/Sonnet below the tier and the Codex/Luna recommendation above it; `ug usage` agrees with backend spend, threshold, and percentage. Explicit `ug claude` displays the backend's Codex/Luna recommendation while its generated model setting and native header select Sonnet. No budget writes or inference tasks. Run in the shared `E2E CUJs` job with both credential pairs documented in `integration/README.md`. | | `test_case_01_*` | Launch managed Claude without defaults after configure and from fresh state | Claude receives the admin MPS header; its gateway cache and replacement picker match the independently fetched provider model IDs; catalog labels are preserved and a model appears in a numbered picker row | | `test_case_03_*`, `test_case_05_*` | Pass a provider or model-location override to managed Claude after configure and from fresh state | ug rejects the override before Claude starts and preserves agent-owned state | | `test_case_02_*` | Launch managed Codex after configure and from fresh state | The scoped and stable catalogs, ug-launched app server, and fresh bare app server match the independently fetched admin MPS model IDs. The configured case uses real `ug revert` to remove ug's shared pointer and stable file while preserving a user setting | diff --git a/tests/e2e_cuj/AGENTS.md b/tests/e2e_cuj/AGENTS.md index 6961d9093..cfdab7bfe 100644 --- a/tests/e2e_cuj/AGENTS.md +++ b/tests/e2e_cuj/AGENTS.md @@ -3,6 +3,8 @@ 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. +A CUJ that needs another identity overrides `CLIENT_ID_ENV` and `CLIENT_SECRET_ENV` +with its credential environment-variable names. The workspaces share one metastore, so metastore-scoped fixture names belong in `helpers/constants.py`; keep each workspace URL in its test class. All workspace configurations are READ ONLY. Never create, publish, update, delete, diff --git a/tests/e2e_cuj/base.py b/tests/e2e_cuj/base.py index a0061649f..cbffca794 100644 --- a/tests/e2e_cuj/base.py +++ b/tests/e2e_cuj/base.py @@ -9,6 +9,8 @@ class BaseCujTest: WORKSPACE_URL: ClassVar[str] = "" + CLIENT_ID_ENV: ClassVar[str] = "UG_CUJ_SP_CLIENT_ID" + CLIENT_SECRET_ENV: ClassVar[str] = "UG_CUJ_SP_CLIENT_SECRET" workspace: WorkspaceClient @pytest.fixture(scope="class", autouse=True) @@ -17,10 +19,12 @@ def setup_workspace(self, request): if not workspace_url: pytest.fail("Set WORKSPACE_URL on your CUJ test class.", pytrace=False) - client_id = os.environ.get("UG_CUJ_SP_CLIENT_ID", "") - client_secret = os.environ.get("UG_CUJ_SP_CLIENT_SECRET", "") + client_id_env = request.cls.CLIENT_ID_ENV + client_secret_env = request.cls.CLIENT_SECRET_ENV + client_id = os.environ.get(client_id_env, "") + client_secret = os.environ.get(client_secret_env, "") if not client_id or not client_secret: - pytest.fail("Set UG_CUJ_SP_CLIENT_ID and UG_CUJ_SP_CLIENT_SECRET.", pytrace=False) + pytest.fail(f"Set {client_id_env} and {client_secret_env}.", pytrace=False) request.cls.workspace = WorkspaceClient( host=workspace_url, diff --git a/tests/e2e_cuj/test_ug_budget_defaults.py b/tests/e2e_cuj/test_ug_budget_defaults.py new file mode 100644 index 000000000..54b6a63fe --- /dev/null +++ b/tests/e2e_cuj/test_ug_budget_defaults.py @@ -0,0 +1,238 @@ +"""Live budget recommendation and default-model checks for a fixed workspace tier.""" + +import json +import re +from decimal import ROUND_HALF_UP, Decimal + +from tests.integration.utils.terminal import AgentTerminal + +from .base import BaseCujTest +from .helpers.constants import CLAUDE, CODEX, CodingAgent + + +class _BudgetDefaultsBase(BaseCujTest): + WORKSPACE_URL = "https://dbc-497adeef-62c0.cloud.databricks.com" + BUDGET_ID = "25c4ce5c-fcd6-4d00-8266-69030cc1a236" + SONNET_MODEL = "system.ai.claude-sonnet-4-6" + SOL_MODEL = "system.ai.gpt-5-6-sol" + LUNA_MODEL = "system.ai.gpt-5-6-luna" + + +class TestCujBudgetDefaultsBelowTier(_BudgetDefaultsBase): + CLIENT_ID_ENV = "UG_BUDGET_CUJ_SP_CLIENT_ID" + CLIENT_SECRET_ENV = "UG_BUDGET_CUJ_SP_CLIENT_SECRET" + + def test_bare_ug_uses_claude_default_below_budget_tier(self, cuj): + """Scenario: launch bare ``ug`` as a fresh principal below the fixed 1% tier. + + Expected: the published config keeps Claude/Sonnet as the default, the real recommendation + also selects Claude/Sonnet, and the real Claude TUI starts on that model without a prompt, + inference request, budget write, or account-login flow. A fresh principal may have no + spend counters yet; when the backend returns an effective threshold, omitted spend is + treated as zero and must remain below the 1% tier. If both figures are absent, this case + verifies default selection without a numeric spend assertion. + """ + session, workspace, _ = cuj + config_path = "/api/ai-gateway/v2/coding-agent-configs" + payload = workspace.client.api_client.do("GET", path=config_path) + config = payload["coding_agent_configs"][0] + assert config["default_agent"] == CodingAgent.CLAUDE_CODE, config + claude = next( + agent for agent in config["enabled_agents"] if agent["agent"] == CodingAgent.CLAUDE_CODE + ) + claude_defaults = claude["config"]["default_models"] + assert claude_defaults["default_model"] == self.SONNET_MODEL, config + assert config["smart_defaults"]["budget_id"] == self.BUDGET_ID, config + assert config["smart_defaults"]["tiers"] == [ + { + "spending_percentage": 0.01, + "recommended_agent": CodingAgent.CODEX, + "recommended_model": self.LUNA_MODEL, + } + ], config + + session.configure( + [ + "configure", + "--workspace", + self.WORKSPACE_URL, + "--skip-upgrade", + "--disable-databricks-ai-tools", + ] + ) + + recommendation = workspace.client.api_client.do( + "POST", path=config_path + ":recommendModel", body={} + ) + assert recommendation.get("recommended_agent") == CodingAgent.CLAUDE_CODE, recommendation + assert recommendation.get("recommended_model") == self.SONNET_MODEL, recommendation + + threshold = recommendation.get("effective_threshold") + spend = recommendation.get("current_spend") + if threshold is not None: + threshold = Decimal(str(threshold)) + assert threshold > 0, recommendation + spend = Decimal("0") if spend is None else Decimal(str(spend)) + assert spend >= 0, recommendation + assert spend / threshold < Decimal("0.01"), recommendation + else: + # A fresh principal can have no usage counter; that does not establish a numeric ratio. + assert spend is None, recommendation + + with AgentTerminal(session, CLAUDE, [str(session.binary)], "budget-default-claude") as tui: + tui.boot(timeout=240) + claude_settings = session.home / ".claude" / "ucode-settings.json" + settings = json.loads(claude_settings.read_text()) + assert settings["env"]["ANTHROPIC_MODEL"] == self.SONNET_MODEL, settings + tui.wait_for( + lambda text: re.search(r"Claude Sonnet 4\.6\s*·", text), + "Claude's native Sonnet model header", + ) + tui.exit_normally() + + +class TestCujBudgetDefaults(_BudgetDefaultsBase): + def test_bare_ug_shows_luna_recommendation_and_starts_codex(self, cuj): + """Scenario: launch bare ``ug`` with the fixed 1% Luna budget tier active. + + Expected: the read-only `usage` command reports the backend's spend and threshold, then + the real backend recommends Codex/Luna, the launch panel displays that recommendation, and + the native Codex TUI starts and exits normally, without budget writes or an inference task. + Applying the recommended Luna model over the managed Sol default is outside this coverage. + """ + session, workspace, _ = cuj + config_path = "/api/ai-gateway/v2/coding-agent-configs" + payload = workspace.client.api_client.do("GET", path=config_path) + config = payload["coding_agent_configs"][0] + assert config["default_agent"] == CodingAgent.CLAUDE_CODE, config + codex_agent = next( + agent for agent in config["enabled_agents"] if agent["agent"] == CodingAgent.CODEX + ) + live_default_model = ( + codex_agent.get("config", {}).get("default_models", {}).get("default_model") + ) + assert live_default_model == self.SOL_MODEL, ( + f"Live Codex config model mismatch: expected {self.SOL_MODEL!r}, " + f"got {live_default_model!r}" + ) + assert config["smart_defaults"]["budget_id"] == self.BUDGET_ID, config + assert config["smart_defaults"]["tiers"] == [ + { + "spending_percentage": 0.01, + "recommended_agent": CodingAgent.CODEX, + "recommended_model": self.LUNA_MODEL, + } + ], config + + session.configure( + [ + "configure", + "--workspace", + self.WORKSPACE_URL, + "--skip-upgrade", + "--disable-databricks-ai-tools", + ] + ) + + before_usage = workspace.client.api_client.do( + "POST", path=config_path + ":recommendModel", body={} + ) + usage = session.run("usage", timeout=120) + after_usage = workspace.client.api_client.do( + "POST", path=config_path + ":recommendModel", body={} + ) + for recommendation in (before_usage, after_usage): + assert recommendation.get("recommended_agent") == CodingAgent.CODEX, recommendation + assert recommendation.get("recommended_model") == self.LUNA_MODEL, recommendation + + match = re.search( + r"Budget spend:\s*\$([\d,]+\.\d{2})\s+of\s+\$([\d,]+\.\d{2})\s+\((\d+)%\)", + session.redact(usage.stdout + usage.stderr), + ) + assert match, usage.stdout + usage.stderr + displayed_spend, displayed_threshold, displayed_percent = match.groups() + displayed_spend = Decimal(displayed_spend.replace(",", "")) + displayed_threshold = Decimal(displayed_threshold.replace(",", "")) + cent = Decimal("0.01") + + before_spend = Decimal(str(before_usage["current_spend"])) + after_spend = Decimal(str(after_usage["current_spend"])) + before_threshold = Decimal(str(before_usage["effective_threshold"])) + after_threshold = Decimal(str(after_usage["effective_threshold"])) + assert before_spend <= after_spend, (before_usage, after_usage) + assert before_threshold == after_threshold, (before_usage, after_usage) + assert displayed_threshold == before_threshold.quantize(cent, rounding=ROUND_HALF_UP) + assert ( + before_spend.quantize(cent, rounding=ROUND_HALF_UP) + <= displayed_spend + <= after_spend.quantize(cent, rounding=ROUND_HALF_UP) + ), (before_usage, after_usage, displayed_spend) + assert ( + int(format(float(before_spend / before_threshold), ".0%").removesuffix("%")) + <= int(displayed_percent) + <= int(format(float(after_spend / after_threshold), ".0%").removesuffix("%")) + ), (before_usage, after_usage, usage) + + with AgentTerminal(session, CODEX, [str(session.binary)], "budget-luna") as tui: + tui.boot(timeout=240) + transcript = session.redact( + "".join(tui.output) + "\n" + tui.visible, + strip_ansi=True, + ) + assert re.search( + rf"Recommended agent is Codex with model {re.escape(self.LUNA_MODEL)}\.", + transcript, + ), ( + "Codex launch did not display the backend recommendation " + f"for {self.LUNA_MODEL!r}:\n{transcript}" + ) + tui.exit_normally() + + +class TestCujBudgetDefaultsExplicitClaude(_BudgetDefaultsBase): + def test_ug_claude_starts_on_sonnet_despite_codex_luna_recommendation(self, cuj): + """Scenario: launch explicit ``ug claude`` for an above-tier principal. + + Expected: the live backend recommends Codex/Luna, and the Claude launch displays + that recommendation while its generated model setting and native header select Sonnet. + No task is submitted; this checks startup selection, not the model used for inference. + """ + session, workspace, _ = cuj + session.configure( + [ + "configure", + "--workspace", + self.WORKSPACE_URL, + "--skip-upgrade", + "--disable-databricks-ai-tools", + ] + ) + recommendation = workspace.client.api_client.do( + "POST", path="/api/ai-gateway/v2/coding-agent-configs:recommendModel", body={} + ) + assert recommendation.get("recommended_agent") == CodingAgent.CODEX, recommendation + assert recommendation.get("recommended_model") == self.LUNA_MODEL, recommendation + + with AgentTerminal( + session, CLAUDE, [str(session.binary), CLAUDE], "explicit-claude" + ) as tui: + tui.boot(timeout=240) + transcript = session.redact( + "".join(tui.output) + "\n" + tui.visible, + strip_ansi=True, + ) + assert re.search( + rf"Recommended agent is Codex with model {re.escape(self.LUNA_MODEL)}\.", + transcript, + ), ( + "Claude launch did not display the backend recommendation " + f"for {self.LUNA_MODEL!r}:\n{transcript}" + ) + claude_settings = session.home / ".claude" / "ucode-settings.json" + settings = json.loads(claude_settings.read_text()) + assert settings["env"]["ANTHROPIC_MODEL"] == self.SONNET_MODEL, settings + tui.wait_for( + lambda text: re.search(r"Claude Sonnet 4\.6\s*·", text), + "Claude's native Sonnet model header", + ) + tui.exit_normally() diff --git a/tests/integration/README.md b/tests/integration/README.md index 542309acc..c6497d503 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -32,6 +32,32 @@ The `smart_defaults` wire schema, legacy `spend_tiers` cache reads, and recommen request gating are covered by unit/component tests listed in `../README.md`. This suite does not yet assert live `recommendModel` request counts for configs with and without tiers. +## Dedicated-workspace CUJ5 + +[`../e2e_cuj/test_ug_budget_defaults.py`](../e2e_cuj/test_ug_budget_defaults.py) runs +three budget cases in the shared `E2E CUJs` job against the fixed 1% tier. The above-tier +cases use `UG_CUJ_SP_CLIENT_ID` / `UG_CUJ_SP_CLIENT_SECRET`; the below-tier case uses +`UG_BUDGET_CUJ_SP_CLIENT_ID` / `UG_BUDGET_CUJ_SP_CLIENT_SECRET`. + +Run from a clean POSIX runner with both credential pairs set: + +```bash +UG_CUJ_SP_CLIENT_ID=... UG_CUJ_SP_CLIENT_SECRET=... \ +UG_BUDGET_CUJ_SP_CLIENT_ID=... UG_BUDGET_CUJ_SP_CLIENT_SECRET=... \ +uv run --with pexpect==4.9.0 --with pyte==0.8.2 \ + pytest --confcutdir=tests/e2e_cuj tests/e2e_cuj -v +``` + +All three cases submit no inference tasks. The above-tier bare-launch case checks +`ug usage`, the Luna recommendation, and Codex startup; it does not assert that Codex +applies Luna. The below-tier case verifies Claude/Sonnet selection and checks the +spend percentage when the backend supplies budget figures. + +The explicit `ug claude` case verifies that startup displays the backend's Codex/Luna +recommendation while its generated `ANTHROPIC_MODEL` setting and native header select +Sonnet, then exits. It does not verify Claude's inference model. All three cases leave +workspace and budget configuration unchanged. + Shared subprocess command resolution is covered by `../test_subprocess_cross_os.py` and enforced by Ruff. These component checks do not establish native Windows coverage for every agent; the Windows journeys below validate their explicitly selected paths. diff --git a/tests/test_cuj_session.py b/tests/test_cuj_session.py index d3beff12f..6da893156 100644 --- a/tests/test_cuj_session.py +++ b/tests/test_cuj_session.py @@ -16,6 +16,8 @@ def test_cuj_session_does_not_inherit_agent_credentials_or_configuration(tmp_pat "CODEX_HOME", "UG_CUJ_SP_CLIENT_ID", "UG_CUJ_SP_CLIENT_SECRET", + "UG_BUDGET_CUJ_SP_CLIENT_ID", + "UG_BUDGET_CUJ_SP_CLIENT_SECRET", ): monkeypatch.setenv(name, "developer-value") session = UserSession(tmp_path, Path("/installed/ug"), tmp_path / "artifacts", "test-token") @@ -27,6 +29,8 @@ def test_cuj_session_does_not_inherit_agent_credentials_or_configuration(tmp_pat assert "UCODE_MANAGED_CONFIG_STUB" not in session.env assert "UG_CUJ_SP_CLIENT_ID" not in session.env assert "UG_CUJ_SP_CLIENT_SECRET" not in session.env + assert "UG_BUDGET_CUJ_SP_CLIENT_ID" not in session.env + assert "UG_BUDGET_CUJ_SP_CLIENT_SECRET" not in session.env assert session.env["CODEX_HOME"] == str(session.home / ".codex") assert not list(session.home.iterdir()), "Fixtures must not manufacture ug/agent state" diff --git a/tests/test_cuj_workspace.py b/tests/test_cuj_workspace.py index 244b75571..fc08baa4f 100644 --- a/tests/test_cuj_workspace.py +++ b/tests/test_cuj_workspace.py @@ -1,11 +1,14 @@ """Offline checks for read-only CUJ workspace configuration and catalog access.""" import copy +from types import SimpleNamespace import pytest from databricks.sdk import WorkspaceClient from databricks.sdk.core import Config +from tests.e2e_cuj import base +from tests.e2e_cuj.base import BaseCujTest from tests.e2e_cuj.helpers.constants import CLAUDE, CODEX from tests.e2e_cuj.helpers.workspace import Workspace @@ -16,6 +19,47 @@ def client(monkeypatch): return WorkspaceClient(host="https://example.test", token="test-token", auth_type="pat") +def test_cuj_workspace_auth_uses_each_class_selected_credentials(monkeypatch): + calls = [] + + class CapturingClient: + def __init__(self, **kwargs): + calls.append(kwargs) + + class Shared(BaseCujTest): + WORKSPACE_URL = "https://shared.example" + + class Budget(BaseCujTest): + WORKSPACE_URL = "https://budget.example" + CLIENT_ID_ENV = "UG_BUDGET_CUJ_SP_CLIENT_ID" + CLIENT_SECRET_ENV = "UG_BUDGET_CUJ_SP_CLIENT_SECRET" + + monkeypatch.setattr(base, "WorkspaceClient", CapturingClient) + monkeypatch.setenv("UG_CUJ_SP_CLIENT_ID", "shared-id") + monkeypatch.setenv("UG_CUJ_SP_CLIENT_SECRET", "shared-secret") + monkeypatch.setenv("UG_BUDGET_CUJ_SP_CLIENT_ID", "budget-id") + monkeypatch.setenv("UG_BUDGET_CUJ_SP_CLIENT_SECRET", "budget-secret") + + setup_workspace = BaseCujTest.setup_workspace.__wrapped__ + setup_workspace(None, SimpleNamespace(cls=Shared)) + setup_workspace(None, SimpleNamespace(cls=Budget)) + + assert calls == [ + { + "host": "https://shared.example", + "client_id": "shared-id", + "client_secret": "shared-secret", + "auth_type": "oauth-m2m", + }, + { + "host": "https://budget.example", + "client_id": "budget-id", + "client_secret": "budget-secret", + "auth_type": "oauth-m2m", + }, + ] + + def test_cuj_catalog_rejects_unknown_agent_before_network_access(client, monkeypatch): workspace = Workspace(client)