Format database and accounting integration tests

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:

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:

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)))

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,)) == []

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,

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},
+                ]