diff --git a/tests/integration/cost_calculation/conftest.py b/tests/integration/cost_calculation/conftest.py index 3536c020515..fb20f5a9cc2 100644 --- a/tests/integration/cost_calculation/conftest.py +++ b/tests/integration/cost_calculation/conftest.py @@ -138,14 +138,7 @@ def poll_rows(key: str, count: int) -> tuple[CostRow, ...]: return result -def poll_rollups( - key: str, - team_id: str, - user_id: str, - end_user_id: str, - requests: int, - target_spend: float, -) -> Rollups: +def poll_rollups(key: str, team_id: str, user_id: str, end_user_id: str, requests: int, spend: float) -> Rollups: digest: Final = sha256(key.encode()).hexdigest() def read() -> Rollups | None: @@ -185,18 +178,21 @@ def poll_rollups( daily_user=DailySpend.model_validate(daily_user_rows[0]), daily_team=DailySpend.model_validate(daily_team_rows[0]), ) + if rollups.daily_user.api_requests < requests or rollups.daily_team.api_requests < requests: + return None + if not all( + approx_equal(actual, spend) + for actual in ( + rollups.key_spend, + rollups.team_spend, + rollups.user_spend, + rollups.end_user_spend, + ) + ): + return None return rollups - 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, - ) + result: Final = eventually(read, lambda value: value is not None, seconds=60) assert result is not None return result diff --git a/tests/integration/cost_calculation/cost_tracking_cases.json b/tests/integration/cost_calculation/cost_tracking_cases.json index bc5b40ccafe..78d00f28381 100644 --- a/tests/integration/cost_calculation/cost_tracking_cases.json +++ b/tests/integration/cost_calculation/cost_tracking_cases.json @@ -30015,7 +30015,8 @@ "input_cost_per_token": 1.75e-06, "output_cost_per_token": 1.4e-05 }, - "prompt_tokens": 8, + "prompt_tokens": 10, + "min_completion_tokens": 9, "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 de55ab396fe..692b76fff15 100644 --- a/tests/integration/cost_calculation/test_cost_tracking.py +++ b/tests/integration/cost_calculation/test_cost_tracking.py @@ -380,11 +380,7 @@ 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, len(responses)) - if len(responses) > 1 or fallback_deployment is not None - else (poll_cost_row(key),) - ) + rows: Final = poll_rows(key, len(responses)) if isinstance(expected, RecountExpected): row: Final = rows[0] _assert_recount(case, expected, row) @@ -417,7 +413,7 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) assert deployment is not None and team_id is not None and user_id is not None assert end_user_id is not None target_spend: Final = expected.spend * 3 - rollups: Final = poll_rollups(key, team_id, user_id, end_user_id, 3, target_spend) + rollups: Final = poll_rollups(key, team_id, user_id, end_user_id, requests=3, spend=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)