diff --git a/tests/integration/_support/process.py b/tests/integration/_support/process.py index 7323152b0b6..0d66ecc9d90 100644 --- a/tests/integration/_support/process.py +++ b/tests/integration/_support/process.py @@ -62,10 +62,28 @@ def owned_proxy(gateway: Gateway, directory: Path, overrides: Mapping[str, str]) output.mkdir(parents=True, exist_ok=True) with (output / f"owned-proxy-{uuid.uuid4().hex}.log").open("w") as log: process: Final = subprocess.Popen( - [sys.executable, "-m", "integration._support.proxy", "--config", "tests/integration/proxy_config.yaml", - "--host", "127.0.0.1", "--port", str(port), "--num_workers", "1", "--telemetry", "False", - "--use_prisma_db_push", "--enforce_prisma_migration_check"], - cwd=root, env=environment, stdout=log, stderr=subprocess.STDOUT, start_new_session=True, + [ + sys.executable, + "-m", + "integration._support.proxy", + "--config", + "tests/integration/proxy_config.yaml", + "--host", + "127.0.0.1", + "--port", + str(port), + "--num_workers", + "1", + "--telemetry", + "False", + "--use_prisma_db_push", + "--enforce_prisma_migration_check", + ], + cwd=root, + env=environment, + stdout=log, + stderr=subprocess.STDOUT, + start_new_session=True, ) try: with httpx.Client(base_url=f"http://127.0.0.1:{port}", timeout=15, trust_env=False) as client: diff --git a/tests/integration/database/test_partition_transactions.py b/tests/integration/database/test_partition_transactions.py index 2759d8dcfb7..dbf54e6962b 100644 --- a/tests/integration/database/test_partition_transactions.py +++ b/tests/integration/database/test_partition_transactions.py @@ -21,17 +21,26 @@ class PartitionConnection: db: Prisma -@pytest.mark.covers("other.database.partitions.lock_wait_outlives_transaction_default", "other.database.partitions.repeat_preserves_rows") +@pytest.mark.covers( + "other.database.partitions.lock_wait_outlives_transaction_default", + "other.database.partitions.repeat_preserves_rows", +) async def test_real_partition_ddl_survives_witnessed_lock_and_is_idempotent() -> None: schema: Final = f"integration_{uuid.uuid4().hex}" url: Final = os.environ["DATABASE_URL"] parsed: Final = urlsplit(url) - scoped_url: Final = urlunsplit(parsed._replace(query=urlencode({**dict(parse_qsl(parsed.query)), "schema": schema}))) + scoped_url: Final = urlunsplit( + parsed._replace(query=urlencode({**dict(parse_qsl(parsed.query)), "schema": schema})) + ) parent: Final = sql.Identifier(schema, "LiteLLM_SpendLogs") with psycopg.connect(url, autocommit=True) as setup: setup.execute(sql.SQL("CREATE SCHEMA {}").format(sql.Identifier(schema))) try: - setup.execute(sql.SQL('CREATE TABLE {} (request_id text, "startTime" timestamp NOT NULL) PARTITION BY RANGE ("startTime")').format(parent)) + setup.execute( + sql.SQL( + 'CREATE TABLE {} (request_id text, "startTime" timestamp NOT NULL) PARTITION BY RANGE ("startTime")' + ).format(parent) + ) database: Final = Prisma(datasource={"url": scoped_url}) await database.connect() try: @@ -39,12 +48,15 @@ async def test_real_partition_ddl_survives_witnessed_lock_and_is_idempotent() -> with psycopg.connect(url) as blocker: blocker.execute(sql.SQL("LOCK TABLE {} IN ACCESS SHARE MODE").format(parent)) blocker_pid: Final = blocker.info.backend_pid - operation: Final = asyncio.create_task(manager.ensure_partitions(PartitionConnection(database), lambda: 7000)) + operation: Final = asyncio.create_task( + manager.ensure_partitions(PartitionConnection(database), lambda: 7000) + ) wait_deadline: Final = time.monotonic() + 3 try: while True: witnesses: Final = read_rows( - "SELECT a.pid, extract(epoch FROM clock_timestamp()-a.query_start)::double precision AS age " + "SELECT a.pid, extract(epoch FROM " + "clock_timestamp()-a.query_start)::double precision AS age " "FROM pg_stat_activity a WHERE %s = ANY(pg_blocking_pids(a.pid)) " "AND a.wait_event_type = 'Lock' AND a.query LIKE 'CREATE TABLE IF NOT EXISTS%%'", (blocker_pid,), @@ -71,11 +83,12 @@ async def test_real_partition_ddl_survives_witnessed_lock_and_is_idempotent() -> catalog: Final = read_rows( "SELECT child.relname FROM pg_inherits i JOIN pg_class child ON child.oid=i.inhrelid " "JOIN pg_class parent ON parent.oid=i.inhparent JOIN pg_namespace n ON n.oid=parent.relnamespace " - "WHERE n.nspname=%s AND parent.relname='LiteLLM_SpendLogs'", (schema,), + "WHERE n.nspname=%s AND parent.relname='LiteLLM_SpendLogs'", + (schema,), ) assert catalog == [{"relname": ensured[0]}] now: Final = datetime.now(timezone.utc).replace(tzinfo=None) - setup.execute(sql.SQL('INSERT INTO {} VALUES (%s, %s)').format(parent), ("retained", now)) + setup.execute(sql.SQL("INSERT INTO {} VALUES (%s, %s)").format(parent), ("retained", now)) assert await manager.ensure_partitions(PartitionConnection(database), lambda: 7000) == ensured assert setup.execute(sql.SQL("SELECT request_id FROM {}").format(parent)).fetchall() == [("retained",)] finally: diff --git a/tests/integration/database/test_reader_writer_regeneration.py b/tests/integration/database/test_reader_writer_regeneration.py index 382dd14fbb0..2b4d93221ba 100644 --- a/tests/integration/database/test_reader_writer_regeneration.py +++ b/tests/integration/database/test_reader_writer_regeneration.py @@ -27,9 +27,15 @@ def test_key_regeneration_uses_writer_with_a_real_readonly_reader(gateway: Gatew role: Final = f"integration_reader_{uuid.uuid4().hex}" url: Final = os.environ["DATABASE_URL"] parsed: Final = urlsplit(url) - reader_url: Final = urlunsplit(parsed._replace(netloc=f"{role}:integration-reader-password@{parsed.hostname}:{parsed.port}")) + reader_url: Final = urlunsplit( + parsed._replace(netloc=f"{role}:integration-reader-password@{parsed.hostname}:{parsed.port}") + ) with psycopg.connect(url, autocommit=True) as admin: - admin.execute(sql.SQL("CREATE ROLE {} LOGIN PASSWORD 'integration-reader-password' NOSUPERUSER NOINHERIT").format(sql.Identifier(role))) + admin.execute( + sql.SQL("CREATE ROLE {} LOGIN PASSWORD 'integration-reader-password' NOSUPERUSER NOINHERIT").format( + sql.Identifier(role) + ) + ) try: admin.execute(sql.SQL("GRANT USAGE ON SCHEMA public TO {}").format(sql.Identifier(role))) admin.execute(sql.SQL("GRANT SELECT ON ALL TABLES IN SCHEMA public TO {}").format(sql.Identifier(role))) @@ -39,7 +45,9 @@ def test_key_regeneration_uses_writer_with_a_real_readonly_reader(gateway: Gatew with pytest.raises(psycopg.errors.ReadOnlySqlTransaction): reader.execute('UPDATE "LiteLLM_VerificationToken" SET blocked = true WHERE false') with owned_proxy(gateway, tmp_path, {"DATABASE_URL_READ_REPLICA": reader_url}) as candidate: - assert read_rows("SELECT pid FROM pg_stat_activity WHERE usename=%s", (role,)), "Candidate reader was never connected" + assert read_rows("SELECT pid FROM pg_stat_activity WHERE usename=%s", (role,)), ( + "Candidate reader was never connected" + ) with gateway.scenario() as scenario: model: Final = scenario.model() outside: Final = scenario.model() @@ -48,12 +56,24 @@ def test_key_regeneration_uses_writer_with_a_real_readonly_reader(gateway: Gatew scenario.cleanups.callback(delete_if_present, gateway, old) scenario.cleanups.callback(delete_if_present, gateway, new) old_hash: Final = sha256(old.encode()).hexdigest() - before: Final = candidate.request("POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "no grant yet"}]}, key=old) - assert before.status_code == 403 and before.json()["error"]["type"] == "key_model_access_denied", before.text - response: Final = candidate.request("POST", "/v1/access_group", { - "access_group_name": f"integration-{uuid.uuid4().hex}", - "access_model_names": [model], "assigned_key_ids": [old_hash], - }) + before: Final = candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "no grant yet"}]}, + key=old, + ) + assert before.status_code == 403 and before.json()["error"]["type"] == "key_model_access_denied", ( + before.text + ) + response: Final = candidate.request( + "POST", + "/v1/access_group", + { + "access_group_name": f"integration-{uuid.uuid4().hex}", + "access_model_names": [model], + "assigned_key_ids": [old_hash], + }, + ) assert response.status_code == 201, response.text group: Final = string_value(response.json()["access_group_id"]) try: @@ -61,30 +81,57 @@ def test_key_regeneration_uses_writer_with_a_real_readonly_reader(gateway: Gatew blocker.execute('LOCK TABLE "LiteLLM_AccessGroupTable" IN ACCESS EXCLUSIVE MODE') pending: Final = executor.submit(candidate.request, "GET", f"/v1/access_group/{group}") try: - reached: Final = eventually(lambda: read_rows( - "SELECT usename FROM pg_stat_activity WHERE %s=ANY(pg_blocking_pids(pid)) " - "AND usename=%s AND query LIKE 'SELECT%%'", (blocker.info.backend_pid, role), - ), bool, seconds=3) + reached: Final = eventually( + lambda: read_rows( + "SELECT usename FROM pg_stat_activity WHERE %s=ANY(pg_blocking_pids(pid)) " + "AND usename=%s AND query LIKE 'SELECT%%'", + (blocker.info.backend_pid, role), + ), + bool, + seconds=3, + ) assert reached == [{"usename": role}] finally: blocker.rollback() selected: Final = pending.result(timeout=5) - assert selected.status_code == 200 and selected.json()["access_group_id"] == group, selected.text + assert selected.status_code == 200 and selected.json()["access_group_id"] == group, ( + selected.text + ) assert candidate.chat(model, key=old)["usage"]["total_tokens"] == 40 - regenerated: Final = candidate.post("/key/regenerate", {"key": old, "new_key": new, "grace_period": "0s"}) + regenerated: Final = candidate.post( + "/key/regenerate", {"key": old, "new_key": new, "grace_period": "0s"} + ) assert regenerated["key"] == new new_hash: Final = sha256(new.encode()).hexdigest() assert new != old - assert read_rows('SELECT assigned_key_ids FROM "LiteLLM_AccessGroupTable" WHERE access_group_id=%s', (group,)) == [{"assigned_key_ids": [new_hash]}] - assert read_rows('SELECT token, access_group_ids FROM "LiteLLM_VerificationToken" WHERE token=ANY(%s)', ([old_hash, new_hash],)) == [{"token": new_hash, "access_group_ids": [group]}] + assert read_rows( + 'SELECT assigned_key_ids FROM "LiteLLM_AccessGroupTable" WHERE access_group_id=%s', (group,) + ) == [{"assigned_key_ids": [new_hash]}] + assert read_rows( + 'SELECT token, access_group_ids FROM "LiteLLM_VerificationToken" WHERE token=ANY(%s)', + ([old_hash, new_hash],), + ) == [{"token": new_hash, "access_group_ids": [group]}] assert candidate.chat(model, key=new)["usage"]["total_tokens"] == 40 assert candidate.chat(outside, key=new)["usage"]["total_tokens"] == 40 - denied: Final = candidate.request("POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "rotated key"}]}, key=old) - assert denied.status_code == 401 and denied.json()["error"]["type"] == "token_not_found_in_db", denied.text + denied: Final = candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "rotated key"}]}, + key=old, + ) + assert ( + denied.status_code == 401 and denied.json()["error"]["type"] == "token_not_found_in_db" + ), denied.text finally: deleted: Final = gateway.request("DELETE", f"/v1/access_group/{group}") assert deleted.status_code == 204, deleted.text - assert read_rows('SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_id=%s', (group,)) == [] + assert ( + read_rows( + 'SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_id=%s', + (group,), + ) + == [] + ) finally: admin.execute(sql.SQL("DROP OWNED BY {}").format(sql.Identifier(role))) admin.execute(sql.SQL("DROP ROLE {}").format(sql.Identifier(role))) diff --git a/tests/integration/database/test_transaction_atomicity.py b/tests/integration/database/test_transaction_atomicity.py index 431d7982482..c150354d9a6 100644 --- a/tests/integration/database/test_transaction_atomicity.py +++ b/tests/integration/database/test_transaction_atomicity.py @@ -24,31 +24,80 @@ def test_access_group_second_key_constraint_failure_rolls_back_all_writes(gatewa witness: Final = constraint + "_seq" check_function: Final = constraint + "_check" body: Final = {"access_group_name": name, "access_model_names": [model], "assigned_key_ids": tokens} + def remove_partial_group() -> None: - for row in read_rows('SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_name=%s', (name,)): + for row in read_rows( + 'SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_name=%s', (name,) + ): response: Final = gateway.request("DELETE", f"/v1/access_group/{row['access_group_id']}") assert response.status_code == 204, response.text - assert read_rows('SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_name=%s', (name,)) == [] + assert ( + read_rows('SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_name=%s', (name,)) + == [] + ) scenario.cleanups.callback(remove_partial_group) - before: Final = read_rows('SELECT token, access_group_ids FROM "LiteLLM_VerificationToken" WHERE token=ANY(%s) ORDER BY token', (tokens,)) + before: Final = read_rows( + 'SELECT token, access_group_ids FROM "LiteLLM_VerificationToken" WHERE token=ANY(%s) ORDER BY token', + (tokens,), + ) with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as connection, ExitStack() as cleanup: connection.execute(sql.SQL("CREATE SEQUENCE {}").format(sql.Identifier(witness))) cleanup.callback(connection.execute, sql.SQL("DROP SEQUENCE {}").format(sql.Identifier(witness))) - connection.execute(sql.SQL("CREATE FUNCTION {}(text[]) RETURNS boolean LANGUAGE plpgsql AS $$ BEGIN IF cardinality($1)>0 THEN PERFORM nextval({}); RETURN false; END IF; RETURN true; END $$").format(sql.Identifier(check_function), sql.Literal(witness))) - cleanup.callback(connection.execute, sql.SQL("DROP FUNCTION {}(text[])").format(sql.Identifier(check_function))) - connection.execute(sql.SQL('ALTER TABLE "LiteLLM_VerificationToken" ADD CONSTRAINT {} CHECK (token <> {} OR {}(access_group_ids))').format(sql.Identifier(constraint), sql.Literal(tokens[1]), sql.Identifier(check_function))) - cleanup.callback(connection.execute, sql.SQL('ALTER TABLE "LiteLLM_VerificationToken" DROP CONSTRAINT {}').format(sql.Identifier(constraint))) + connection.execute( + sql.SQL( + "CREATE FUNCTION {}(text[]) RETURNS boolean LANGUAGE plpgsql AS $$ BEGIN IF " + "cardinality($1)>0 THEN PERFORM nextval({}); RETURN false; END IF; RETURN true; END $$" + ).format(sql.Identifier(check_function), sql.Literal(witness)) + ) + cleanup.callback( + connection.execute, sql.SQL("DROP FUNCTION {}(text[])").format(sql.Identifier(check_function)) + ) + connection.execute( + sql.SQL( + 'ALTER TABLE "LiteLLM_VerificationToken" ADD ' + "CONSTRAINT {} CHECK (token <> {} OR {}(access_group_ids))" + ).format(sql.Identifier(constraint), sql.Literal(tokens[1]), sql.Identifier(check_function)) + ) + cleanup.callback( + connection.execute, + sql.SQL('ALTER TABLE "LiteLLM_VerificationToken" DROP CONSTRAINT {}').format( + sql.Identifier(constraint) + ), + ) try: - assert connection.execute(sql.SQL("SELECT is_called FROM {}").format(sql.Identifier(witness))).fetchone() == (False,) + assert connection.execute( + sql.SQL("SELECT is_called FROM {}").format(sql.Identifier(witness)) + ).fetchone() == (False,) failed: Final = gateway.request("POST", "/v1/access_group", body) assert failed.status_code == 500, failed.text - assert connection.execute(sql.SQL("SELECT is_called FROM {}").format(sql.Identifier(witness))).fetchone() == (True,) - assert read_rows('SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_name=%s', (name,)) == [] - assert read_rows('SELECT token, access_group_ids FROM "LiteLLM_VerificationToken" WHERE token=ANY(%s) ORDER BY token', (tokens,)) == before + assert connection.execute( + sql.SQL("SELECT is_called FROM {}").format(sql.Identifier(witness)) + ).fetchone() == (True,) + assert ( + read_rows( + 'SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_name=%s', (name,) + ) + == [] + ) + assert ( + read_rows( + "SELECT token, access_group_ids FROM " + '"LiteLLM_VerificationToken" WHERE token=ANY(%s) ORDER BY token', + (tokens,), + ) + == before + ) for key in keys: - denied: Final = gateway.request("POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "rolled back grant"}]}, key=key) - assert denied.status_code == 403 and denied.json()["error"]["type"] == "key_model_access_denied", denied.text + denied: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "rolled back grant"}]}, + key=key, + ) + assert denied.status_code == 403 and denied.json()["error"]["type"] == "key_model_access_denied", ( + denied.text + ) finally: cleanup.close() created: Final = gateway.request("POST", "/v1/access_group", body) @@ -60,6 +109,17 @@ def test_access_group_second_key_constraint_failure_rolls_back_all_writes(gatewa finally: deleted: Final = gateway.request("DELETE", f"/v1/access_group/{identity}") assert deleted.status_code == 204, deleted.text - assert read_rows('SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_id=%s', (identity,)) == [] - assert read_rows('SELECT token, access_group_ids FROM "LiteLLM_VerificationToken" WHERE token=ANY(%s) ORDER BY token', (tokens,)) == before - assert read_rows('SELECT conname FROM pg_constraint WHERE conname=%s', (constraint,)) == [] + assert ( + read_rows( + 'SELECT access_group_id FROM "LiteLLM_AccessGroupTable" WHERE access_group_id=%s', (identity,) + ) + == [] + ) + assert ( + read_rows( + 'SELECT token, access_group_ids FROM "LiteLLM_VerificationToken" WHERE token=ANY(%s) ORDER BY token', + (tokens,), + ) + == before + ) + assert read_rows("SELECT conname FROM pg_constraint WHERE conname=%s", (constraint,)) == [] diff --git a/tests/integration/pricing/test_price_precedence.py b/tests/integration/pricing/test_price_precedence.py index 9f312c1e670..0d73558d8a1 100644 --- a/tests/integration/pricing/test_price_precedence.py +++ b/tests/integration/pricing/test_price_precedence.py @@ -55,7 +55,8 @@ def test_generated_zero_null_and_omitted_prices_follow_independent_arithmetic(ga assert response.headers.get("x-litellm-response-cost") in (None, "0", "0.0") rows: Final = eventually( lambda: read_rows( - 'SELECT spend, metadata, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + "SELECT spend, metadata, prompt_tokens, " + 'completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (response.json()["id"],), ), lambda values: len(values) == 1, diff --git a/tests/integration/spend/test_cache_and_quota.py b/tests/integration/spend/test_cache_and_quota.py index 05bf100f091..114aabbae33 100644 --- a/tests/integration/spend/test_cache_and_quota.py +++ b/tests/integration/spend/test_cache_and_quota.py @@ -22,7 +22,9 @@ def test_generated_cache_sequences_preserve_content_usage_and_zero_hit_cost(gate self.resources = ExitStack() try: self.scenario = self.resources.enter_context(gateway.scenario()) - self.upstream = self.resources.enter_context(httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False)) + self.upstream = self.resources.enter_context( + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) + ) self.model = self.scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) self.key = self.scenario.key(models=[self.model]) self.prefix = uuid.uuid4().hex @@ -46,13 +48,22 @@ def test_generated_cache_sequences_preserve_content_usage_and_zero_hit_cost(gate def perform_request(self, marker: int) -> None: self.upstream.get("/__observations").raise_for_status() - response: Final = gateway.request("POST", "/v1/chat/completions", { - "model": self.model, "messages": [{"role": "user", "content": f"{self.prefix}-{marker}"}], - }, key=self.key) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": self.model, + "messages": [{"role": "user", "content": f"{self.prefix}-{marker}"}], + }, + key=self.key, + ) assert response.status_code == 200, response.text self.requests += 1 body: Final = response.json() - assert body["choices"][0]["message"]["content"] == "Hello! This is a mock response from the fake OpenAI endpoint." + assert ( + body["choices"][0]["message"]["content"] + == "Hello! This is a mock response from the fake OpenAI endpoint." + ) assert body["usage"]["total_tokens"] == 40 observed: Final = self.upstream.get("/__observations").json()["requests"] expected_calls: Final = 0 if marker in self.seen else 1 @@ -70,10 +81,15 @@ def test_generated_cache_sequences_preserve_content_usage_and_zero_hit_cost(gate def teardown(self) -> None: try: if self.requests and not self.failed: - rows: Final = eventually(lambda: read_rows( - 'SELECT request_id, spend, cache_hit, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE api_key=%s', - (sha256(self.key.encode()).hexdigest(),), - ), lambda values: len(values) == self.requests, seconds=70) + rows: Final = eventually( + lambda: read_rows( + "SELECT request_id, spend, cache_hit, prompt_tokens, " + 'completion_tokens FROM "LiteLLM_SpendLogs" WHERE api_key=%s', + (sha256(self.key.encode()).hexdigest(),), + ), + lambda values: len(values) == self.requests, + seconds=70, + ) assert len({row["request_id"] for row in rows}) == self.requests assert sum(float(row["spend"]) for row in rows) == pytest.approx(self.paid * 0.06) assert sum(row["cache_hit"] == "True" for row in rows) == self.requests - self.paid @@ -81,7 +97,10 @@ def test_generated_cache_sequences_preserve_content_usage_and_zero_hit_cost(gate assert row["prompt_tokens"] == 20 and row["completion_tokens"] == 20 if row["cache_hit"] == "True": assert float(row["spend"]) == 0 and "_cache_hit" in row["request_id"] - assert any(row["request_id"].startswith(identity + "_cache_hit") for identity in self.identities.values()) + assert any( + row["request_id"].startswith(identity + "_cache_hit") + for identity in self.identities.values() + ) else: assert row["request_id"] in self.identities.values() assert float(row["spend"]) == pytest.approx(0.06) @@ -95,7 +114,10 @@ def test_generated_cache_sequences_preserve_content_usage_and_zero_hit_cost(gate @pytest.mark.covers("quota_management.response_cache.repeated_hits_preserve_identity_and_single_charge") def test_repeated_hits_keep_response_identity_and_create_distinct_zero_cost_rows(gateway: Gateway) -> None: - with gateway.scenario() as scenario, httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream: + with ( + gateway.scenario() as scenario, + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + ): model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) key: Final = scenario.key(models=[model]) prompt: Final = f"repeated cache {uuid.uuid4().hex}" @@ -104,9 +126,19 @@ def test_repeated_hits_keep_response_identity_and_create_distinct_zero_cost_rows assert len(upstream.get("/__observations").json()["requests"]) == 1 assert len({result["id"] for result in results}) == 1 for result in results: - assert result["choices"][0]["message"]["content"] == "Hello! This is a mock response from the fake OpenAI endpoint." + assert ( + result["choices"][0]["message"]["content"] + == "Hello! This is a mock response from the fake OpenAI endpoint." + ) assert result["usage"]["total_tokens"] == 40 - rows: Final = eventually(lambda: read_rows('SELECT request_id, spend, cache_hit FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (sha256(key.encode()).hexdigest(),)), lambda values: len(values) == 3, seconds=70) + rows: Final = eventually( + lambda: read_rows( + 'SELECT request_id, spend, cache_hit FROM "LiteLLM_SpendLogs" WHERE api_key=%s', + (sha256(key.encode()).hexdigest(),), + ), + lambda values: len(values) == 3, + seconds=70, + ) assert len({row["request_id"] for row in rows}) == 3 assert sorted(float(row["spend"]) for row in rows) == [0, 0, 0.06] for row in rows: @@ -119,39 +151,74 @@ def test_repeated_hits_keep_response_identity_and_create_distinct_zero_cost_rows @pytest.mark.covers("quota_management.budget.key.boundary_blocks_before_provider_and_reset_restores") def test_key_budget_at_boundary_blocks_provider_then_explicit_reset_restores(gateway: Gateway) -> None: - with gateway.scenario() as scenario, httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream: + with ( + gateway.scenario() as scenario, + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + ): model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) key: Final = scenario.key(models=[model], max_budget=0.06) control: Final = scenario.key(models=[model]) first: Final = gateway.chat(model, key=key, text=f"budget {uuid.uuid4().hex}") assert first["usage"]["total_tokens"] == 40 digest: Final = sha256(key.encode()).hexdigest() - spent: Final = eventually(lambda: read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)), lambda values: len(values) == 1 and float(values[0]["spend"]) >= 0.06, seconds=70) + spent: Final = eventually( + lambda: read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)), + lambda values: len(values) == 1 and float(values[0]["spend"]) >= 0.06, + seconds=70, + ) assert float(spent[0]["spend"]) == pytest.approx(0.06) upstream.get("/__observations").raise_for_status() - denied: Final = gateway.request("POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": f"over budget {uuid.uuid4().hex}"}]}, key=key) + denied: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": f"over budget {uuid.uuid4().hex}"}]}, + key=key, + ) assert denied.status_code == 429 and denied.json()["error"]["type"] == "budget_exceeded", denied.text assert upstream.get("/__observations").json()["requests"] == [] assert gateway.chat(model, key=control, text=f"control {uuid.uuid4().hex}")["usage"]["total_tokens"] == 40 gateway.post("/key/update", {"key": key, "spend": 0}) - assert read_rows('SELECT spend, max_budget FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)) == [{"spend": 0.0, "max_budget": 0.06}] + assert read_rows('SELECT spend, max_budget FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)) == [ + {"spend": 0.0, "max_budget": 0.06} + ] assert gateway.chat(model, key=key, text=f"reset {uuid.uuid4().hex}")["usage"]["total_tokens"] == 40 - eventually(lambda: read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)), lambda values: len(values) == 1 and float(values[0]["spend"]) >= 0.06, seconds=70) + eventually( + lambda: read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)), + lambda values: len(values) == 1 and float(values[0]["spend"]) >= 0.06, + seconds=70, + ) upstream.get("/__observations").raise_for_status() - denied_again: Final = gateway.request("POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": f"boundary again {uuid.uuid4().hex}"}]}, key=key) - assert denied_again.status_code == 429 and denied_again.json()["error"]["type"] == "budget_exceeded", denied_again.text + denied_again: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": f"boundary again {uuid.uuid4().hex}"}]}, + key=key, + ) + assert denied_again.status_code == 429 and denied_again.json()["error"]["type"] == "budget_exceeded", ( + denied_again.text + ) assert upstream.get("/__observations").json()["requests"] == [] @pytest.mark.covers("quota_management.response_cache.system_messages_partition_cache_identity") def test_different_system_messages_do_not_share_a_cached_response(gateway: Gateway) -> None: - with gateway.scenario() as scenario, httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream: + with ( + gateway.scenario() as scenario, + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + ): model: Final = scenario.model() prompt: Final = uuid.uuid4().hex identities: dict[str, str] = {} for system, expected_calls in (("first policy", 1), ("second policy", 1), ("first policy", 0)): upstream.get("/__observations").raise_for_status() - response: Final = gateway.request("POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "system", "content": system}, {"role": "user", "content": prompt}]}) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "system", "content": system}, {"role": "user", "content": prompt}], + }, + ) assert response.status_code == 200 and response.json()["usage"]["total_tokens"] == 40, response.text calls: Final = upstream.get("/__observations").json()["requests"] assert len(calls) == expected_calls @@ -161,4 +228,7 @@ def test_different_system_messages_do_not_share_a_cached_response(gateway: Gatew assert response.json()["id"] not in identities.values() identities = {**identities, system: response.json()["id"]} if calls: - assert calls[0]["body"]["messages"] == [{"role": "system", "content": system}, {"role": "user", "content": prompt}] + assert calls[0]["body"]["messages"] == [ + {"role": "system", "content": system}, + {"role": "user", "content": prompt}, + ]