From f258ae5c9537e5d4a29e621967bdb9084877af89 Mon Sep 17 00:00:00 2001 From: Andy Xu Date: Fri, 2 Oct 2026 13:31:39 +0000 Subject: [PATCH] [AIGTWY-4880] Share integration credentials with budget tests --- pyproject.toml | 1 + scripts/run_integration.py | 75 +++++- tests/integration/utils/harness.py | 8 +- tests/test_integration_runner.py | 362 +++++++++++++++++++++++++++++ 4 files changed, 433 insertions(+), 13 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 94543bd6c..2e8dc7973 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -84,6 +84,7 @@ ignore = [ [tool.ruff.lint.per-file-ignores] "tests/**" = ["B011", "TID251"] # tests may exercise subprocess directly +"scripts/run_integration.py" = ["TID251"] # standalone runner runs before the application is installed "src/ucode/os_compatibility/file_lock_cross_os.py" = ["TID251"] "src/ucode/os_compatibility/subprocess_cross_os.py" = ["TID251"] diff --git a/scripts/run_integration.py b/scripts/run_integration.py index bc81650d3..c915dca40 100644 --- a/scripts/run_integration.py +++ b/scripts/run_integration.py @@ -45,6 +45,7 @@ "UG_PARENT_SCHEMA_DEFAULTS_CLIENT_SECRET", ), ) + WINDOWS_PATHEXT = ".COM;.EXE;.BAT;.CMD" UV_INDEX_CREDENTIAL_ENV = ( "UV_INDEX_DATABRICKS_PYPI_USERNAME", @@ -196,7 +197,7 @@ def npm_executable(bin_dir: Path, name: str) -> Path: def mint_m2m_token(workspace: str, client_id: str, client_secret: str) -> str: - """Mint a short-lived workspace token for a service principal via OAuth client credentials. + """Mint a short-lived workspace token via OAuth client credentials. Managed-workspace M2M tokens expire hourly, so the runner mints them from client credentials rather than storing long-lived bearers for the base or Claude defaults workspaces. @@ -206,7 +207,7 @@ def mint_m2m_token(workspace: str, client_id: str, client_secret: str) -> str: {"grant_type": "client_credentials", "scope": "all-apis"} ).encode() request = urllib.request.Request( - f"{workspace.rstrip('/')}/oidc/v1/token", + workspace.rstrip("/") + "/oidc/v1/token", data=body, headers={ "Authorization": f"Basic {basic}", @@ -220,9 +221,34 @@ def mint_m2m_token(workspace: str, client_id: str, client_secret: str) -> str: return token +def mint_account_m2m_token( + account_host: str, client_id: str, client_secret: str, *, account_id: str +) -> str: + """Mint a short-lived account token for the account-scoped budget API.""" + basic = base64.b64encode(f"{client_id}:{client_secret}".encode()).decode() + body = urllib.parse.urlencode( + {"grant_type": "client_credentials", "scope": "all-apis"} + ).encode() + request = urllib.request.Request( + account_host.rstrip("/") + + f"/oidc/accounts/{urllib.parse.quote(account_id, safe='')}/v1/token", + data=body, + headers={ + "Authorization": f"Basic {basic}", + "Content-Type": "application/x-www-form-urlencoded", + }, + ) + with urllib.request.urlopen(request, timeout=30) as response: # noqa: S310 (https account URL) + token = json.load(response).get("access_token", "") + if not token: + raise RuntimeError("Service-principal client credentials returned no access token.") + return token + + @contextlib.contextmanager -def managed_process(command, *, interrupt=False, **kwargs): +def managed_process(command, *, interrupt=False, interrupt_grace=15, **kwargs): """Bound child lifetimes, including descendants that outlive their parent.""" + # The black-box installer runs before the application is available to import. proc = subprocess.Popen(command, **process_group_options(), **kwargs) try: yield proc @@ -234,7 +260,7 @@ def managed_process(command, *, interrupt=False, **kwargs): with contextlib.suppress(ProcessLookupError): os.killpg(proc.pid, first_signal) try: - proc.wait(timeout=15 if interrupt else 5) + proc.wait(timeout=interrupt_grace if interrupt else 5) except subprocess.TimeoutExpired: pass with contextlib.suppress(ProcessLookupError): @@ -521,10 +547,17 @@ def terminate(signum, frame): second_bearer = os.environ.get("DATABRICKS_SECOND_BEARER", "").strip() oauth_token = os.environ.get("CLAUDE_CODE_OAUTH_TOKEN", "").strip() target_bearers: dict[str, str] = {} + budget_environment = { + key: value.strip() + for key, value in os.environ.items() + if key.startswith("UG_BUDGET_") + and not key.endswith(("_TOKEN", "_SECRET")) + and value.strip() + } client_secrets = ( - os.environ.get("DATABRICKS_CLIENT_SECRET", ""), - os.environ.get("UG_MPS_DEFAULTS_CLIENT_SECRET", ""), - os.environ.get("UG_PARENT_SCHEMA_DEFAULTS_CLIENT_SECRET", ""), + os.environ.get("DATABRICKS_CLIENT_SECRET", "").strip(), + os.environ.get("UG_MPS_DEFAULTS_CLIENT_SECRET", "").strip(), + os.environ.get("UG_PARENT_SCHEMA_DEFAULTS_CLIENT_SECRET", "").strip(), ) def redact(value: str) -> str: @@ -535,6 +568,7 @@ def redact(value: str) -> str: second_bearer, oauth_token, *target_bearers.values(), + budget_environment.get("UG_BUDGET_ACCOUNT_TOKEN", ""), *client_secrets, *installer_secrets, ), @@ -825,13 +859,22 @@ 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 args.installation_only: client_id = os.environ.get("DATABRICKS_CLIENT_ID", "").strip() client_secret = os.environ.get("DATABRICKS_CLIENT_SECRET", "").strip() - if client_id and client_secret: + if not bearer and not args.profile and client_id and client_secret: bearer = mint_m2m_token(args.workspace, client_id, client_secret) - if not args.installation_only: + if budget_environment: + account_host = budget_environment.get("UG_BUDGET_ACCOUNT_HOST", "") + account_id = budget_environment.get("UG_BUDGET_ACCOUNT_ID", "") + if account_host and account_id and client_id and client_secret: + budget_environment["UG_BUDGET_ACCOUNT_TOKEN"] = mint_account_m2m_token( + account_host, + client_id, + client_secret, + account_id=account_id, + ) 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: @@ -889,6 +932,9 @@ def run(command, *, cwd=output, env=base_env, timeout=600) -> str: runtime_env[f"UG_INTEGRATION_{agent.upper()}_MODEL"] = ( getattr(args, f"{agent}_model") or "" ) + # The workflow target metadata and account token are needed by the CUJ5 + # fixture, but must never reach native agent/version processes. + test_env = {**runtime_env, **budget_environment} suite = ROOT / "tests/integration" suite_hash = hashlib.sha256() for path in [Path(__file__), *sorted(suite.rglob("*.py")), suite / "pytest.ini"]: @@ -920,12 +966,17 @@ 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, + # Budget cleanup can require several bounded account API calls. Give + # its finally block time to restore before killing the process group. + interrupt_grace=180 if budget_environment else 15, ) as result: - result.wait(timeout=3600) + # Fresh workspace and account M2M tokens expire hourly. Reserve time + # for budget restoration while those tokens are still valid. + result.wait(timeout=2400 if budget_environment else 3600) exitcode = result.returncode junit = output / "junit.xml" if junit.is_file(): diff --git a/tests/integration/utils/harness.py b/tests/integration/utils/harness.py index b26a9528b..5f8802d81 100644 --- a/tests/integration/utils/harness.py +++ b/tests/integration/utils/harness.py @@ -109,7 +109,13 @@ def __init__(self, root: Path, project_root: Path, binary: Path, artifacts: Path def redact(self, text: str, *, strip_ansi: bool = True) -> str: # Also scrub the relayed launch's subscription OAuth token, not just the bearer. - for name in ("DATABRICKS_BEARER", "DATABRICKS_SECOND_BEARER", "CLAUDE_CODE_OAUTH_TOKEN"): + for name in ( + "DATABRICKS_BEARER", + "DATABRICKS_SECOND_BEARER", + "DATABRICKS_CLIENT_SECRET", + "CLAUDE_CODE_OAUTH_TOKEN", + "UG_BUDGET_ACCOUNT_TOKEN", + ): for token in (os.environ.get(name), self.env.get(name)): if token: text = text.replace(token, "") diff --git a/tests/test_integration_runner.py b/tests/test_integration_runner.py index 49f749c95..7779bf9e1 100644 --- a/tests/test_integration_runner.py +++ b/tests/test_integration_runner.py @@ -1,8 +1,370 @@ +import base64 +import json +import urllib.error +import urllib.parse +from contextlib import contextmanager from pathlib import Path +from unittest.mock import call, patch import pytest from scripts import run_integration as runner +from tests.integration.utils.harness import UserSession, clean_environment + +ACCOUNT_ID = "11111111-2222-3333-4444-555555555555" +CUJ5_TARGET_ENV = { + "UCODE_TEST_WORKSPACE": "https://cuj5-workspace.example", + "UG_BUDGET_WORKSPACE_ID": "workspace-id", + "UG_BUDGET_ACCOUNT_HOST": "https://accounts.example", + "UG_BUDGET_ACCOUNT_ID": "account-id", + "UG_BUDGET_ID": "budget-id", + "UG_BUDGET_SONNET_MODEL": "sonnet-model", + "UG_BUDGET_SOL_MODEL": "sol-model", + "UG_BUDGET_LUNA_MODEL": "luna-model", +} + + +def test_mint_m2m_token_uses_workspace_token_endpoint(): + with patch.object(runner.urllib.request, "urlopen") as request: + request.return_value.__enter__.return_value.read.return_value = ( + b'{"access_token":"workspace-token"}' + ) + token = runner.mint_m2m_token( + "https://workspace.example/", + "client", + "secret", + ) + assert token == "workspace-token" + token_request = request.call_args.args[0] + assert token_request.full_url == "https://workspace.example/oidc/v1/token" + assert token_request.get_method() == "POST" + + +def test_mint_account_m2m_token_uses_account_token_endpoint(): + with patch.object(runner.urllib.request, "urlopen") as request: + request.return_value.__enter__.return_value.read.return_value = ( + b'{"access_token":"account-token"}' + ) + token = runner.mint_account_m2m_token( + "https://accounts.cloud.databricks.com", + "client", + "secret", + account_id=ACCOUNT_ID, + ) + assert token == "account-token" + token_request = request.call_args.args[0] + assert token_request.full_url == ( + f"https://accounts.cloud.databricks.com/oidc/accounts/{ACCOUNT_ID}/v1/token" + ) + assert token_request.get_method() == "POST" + assert ( + token_request.headers["Authorization"] + == "Basic " + base64.b64encode(b"client:secret").decode() + ) + assert token_request.headers["Content-type"] == "application/x-www-form-urlencoded" + assert urllib.parse.parse_qs(token_request.data.decode()) == { + "grant_type": ["client_credentials"], + "scope": ["all-apis"], + } + + +def test_mint_account_m2m_token_quotes_account_id_path_segment(): + with patch.object(runner.urllib.request, "urlopen") as request: + request.return_value.__enter__.return_value.read.return_value = ( + b'{"access_token":"account-token"}' + ) + token = runner.mint_account_m2m_token( + "https://accounts.cloud.databricks.com", + "client", + "secret", + account_id="account/id?query", + ) + assert token == "account-token" + assert request.call_args.args[0].full_url == ( + "https://accounts.cloud.databricks.com/oidc/accounts/account%2Fid%3Fquery/v1/token" + ) + + +@pytest.mark.skipif(runner.os.name != "posix", reason="POSIX process-group interrupt handling") +def test_budget_cleanup_receives_its_interrupt_grace_period(): + with ( + patch.object(runner.subprocess, "Popen") as start, + patch.object(runner.os, "killpg"), + ): + with runner.managed_process(["pytest"], interrupt=True, interrupt_grace=180): + pass + assert start.return_value.wait.call_args_list == [call(timeout=180), call(timeout=5)] + + +class _FakeProcess: + returncode = 0 + + def __init__(self, command): + self.command = command + + def communicate(self, timeout=None): + if "-c" in self.command: + runtime = Path(self.command[0]).parent.parent + output = json.dumps( + { + "distribution": "unity-gateway", + "version": "test", + "path": str(runtime / "site-packages/ucode/__init__.py"), + } + ) + elif "ls" in self.command and "--json" in self.command: + output = '{"dependencies":{}}' + elif "--version" in self.command: + output = "1.2.3" + else: + output = "" + return output, "" + + def wait(self, timeout=None): + return self.returncode + + +def _main_setup( + tmp_path, + monkeypatch, + *, + workspace, + client_id, + client_secret, + target_environment=None, +): + wheel = tmp_path / "unity_gateway.whl" + wheel.write_bytes(b"wheel") + npm_lock = tmp_path / "npm-lock.json" + npm_lock.write_text("{}") + arguments_environment = { + "DATABRICKS_CLIENT_ID": client_id, + "DATABRICKS_CLIENT_SECRET": client_secret, + } + args = runner.arguments( + [ + "--ug-wheel", + str(wheel), + "--npm-lock", + str(npm_lock), + "--claude-version", + "1.2.3", + "--headless-only", + "--workspace", + workspace, + ], + environment=arguments_environment, + ) + args.output = tmp_path / "run" + monkeypatch.setattr(runner, "arguments", lambda: args) + monkeypatch.setattr(runner.signal, "signal", lambda *ignored: None) + monkeypatch.setattr(runner, "present_policy_paths", lambda paths: []) + + binaries = {} + for name in ("uv", "npm", "node", "databricks"): + binary = tmp_path / name + binary.write_text("") + binaries[name] = binary + monkeypatch.setattr(runner.shutil, "which", lambda name: str(binaries[name])) + + def fake_venv_executable(environment_path, name): + executable = Path(environment_path) / "bin" / name + executable.parent.mkdir(parents=True, exist_ok=True) + executable.touch() + return executable + + monkeypatch.setattr(runner, "venv_executable", fake_venv_executable) + monkeypatch.setattr(runner, "npm_executable", lambda bin_dir, name: Path(bin_dir) / name) + + mint_calls = [] + account_calls = [] + + def fake_mint(workspace, client_id, client_secret): + mint_calls.append((workspace, client_id, client_secret)) + return "workspace-token" + + monkeypatch.setattr(runner, "mint_m2m_token", fake_mint) + + def fake_account_token(account_host, client_id, client_secret, *, account_id): + account_calls.append((account_host, client_id, client_secret, account_id)) + return "account-token" + + monkeypatch.setattr(runner, "mint_account_m2m_token", fake_account_token) + process_calls = [] + + @contextmanager + def fake_managed_process(command, **kwargs): + command = [str(item) for item in command] + process_calls.append((command, kwargs)) + junit = next((item for item in command if item.startswith("--junitxml=")), None) + if junit: + Path(junit.split("=", 1)[1]).write_text( + '' + "" + ) + yield _FakeProcess(command) + + monkeypatch.setattr(runner, "managed_process", fake_managed_process) + + host_environment = { + key: runner.os.environ[key] + for key in ( + "PATH", + "SYSTEMROOT", + "WINDIR", + "COMSPEC", + "PATHEXT", + "PROGRAMDATA", + "PROGRAMFILES", + ) + if key in runner.os.environ + } + environment = { + **host_environment, + **(target_environment or {}), + "DATABRICKS_CLIENT_ID": client_id, + } + environment["DATABRICKS_CLIENT_SECRET"] = client_secret + return args, environment, mint_calls, account_calls, process_calls + + +def _pytest_call(process_calls): + pytest_calls = [ + (command, kwargs) + for command, kwargs in process_calls + if "-m" in command and "pytest" in command + ] + assert len(pytest_calls) == 1 + return pytest_calls[0][1] + + +def test_main_forwards_budget_target_into_pytest_only(tmp_path, monkeypatch): + workspace = CUJ5_TARGET_ENV["UCODE_TEST_WORKSPACE"] + args, environment, mint_calls, account_calls, process_calls = _main_setup( + tmp_path, + monkeypatch, + workspace=workspace, + client_id="cuj-client", + client_secret="cuj-secret", + target_environment=CUJ5_TARGET_ENV, + ) + environment["UG_BUDGET_FUTURE_METADATA"] = "future-value" + environment["UG_BUDGET_CALLER_TOKEN"] = "caller-token" + environment["UG_BUDGET_CALLER_SECRET"] = "caller-secret" + with patch.dict(runner.os.environ, environment, clear=True): + assert runner.main() == 0 + + assert mint_calls == [(workspace, "cuj-client", "cuj-secret")] + assert account_calls == [ + ( + CUJ5_TARGET_ENV["UG_BUDGET_ACCOUNT_HOST"], + "cuj-client", + "cuj-secret", + CUJ5_TARGET_ENV["UG_BUDGET_ACCOUNT_ID"], + ), + ] + pytest_kwargs = _pytest_call(process_calls) + pytest_environment = pytest_kwargs["env"] + assert pytest_kwargs["interrupt_grace"] == 180 + assert pytest_environment["UCODE_TEST_WORKSPACE"] == CUJ5_TARGET_ENV["UCODE_TEST_WORKSPACE"] + assert pytest_environment["DATABRICKS_BEARER"] == "workspace-token" + assert "DATABRICKS_CLIENT_ID" not in pytest_environment + assert "DATABRICKS_CLIENT_SECRET" not in pytest_environment + assert pytest_environment["UG_BUDGET_ACCOUNT_TOKEN"] == "account-token" + assert {key: pytest_environment[key] for key in CUJ5_TARGET_ENV} == CUJ5_TARGET_ENV + assert pytest_environment["UG_BUDGET_FUTURE_METADATA"] == "future-value" + assert "UG_BUDGET_CALLER_TOKEN" not in pytest_environment + assert "UG_BUDGET_CALLER_SECRET" not in pytest_environment + assert all( + "DATABRICKS_CLIENT_ID" not in kwargs.get("env", {}) + and "DATABRICKS_CLIENT_SECRET" not in kwargs.get("env", {}) + and "UG_BUDGET_ACCOUNT_TOKEN" not in kwargs.get("env", {}) + and not any(key.startswith("UG_BUDGET_") for key in kwargs.get("env", {})) + for command, kwargs in process_calls + if not ("-m" in command and "pytest" in command) + ) + + +def test_main_redacts_budget_account_token_errors(tmp_path, monkeypatch): + args, environment, mint_calls, account_calls, process_calls = _main_setup( + tmp_path, + monkeypatch, + workspace=CUJ5_TARGET_ENV["UCODE_TEST_WORKSPACE"], + client_id="cuj-client", + client_secret="cuj-secret", + target_environment=CUJ5_TARGET_ENV, + ) + + def reject_account_token(account_host, client_id, client_secret, *, account_id): + account_calls.append((account_host, client_id, client_secret, account_id)) + raise urllib.error.HTTPError( + account_host, + 401, + f"account token endpoint rejected: {client_secret}", + {}, + None, + ) + + monkeypatch.setattr(runner, "mint_account_m2m_token", reject_account_token) + with patch.dict(runner.os.environ, environment, clear=True): + assert runner.main() == 1 + report = json.loads((args.output / "versions.json").read_text()) + assert "account token endpoint rejected" in report["error"] + assert "cuj-secret" not in report["error"] + assert mint_calls == [(CUJ5_TARGET_ENV["UCODE_TEST_WORKSPACE"], "cuj-client", "cuj-secret")] + assert account_calls == [ + ( + CUJ5_TARGET_ENV["UG_BUDGET_ACCOUNT_HOST"], + "cuj-client", + "cuj-secret", + CUJ5_TARGET_ENV["UG_BUDGET_ACCOUNT_ID"], + ) + ] + assert not any("-m" in command and "pytest" in command for command, _ in process_calls) + + +def test_main_omits_budget_environment_when_not_provided(tmp_path, monkeypatch): + args, environment, mint_calls, account_calls, process_calls = _main_setup( + tmp_path, + monkeypatch, + workspace="https://workspace.example", + client_id="client", + client_secret="secret", + ) + with patch.dict(runner.os.environ, environment, clear=True): + assert runner.main() == 0 + assert mint_calls == [("https://workspace.example", "client", "secret")] + assert account_calls == [] + pytest_kwargs = _pytest_call(process_calls) + pytest_environment = pytest_kwargs["env"] + assert pytest_kwargs["interrupt_grace"] == 15 + assert pytest_environment["UCODE_TEST_WORKSPACE"] == "https://workspace.example" + assert not any(key.startswith("UG_BUDGET_") for key in pytest_environment) + assert all( + not any(key.startswith("UG_BUDGET_") for key in kwargs.get("env", {})) + for command, kwargs in process_calls + if not ("-m" in command and "pytest" in command) + ) + + +def test_budget_client_secret_is_redacted_and_excluded_from_agent_environment(tmp_path): + with patch.dict( + "os.environ", + { + "DATABRICKS_CLIENT_SECRET": "sensitive-client-secret", + "UG_BUDGET_ACCOUNT_TOKEN": "sensitive-account-token", + }, + ): + environment = clean_environment(tmp_path) + session = UserSession(tmp_path, tmp_path, Path("ug"), tmp_path / "artifacts") + assert "DATABRICKS_CLIENT_SECRET" not in environment + assert "UG_BUDGET_ACCOUNT_TOKEN" not in environment + assert session.redact( + "request failed: sensitive-client-secret; token: sensitive-account-token" + ) == ("request failed: ; token: ") + assert session.redact("request failed: sensitive-account-token") == ( + "request failed: " + ) def windows_live_environment():