test(integration): restore the concurrent rollup and disconnect fix from cc93a37

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

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
 

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

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)