mirror of
https://github.com/superdesigndev/treg.git
synced 2026-10-02 03:24:35 +08:00
104 lines
3.7 KiB
Python
104 lines
3.7 KiB
Python
"""Atomic daily accounting for overflow aggregator spend."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from treg.infra.db import session_maker
|
|
from treg.domain.capacity import overflow_spend as spend_ledger
|
|
from treg.domain.capacity.overflow_spend import add_in_transaction
|
|
from treg.models import OverflowSpend
|
|
|
|
|
|
async def test_add_in_transaction_does_not_commit_for_its_caller(clients):
|
|
async with session_maker() as db:
|
|
row = await add_in_transaction(db, "orthogonal", 2_000, 100, day="2026-08-28")
|
|
row = await add_in_transaction(db, "orthogonal", 3_000, -50, day="2026-08-28")
|
|
assert (row.calls, row.cost_micro, row.delta_micro) == (2, 5_000, 50)
|
|
await db.rollback()
|
|
|
|
async with session_maker() as db:
|
|
assert await db.get(OverflowSpend, ("orthogonal", "2026-08-28")) is None
|
|
|
|
|
|
async def test_concurrent_first_adds_are_atomic(clients, monkeypatch):
|
|
original_get = AsyncSession.get
|
|
both_read_missing = asyncio.Event()
|
|
reads = 0
|
|
gate_reads = True
|
|
|
|
async def get_after_both_read(self, entity, ident, *args, **kwargs):
|
|
nonlocal reads
|
|
row = await original_get(self, entity, ident, *args, **kwargs)
|
|
if (
|
|
gate_reads and entity is OverflowSpend
|
|
and ident == ("orthogonal", "2026-08-28")
|
|
):
|
|
reads += 1
|
|
if reads == 2:
|
|
both_read_missing.set()
|
|
await asyncio.wait_for(both_read_missing.wait(), timeout=2)
|
|
return row
|
|
|
|
monkeypatch.setattr(AsyncSession, "get", get_after_both_read)
|
|
|
|
async def add(cost_micro: int, delta_micro: int) -> None:
|
|
async with session_maker() as db:
|
|
await add_in_transaction(
|
|
db, "orthogonal", cost_micro, delta_micro, day="2026-08-28",
|
|
)
|
|
await db.commit()
|
|
|
|
results = await asyncio.gather(add(2_000, 100), add(3_000, -50), return_exceptions=True)
|
|
gate_reads = False
|
|
|
|
assert results == [None, None]
|
|
async with session_maker() as db:
|
|
row = await db.get(OverflowSpend, ("orthogonal", "2026-08-28"))
|
|
assert row is not None
|
|
assert (row.calls, row.cost_micro, row.delta_micro) == (2, 5_000, 50)
|
|
|
|
|
|
async def test_concurrent_budget_reservations_allow_at_most_one_call_at_the_cap(clients):
|
|
async def reserve() -> bool:
|
|
async with session_maker() as db:
|
|
row = await spend_ledger.reserve_in_transaction(
|
|
db, "orthogonal", 3_000, 4_000, day="2026-08-28",
|
|
)
|
|
await db.commit()
|
|
return row is not None
|
|
|
|
admitted = await asyncio.gather(reserve(), reserve())
|
|
|
|
assert sum(admitted) == 1
|
|
async with session_maker() as db:
|
|
row = await db.get(OverflowSpend, ("orthogonal", "2026-08-28"))
|
|
assert row is not None
|
|
assert (row.calls, row.cost_micro) == (0, 3_000)
|
|
|
|
|
|
async def test_first_budget_reservation_larger_than_the_cap_is_rejected(clients):
|
|
async with session_maker() as db:
|
|
row = await spend_ledger.reserve_in_transaction(
|
|
db, "orthogonal", 5_000, 4_000, day="2026-08-28",
|
|
)
|
|
await db.commit()
|
|
|
|
assert row is None
|
|
async with session_maker() as db:
|
|
assert await db.get(OverflowSpend, ("orthogonal", "2026-08-28")) is None
|
|
|
|
|
|
async def test_budget_reservation_does_not_commit_for_its_caller(clients):
|
|
async with session_maker() as db:
|
|
row = await spend_ledger.reserve_in_transaction(
|
|
db, "orthogonal", 3_000, 4_000, day="2026-08-28",
|
|
)
|
|
assert row is not None
|
|
await db.rollback()
|
|
|
|
async with session_maker() as db:
|
|
assert await db.get(OverflowSpend, ("orthogonal", "2026-08-28")) is None
|