From 5cbe0eb02e6dea20346803596b7fefe03f4637ee Mon Sep 17 00:00:00 2001 From: Mrinal Chanshetty Date: Mon, 19 Oct 2026 13:41:40 -0700 Subject: [PATCH] test(proxy): cover team-attributed daily tag spend Unit: one tag from two teams stays two upsert rows; an entry with no team_id (queued by an older pod) folds into the "" row; only the tag table is keyed by team_id. The writer queues one tag from two teams as separate team-attributed entries, read back through the real queue. Integration: two teams send the same tag through the proxy and LiteLLM_DailyTagSpend holds one row per team with the hand-computed spend. Mutations (each red, then restored): - team_id dropped from the upsert conflict key: the two-team and legacy-entry upsert tests fail - team_id added to every table's conflict key: all five per-table tests fail - team_id dropped from the writer's queue key: the writer test fails, with all three requests merged into one entry credited to team-a - the writer stops recording the team: the integration test fails, with both teams' spend under team_id '' LIT-8516 --- .../spend/test_team_tag_spend_attribution.py | 44 +++++++++++++++++++ .../proxy/db/test_daily_spend_bulk_upsert.py | 39 ++++++++++++++++ .../proxy/db/test_db_spend_update_writer.py | 39 ++++++++++++++++ 3 files changed, 122 insertions(+) create mode 100644 tests/integration/spend/test_team_tag_spend_attribution.py diff --git a/tests/integration/spend/test_team_tag_spend_attribution.py b/tests/integration/spend/test_team_tag_spend_attribution.py new file mode 100644 index 00000000000..b0bd50ca44e --- /dev/null +++ b/tests/integration/spend/test_team_tag_spend_attribution.py @@ -0,0 +1,44 @@ +import uuid +from typing import Final + +import pytest +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows + +_SPEND_PER_CALL: Final = 20 * 0.001 + 20 * 0.002 + + +def test_one_tag_sent_by_two_teams_is_recorded_once_per_team(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + tag: Final = f"cost-center:{uuid.uuid4().hex}" + teams: Final = tuple(scenario.team(models=[model]) for _ in range(2)) + calls_per_team: Final = {teams[0]: 1, teams[1]: 2} + for team, calls in calls_per_team.items(): + key = scenario.key(team_id=team, models=[model]) + for _ in range(calls): + reply = gateway.post( + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": f"tagged {uuid.uuid4().hex}"}], + "metadata": {"tags": [tag]}, + }, + key=key, + ) + assert reply["usage"]["total_tokens"] == 40, reply + + rows: Final = eventually( + lambda: read_rows( + 'SELECT team_id, spend, api_requests FROM "LiteLLM_DailyTagSpend" WHERE tag = %s ORDER BY team_id', + (tag,), + ), + lambda values: sum(int(row["api_requests"]) for row in values) == 3, + seconds=70, + ) + + assert sorted((row["team_id"], int(row["api_requests"])) for row in rows) == sorted(calls_per_team.items()), ( + rows + ) + for row in rows: + assert float(row["spend"]) == pytest.approx(calls_per_team[row["team_id"]] * _SPEND_PER_CALL), rows diff --git a/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py b/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py index 0912c0c406d..9886ec91cda 100644 --- a/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py +++ b/tests/unit/proxy/db/test_daily_spend_bulk_upsert.py @@ -144,6 +144,45 @@ def test_non_tag_tables_carry_no_request_id_column(): assert '"user_id"' in sql +def test_one_tag_from_two_teams_stays_two_rows(): + merged = merge_by_conflict_key( + TAG_TABLE, (tag_txn(team_id="team-b", spend=2.0), tag_txn(team_id="team-a", spend=1.0)) + ) + + assert [(txn["team_id"], txn["spend"]) for _, txn in merged] == [("team-a", 1.0), ("team-b", 2.0)], merged + + +def test_an_entry_without_team_id_folds_into_the_unattributed_row(): + """Entries queued by an older pod, or replayed from Redis across the upgrade, carry no + team_id key. They arbitrate to the same row as an explicit "" rather than a second one.""" + legacy = tag_txn(spend=0.25) + unattributed = tag_txn(team_id="", spend=0.75) + + merged = merge_by_conflict_key(TAG_TABLE, (legacy, unattributed)) + + assert len(merged) == 1, merged + key, folded = merged[0] + assert key[:2] == ("team-a", ""), key + assert folded["spend"] == pytest.approx(1.0) + + +@pytest.mark.parametrize("entity", ["user", "team", "org", "end_user", "agent"]) +def test_only_the_tag_table_is_keyed_by_team(entity): + """Only LiteLLM_DailyTagSpend has the extra team_id column; adding it to any other + table's conflict target names a column that table does not have.""" + table = DAILY_SPEND_TABLES[entity] + row = {**tag_txn(), table.entity_id_column: "id-1"} + + sql, _ = build_bulk_upsert(table, merge_by_conflict_key(table, (row,))) + + conflict_target = re.search(r"ON CONFLICT \(([^)]*)\)", sql) + assert conflict_target is not None + assert conflict_target.group(1) == ( + f'"{table.entity_id_column}", "date", "api_key", "model", "custom_llm_provider", ' + '"mcp_namespaced_tool_name", "endpoint"' + ) + + class _RecordingDb: def __init__(self) -> None: self.statements: list[tuple[str, tuple[object, ...]]] = [] diff --git a/tests/unit/proxy/db/test_db_spend_update_writer.py b/tests/unit/proxy/db/test_db_spend_update_writer.py index 1ec65a4fdcd..af7a00bae25 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -1488,6 +1488,45 @@ async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_i ) +@pytest.mark.asyncio +async def test_one_tag_from_two_teams_is_queued_as_two_team_attributed_entries(): + """The queue merges entries by key before the upsert ever sees them, so a key without the + team would fold both teams' spend on a shared tag into one entry.""" + writer = DBSpendUpdateWriter() + prisma = MagicMock() + prisma.get_request_status = MagicMock(return_value="success") + + def tagged(team_id: str | None, spend: float) -> dict: + return { + "request_id": f"req-{team_id}", + "request_tags": '["cost-center:cc-42"]', + "team_id": team_id, + "user": "user-1", + "startTime": "2026-10-17T00:00:00", + "api_key": "key-1", + "model": "gpt-4o", + "custom_llm_provider": "openai", + "model_group": "gpt-4o", + "prompt_tokens": 10, + "completion_tokens": 5, + "spend": spend, + "metadata": '{"usage_object": {}}', + } + + for team_id, spend in (("team-a", 1.0), ("team-b", 2.0), (None, 0.5)): + await writer.add_spend_log_transaction_to_daily_tag_transaction( + payload=tagged(team_id, spend), prisma_client=prisma + ) + + queued = await writer.daily_tag_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + + assert sorted((txn["tag"], txn["team_id"], txn["spend"]) for txn in queued.values()) == [ + ("cost-center:cc-42", "", 0.5), + ("cost-center:cc-42", "team-a", 1.0), + ("cost-center:cc-42", "team-b", 2.0), + ], queued + + @pytest.mark.asyncio async def test_add_spend_log_transaction_to_daily_org_transaction_injects_org_id_and_queues_update(): """