test(usage): assert per-user entity spend and metadata (#50251)

Co-authored-by: mateo <277851410+mateo-berri@users.noreply.github.com>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

tests/integration/_support/daily_activity.py

@@ -272,16 +272,32 @@ def assert_key_reported(
     date: str,
     metadata: Mapping[str, JsonValue],
     metrics: Mapping[str, float],
+    *,
+    entity_rows: Mapping[str, int] | None = None,
 ) -> None:
     assert response.status_code == 200, response.text
     body: Final = object_value(response.json())
     records: Final = tuple(object_value(record) for record in records_of_key(body, api_key))
     assert records, response.text
     assert all(record["metadata"] == metadata for record in records), response.text
-    assert all(counted(record["metrics"]) == pytest.approx(metrics) for record in records), response.text
     days: Final = body["results"]
     assert isinstance(days, list) and len(days) == 1, response.text
     day: Final = object_value(days[0])
+    breakdown: Final = object_value(day["breakdown"])
+    if entity_rows is None:
+        assert all(counted(record["metrics"]) == pytest.approx(metrics) for record in records), response.text
+    else:
+        totals: Final = records_of_key({name: value for name, value in breakdown.items() if name != "entities"}, api_key)
+        assert totals, response.text
+        assert all(counted(object_value(record)["metrics"]) == pytest.approx(metrics) for record in totals), response.text
+        entities: Final = object_value(breakdown["entities"])
+        assert set(entities) == set(entity_rows), response.text
+        for entity, rows in entity_rows.items():
+            bucket: Final = object_value(entities[entity])
+            assert counted(bucket["metrics"]) == pytest.approx(seeded_metrics(rows)), response.text
+            keys: Final = object_value(bucket["api_key_breakdown"])
+            assert set(keys) == {api_key}, response.text
+            assert counted(object_value(keys[api_key])["metrics"]) == pytest.approx(seeded_metrics(rows)), response.text
     assert day["date"] == date, response.text
     assert counted(day["metrics"]) == pytest.approx(metrics), response.text
     assert object_value(body["metadata"])["total_spend"] == pytest.approx(metrics["spend"]), response.text

tests/integration/spend/test_daily_activity_key_owner.py

@@ -57,6 +57,7 @@ def test_key_whose_daily_spend_names_two_users_is_reported_with_no_owner(gateway
                 DAY,
                 key_metadata(),
                 seeded_metrics(2),
+                entity_rows={first: 1, second: 1},
             )
 
 
@@ -74,6 +75,7 @@ def test_daily_spend_rows_naming_no_user_do_not_hide_the_one_user_the_others_nam
                 DAY,
                 key_metadata(user=owner, email=email),
                 seeded_metrics(2),
+                entity_rows={owner: 1, "Unassigned": 1},
             )
 
 

tests/integration/spend/test_daily_activity_key_owner_faults.py

@@ -199,7 +199,7 @@ def test_key_stops_being_reported_with_an_owner_once_a_second_user_spends_with_i
             with daily_rows((user_row(second, api_key, DAY),)):
                 shared: Final = activity_of_key(gateway, AGGREGATED_USER_ACTIVITY, api_key)
     assert_key_reported(alone, api_key, DAY, key_metadata(user=first, email=email), seeded_metrics(1))
-    assert_key_reported(shared, api_key, DAY, key_metadata(), seeded_metrics(2))
+    assert_key_reported(shared, api_key, DAY, key_metadata(), seeded_metrics(2), entity_rows={first: 1, second: 1})
 
 
 @pytest.mark.timeout(300)