diff --git a/tests/integration/cost_calculation/conftest.py b/tests/integration/cost_calculation/conftest.py index d7f817efc3a..4f1f723c2f7 100644 --- a/tests/integration/cost_calculation/conftest.py +++ b/tests/integration/cost_calculation/conftest.py @@ -123,7 +123,7 @@ def poll_cost_row(key: str) -> CostRow: return result -def poll_rows(key: str) -> tuple[CostRow, ...]: +def poll_rows(key: str, count: int) -> tuple[CostRow, ...]: digest: Final = sha256(key.encode()).hexdigest() def read() -> tuple[CostRow, ...]: @@ -134,11 +134,18 @@ def poll_rows(key: str) -> tuple[CostRow, ...]: ) return tuple(parsed for row in rows if (parsed := _row(row)) is not None) - result: Final = eventually(read, lambda rows: len(rows) > 0, seconds=20) + result: Final = eventually(read, lambda rows: len(rows) >= count, seconds=60) return result -def poll_rollups(key: str, team_id: str, user_id: str, end_user_id: str) -> Rollups: +def poll_rollups( + key: str, + team_id: str, + user_id: str, + end_user_id: str, + requests: int, + target_spend: float, +) -> Rollups: digest: Final = sha256(key.encode()).hexdigest() def read() -> Rollups | None: @@ -179,7 +186,16 @@ def poll_rollups(key: str, team_id: str, user_id: str, end_user_id: str) -> Roll daily_team=DailySpend.model_validate(daily_team_rows[0]), ) - result: Final = eventually(read, lambda value: value is not None, seconds=20) + result: Final = eventually( + read, + lambda value: ( + value is not None + and value.daily_user.api_requests >= requests + and value.daily_team.api_requests >= requests + and approx_equal(value.key_spend, target_spend) + ), + seconds=60, + ) assert result is not None return result diff --git a/tests/integration/cost_calculation/cost_tracking_case.py b/tests/integration/cost_calculation/cost_tracking_case.py index effb6f2ed35..56a7cefb443 100644 --- a/tests/integration/cost_calculation/cost_tracking_case.py +++ b/tests/integration/cost_calculation/cost_tracking_case.py @@ -178,6 +178,7 @@ class RecountExpected(BaseModel): prompt_tokens: int | None = None completion_tokens: int | None = None min_completion_tokens: int | None = None + max_completion_tokens: int | None = None class FailureDetails(BaseModel): diff --git a/tests/integration/cost_calculation/cost_tracking_cases.json b/tests/integration/cost_calculation/cost_tracking_cases.json index facab77828c..bc5b40ccafe 100644 --- a/tests/integration/cost_calculation/cost_tracking_cases.json +++ b/tests/integration/cost_calculation/cost_tracking_cases.json @@ -30014,7 +30014,9 @@ "recount": { "input_cost_per_token": 1.75e-06, "output_cost_per_token": 1.4e-05 - } + }, + "prompt_tokens": 8, + "max_completion_tokens": 30 } } ] diff --git a/tests/integration/cost_calculation/test_cost_tracking.py b/tests/integration/cost_calculation/test_cost_tracking.py index 15b5921c94e..de55ab396fe 100644 --- a/tests/integration/cost_calculation/test_cost_tracking.py +++ b/tests/integration/cost_calculation/test_cost_tracking.py @@ -228,6 +228,10 @@ def _assert_recount(case: CostTrackingTestCase, expected: RecountExpected, row: assert row.completion_tokens >= expected.min_completion_tokens, ( f"{case.name}: completion_tokens {row.completion_tokens} < minimum {expected.min_completion_tokens}" ) + if expected.max_completion_tokens is not None: + assert row.completion_tokens <= expected.max_completion_tokens, ( + f"{case.name}: completion_tokens {row.completion_tokens} > maximum {expected.max_completion_tokens}" + ) recount: Final = row.prompt_tokens * expected.recount.input_cost_per_token + ( row.completion_tokens * expected.recount.output_cost_per_token ) @@ -376,7 +380,11 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) assert response.is_success, f"{case.name}: proxy returned {response.status_code}: {response.text[:400]}" if case.response.content_type == "text/event-stream": _assert_stream_has_no_error(response.text) - rows: Final = poll_rows(key) if len(responses) > 1 else (poll_cost_row(key),) + rows: Final = ( + poll_rows(key, len(responses)) + if len(responses) > 1 or fallback_deployment is not None + else (poll_cost_row(key),) + ) if isinstance(expected, RecountExpected): row: Final = rows[0] _assert_recount(case, expected, row) @@ -408,8 +416,8 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) if expected.rollups: assert deployment is not None and team_id is not None and user_id is not None assert end_user_id is not None - rollups: Final = poll_rollups(key, team_id, user_id, end_user_id) target_spend: Final = expected.spend * 3 + rollups: Final = poll_rollups(key, team_id, user_id, end_user_id, 3, target_spend) assert approx_equal(rollups.key_spend, target_spend) assert approx_equal(rollups.team_spend, target_spend) assert approx_equal(rollups.user_spend, target_spend)