diff --git a/tests/integration/security/_sweeps.py b/tests/integration/security/_sweeps.py index f97a0a7fcc6..5b9fa9a613f 100644 --- a/tests/integration/security/_sweeps.py +++ b/tests/integration/security/_sweeps.py @@ -106,6 +106,7 @@ ROUTE_DENY_LIST: Final = MappingProxyType( "/openai_passthrough/{endpoint:path}": "forwards to a provider, not a proxy read", "/get/latest_release_info": "fetches the latest release from api.github.com", "/roi-calculator/repositories": "lists repositories from the configured GitHub API, api.github.com by default", + "/roi-calculator/observed/repositories": "lists repositories from the connected GitHub or GitLab API, api.github.com by default", } ) @@ -140,6 +141,7 @@ NOT_FOUND_EXPECTED: Final = MappingProxyType( "/fallback/{model}": "answers 404 when the model has no fallbacks configured", "/team/{team_id}/members/me": "answers 404 when the caller is not a member, which the admin is not", "/guardrails/submissions/{guardrail_id}": "answers 404 for a guardrail no team submitted for review", + "/credentials/{credential_name:path}/jwks": "answers 404 for any credential that is not an anthropic internal_issuer credential, which the canary credential is not", } ) diff --git a/tests/integration/spend/test_lens_billing.py b/tests/integration/spend/test_lens_billing.py index d8eded62b39..69a1899ba6d 100644 --- a/tests/integration/spend/test_lens_billing.py +++ b/tests/integration/spend/test_lens_billing.py @@ -11,6 +11,8 @@ from tests.integration._support.database import read_rows, write_rows from tests.integration._support.process import owned_proxy from tests.integration.pricing.test_off_peak_pricing import off_peak_window +RELEASE_TAG: Final = "v0.0.0-lens-integration" + def delete_lens(lens_id: str) -> None: write_rows('DELETE FROM "LiteLLM_LensRun" WHERE lens_id=%s', (lens_id,)) @@ -19,8 +21,13 @@ def delete_lens(lens_id: str) -> None: @pytest.mark.parametrize("off_peak", (False, True)) -def test_lens_bills_selected_key_and_rechecks_its_permissions(gateway: Gateway, off_peak: bool) -> None: - with gateway.scenario() as scenario: +def test_lens_bills_selected_key_and_rechecks_its_permissions( + gateway: Gateway, tmp_path: Path, off_peak: bool +) -> None: + with ( + owned_proxy(gateway, tmp_path, {"LITELLM_RELEASE_TAG": RELEASE_TAG}) as isolated, + isolated.scenario() as scenario, + ): model: Final = scenario.model( input_cost_per_token=0.000001, output_cost_per_token=0.000002, @@ -36,12 +43,13 @@ def test_lens_bills_selected_key_and_rechecks_its_permissions(gateway: Gateway, ) key: Final = scenario.key(models=[model], max_budget=1) key_id: Final = sha256(key.encode()).hexdigest() - worker: Final = gateway.post( + worker: Final = isolated.post( "/lens/workers/register", {"name": "Billing regression", "analysis_key_id": key_id} ) + assert worker["image"] == "ghcr.io/berriai/litellm-lens-worker:" + RELEASE_TAG worker_id: Final = string_value(object_value(worker["worker"])["id"]) scenario.cleanups.callback(write_rows, 'DELETE FROM "LiteLLM_LensWorker" WHERE id=%s', (worker_id,)) - lens: Final = gateway.post( + lens: Final = isolated.post( "/lens", { "name": "Billing regression", @@ -54,14 +62,19 @@ def test_lens_bills_selected_key_and_rechecks_its_permissions(gateway: Gateway, lens_id: Final = string_value(lens["id"]) scenario.cleanups.callback(delete_lens, lens_id) worker_key: Final = string_value(worker["token"]) - unauthorized: Final = gateway.request( + unauthorized: Final = isolated.request( "POST", "/lens/workers/register", {"name": "Denied", "analysis_key_id": key_id}, key=key ) assert unauthorized.status_code == 403, unauthorized.text with ThreadPoolExecutor(max_workers=8) as pool: claims: Final = tuple( pool.map( - lambda _: gateway.request("POST", "/lens/worker/claim?protocol_version=2", {}, key=worker_key), + lambda _: isolated.request( + "POST", + "/lens/worker/claim?protocol_version=4&worker_release=" + RELEASE_TAG, + {}, + key=worker_key, + ), range(8), ) ) @@ -72,7 +85,7 @@ def test_lens_bills_selected_key_and_rechecks_its_permissions(gateway: Gateway, assert claim["lens_id"] == lens_id job_id: Final = string_value(object_value(claim["job"])["id"]) path: Final = f"/lens/worker/{lens_id}/{job_id}/model" - result: Final = gateway.post(path, {"prompt": "Inspect this run", "purpose": "extract"}, key=worker_key) + result: Final = isolated.post(path, {"prompt": "Inspect this run", "purpose": "extract"}, key=worker_key) expected: Final = (20 * 0.000001 + 20 * 0.000002) * (0.5 if off_peak else 1) assert result["cost"] == pytest.approx(expected) rows: Final = eventually( @@ -81,8 +94,8 @@ def test_lens_bills_selected_key_and_rechecks_its_permissions(gateway: Gateway, seconds=70, ) assert rows[0]["spend"] == pytest.approx(expected) - assert gateway.get(f"/lens/{lens_id}")["spent"] == pytest.approx(expected) - raw_hash: Final = gateway.request( + assert isolated.get(f"/lens/{lens_id}")["spent"] == pytest.approx(expected) + raw_hash: Final = isolated.request( "POST", "/v1/chat/completions", { @@ -92,31 +105,35 @@ def test_lens_bills_selected_key_and_rechecks_its_permissions(gateway: Gateway, key=key_id, ) assert raw_hash.status_code == 401, raw_hash.text - gateway.post("/key/update", {"key": key, "max_budget": expected / 2}) - exhausted: Final = gateway.request( + isolated.post("/key/update", {"key": key, "max_budget": expected / 2}) + exhausted: Final = isolated.request( "POST", path, {"prompt": "Must not run", "purpose": "extract"}, key=worker_key ) assert exhausted.status_code == 402, exhausted.text - gateway.post("/key/update", {"key": key, "max_budget": 1, "models": ["unavailable-analysis-model"]}) - restricted: Final = gateway.request( + isolated.post("/key/update", {"key": key, "max_budget": 1, "models": ["unavailable-analysis-model"]}) + restricted: Final = isolated.request( "POST", path, {"prompt": "Must not run", "purpose": "extract"}, key=worker_key ) assert restricted.status_code == 403, restricted.text - gateway.post("/key/block", {"key": key}) - blocked: Final = gateway.request("POST", path, {"prompt": "Must not run", "purpose": "extract"}, key=worker_key) + isolated.post("/key/block", {"key": key}) + blocked: Final = isolated.request( + "POST", path, {"prompt": "Must not run", "purpose": "extract"}, key=worker_key + ) assert blocked.status_code == 400, blocked.text - assert gateway.get(f"/lens/{lens_id}")["spent"] == pytest.approx(expected) + assert isolated.get(f"/lens/{lens_id}")["spent"] == pytest.approx(expected) replacement: Final = scenario.key(models=[model], rpm_limit=1) replacement_id: Final = sha256(replacement.encode()).hexdigest() - changed: Final = gateway.request( + changed: Final = isolated.request( "PUT", f"/lens/workers/{worker_id}/billing-key", {"analysis_key_id": replacement_id} ) assert changed.status_code == 200, changed.text - billed_replacement: Final = gateway.post( + billed_replacement: Final = isolated.post( path, {"prompt": "Inspect another run", "purpose": "extract"}, key=worker_key ) assert billed_replacement["cost"] == pytest.approx(expected) - limited: Final = gateway.request("POST", path, {"prompt": "Must not run", "purpose": "extract"}, key=worker_key) + limited: Final = isolated.request( + "POST", path, {"prompt": "Must not run", "purpose": "extract"}, key=worker_key + ) assert limited.status_code == 429, limited.text second_rows: Final = eventually( lambda: read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (replacement_id,)), @@ -124,16 +141,16 @@ def test_lens_bills_selected_key_and_rechecks_its_permissions(gateway: Gateway, seconds=70, ) assert second_rows[0]["spend"] == pytest.approx(expected) - active_revoke: Final = gateway.request("DELETE", f"/lens/workers/{worker_id}") + active_revoke: Final = isolated.request("DELETE", f"/lens/workers/{worker_id}") assert active_revoke.status_code == 409, active_revoke.text - gateway.post(f"/lens/{lens_id}/cancel", {}) - revoked: Final = gateway.request("DELETE", f"/lens/workers/{worker_id}") + isolated.post(f"/lens/{lens_id}/cancel", {}) + revoked: Final = isolated.request("DELETE", f"/lens/workers/{worker_id}") assert revoked.status_code == 200, revoked.text - denied_worker: Final = gateway.request( + denied_worker: Final = isolated.request( "POST", path, {"prompt": "Must not run", "purpose": "extract"}, key=worker_key ) assert denied_worker.status_code == 401, denied_worker.text - forbidden_change: Final = gateway.request( + forbidden_change: Final = isolated.request( "PUT", f"/lens/workers/{worker_id}/billing-key", {"analysis_key_id": replacement_id} ) assert forbidden_change.status_code == 409, forbidden_change.text @@ -161,7 +178,10 @@ def test_worker_spend_logs_do_not_expose_investigation_content( } ) ) - with owned_proxy(gateway, tmp_path, {}, config=config) as isolated, isolated.scenario() as scenario: + with ( + owned_proxy(gateway, tmp_path, {"LITELLM_RELEASE_TAG": RELEASE_TAG}, config=config) as isolated, + isolated.scenario() as scenario, + ): model: Final = scenario.model(input_cost_per_token=0.000001, output_cost_per_token=0.000002) key: Final = scenario.key(models=[model]) key_id: Final = sha256(key.encode()).hexdigest() @@ -185,7 +205,9 @@ def test_worker_spend_logs_do_not_expose_investigation_content( lens_id: Final = string_value(lens["id"]) scenario.cleanups.callback(delete_lens, lens_id) worker_token: Final = string_value(worker["token"]) - claim: Final = isolated.post("/lens/worker/claim?protocol_version=2", {}, key=worker_token) + claim: Final = isolated.post( + "/lens/worker/claim?protocol_version=4&worker_release=" + RELEASE_TAG, {}, key=worker_token + ) job_id: Final = string_value(object_value(claim["job"])["id"]) result: Final = isolated.post( f"/lens/worker/{lens_id}/{job_id}/model", {"prompt": marker, "purpose": "extract"}, key=worker_token