mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-13 03:59:12 +00:00
a93c396a5c
* test(integration): move legacy proxy, router and Redis tests into tests/integration Port 39 legacy tests to the integration tier that owns them, running against the scripted upstream, local Postgres and Redis, test-owned wire peers and owned proxies. Delete 5 legacy tests whose contract is already owned by an existing integration test, and remove the legacy functions, files and helpers left unused. * test(integration): cover recovery of a spent key after its budget is raised
57 lines
2.5 KiB
Python
57 lines
2.5 KiB
Python
import json
|
|
from hashlib import sha256
|
|
from typing import Final
|
|
|
|
import pytest
|
|
from integration._support.client import Gateway, eventually
|
|
from integration._support.database import read_rows
|
|
from integration._support.wire import Reply, Request, wire_server
|
|
|
|
PROMPT: Final = "a scripted sea otter"
|
|
PRICE_PER_IMAGE: Final = 0.25
|
|
|
|
|
|
def _image(request: Request) -> Reply:
|
|
assert (request.method, request.target) == ("POST", "/images/generations")
|
|
return Reply(body=json.dumps({"created": 1700000000, "data": [{"b64_json": "aW1n"}]}).encode())
|
|
|
|
|
|
def test_identical_image_generations_each_charge_the_key(gateway: Gateway) -> None:
|
|
with wire_server(_image) as wire, gateway.scenario() as scenario:
|
|
model: Final = scenario.model(
|
|
model="openai/dall-e-3",
|
|
api_base=wire.url,
|
|
api_key="synthetic-image-key",
|
|
output_cost_per_image=PRICE_PER_IMAGE,
|
|
)
|
|
key: Final = scenario.key(models=[model])
|
|
digest: Final = sha256(key.encode()).hexdigest()
|
|
body: Final = {"model": model, "prompt": PROMPT, "size": "1024x1024", "n": 1}
|
|
first: Final = gateway.request("POST", "/v1/images/generations", body, key=key)
|
|
assert first.status_code == 200, first.text
|
|
logged: Final = eventually(
|
|
lambda: read_rows('SELECT spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (digest,)),
|
|
lambda rows: len(rows) == 1,
|
|
seconds=70,
|
|
)
|
|
charge: Final = float(str(logged[0]["spend"]))
|
|
assert charge == pytest.approx(PRICE_PER_IMAGE)
|
|
eventually(
|
|
lambda: read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)),
|
|
lambda rows: float(str(rows[0]["spend"])) == pytest.approx(charge),
|
|
seconds=70,
|
|
)
|
|
repeat: Final = gateway.request("POST", "/v1/images/generations", body, key=key)
|
|
assert repeat.status_code == 200, repeat.text
|
|
rows: Final = eventually(
|
|
lambda: read_rows('SELECT spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (digest,)),
|
|
lambda values: len(values) == 2,
|
|
seconds=70,
|
|
)
|
|
assert [float(str(row["spend"])) for row in rows] == pytest.approx([charge, charge])
|
|
eventually(
|
|
lambda: read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)),
|
|
lambda values: float(str(values[0]["spend"])) == pytest.approx(2 * charge),
|
|
seconds=70,
|
|
)
|
|
assert len(wire.drain()) == 2
|