test(integration): wait for all rollup writes and pin fallback and disconnect rows

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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
 

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

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

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)