2e1dfb7553
Replaces the BigQuery wrap-view pattern with a discovery + scoped-fetch toolkit driven by the analyst's Claude session. Adds /api/v2/{catalog,schema,sample,scan,scan/estimate}, da catalog/schema/describe/fetch/snapshot/disk-info CLI commands, sqlglot-backed WHERE validator, process-local quota tracker, agent rails skill (cli/skills/agnes-data-querying.md). BREAKING: BQ wrap views off by default — set data_source.bigquery.legacy_wrap_views=true for one cycle. Backward-compat field_validator on primary_key. Catalog cache now matches documented 300s TTL with RBAC fresh per request. Cuts release v0.14.0.
120 lines
4.5 KiB
Python
120 lines
4.5 KiB
Python
"""Tests for the process-local v2 scan quota tracker (spec §3.8)."""
|
|
|
|
from datetime import datetime, timedelta, timezone
|
|
import pytest
|
|
|
|
from app.api.v2_quota import (
|
|
QuotaTracker,
|
|
QuotaExceededError,
|
|
KIND_CONCURRENT,
|
|
KIND_DAILY_BYTES,
|
|
)
|
|
|
|
|
|
def make_tracker(max_concurrent=5, max_daily_bytes=100):
|
|
return QuotaTracker(
|
|
max_concurrent_per_user=max_concurrent,
|
|
max_daily_bytes_per_user=max_daily_bytes,
|
|
)
|
|
|
|
|
|
class TestConcurrent:
|
|
def test_acquire_within_cap_succeeds(self):
|
|
q = make_tracker(max_concurrent=3)
|
|
with q.acquire(user="alice"):
|
|
with q.acquire(user="alice"):
|
|
with q.acquire(user="alice"):
|
|
pass
|
|
|
|
def test_acquire_above_cap_raises(self):
|
|
q = make_tracker(max_concurrent=2)
|
|
with q.acquire(user="alice"):
|
|
with q.acquire(user="alice"):
|
|
with pytest.raises(QuotaExceededError) as e:
|
|
with q.acquire(user="alice"):
|
|
pass
|
|
assert e.value.kind == KIND_CONCURRENT
|
|
assert e.value.current == 2
|
|
assert e.value.limit == 2
|
|
|
|
def test_release_on_context_exit(self):
|
|
q = make_tracker(max_concurrent=1)
|
|
with q.acquire(user="alice"):
|
|
pass
|
|
# Counter dropped on exit; new acquire works
|
|
with q.acquire(user="alice"):
|
|
pass
|
|
|
|
def test_release_on_exception(self):
|
|
q = make_tracker(max_concurrent=1)
|
|
try:
|
|
with q.acquire(user="alice"):
|
|
raise RuntimeError("boom")
|
|
except RuntimeError:
|
|
pass
|
|
with q.acquire(user="alice"):
|
|
pass
|
|
|
|
def test_per_user_isolation(self):
|
|
q = make_tracker(max_concurrent=1)
|
|
with q.acquire(user="alice"):
|
|
with q.acquire(user="bob"):
|
|
pass
|
|
|
|
|
|
class TestDailyBytes:
|
|
def test_record_within_cap(self):
|
|
q = make_tracker(max_daily_bytes=1000)
|
|
q.record_bytes(user="alice", n=300)
|
|
q.record_bytes(user="alice", n=400)
|
|
assert q.bytes_used_today(user="alice") == 700
|
|
|
|
def test_record_above_cap_no_longer_raises(self):
|
|
"""Post-scan recording NEVER raises — the user already paid for the
|
|
BQ scan, refusing to return the bytes they fetched would be perverse.
|
|
Pre-flight enforcement lives in check_daily_budget (called before
|
|
the scan runs)."""
|
|
q = make_tracker(max_daily_bytes=1000)
|
|
q.record_bytes(user="alice", n=600)
|
|
# Push over cap — record completes without raising.
|
|
q.record_bytes(user="alice", n=500)
|
|
assert q.bytes_used_today(user="alice") == 1100
|
|
|
|
def test_check_daily_budget_blocks_when_over_cap(self):
|
|
"""Once recorded bytes push past the cap, check_daily_budget refuses
|
|
the next request pre-flight — server doesn't run the BQ scan."""
|
|
q = make_tracker(max_daily_bytes=1000)
|
|
q.record_bytes(user="alice", n=600)
|
|
q.check_daily_budget(user="alice") # 600 < 1000 → ok
|
|
q.record_bytes(user="alice", n=500) # now at 1100
|
|
with pytest.raises(QuotaExceededError) as e:
|
|
q.check_daily_budget(user="alice")
|
|
assert e.value.kind == KIND_DAILY_BYTES
|
|
|
|
def test_check_daily_budget_at_exact_cap_rejects(self):
|
|
q = make_tracker(max_daily_bytes=1000)
|
|
q.record_bytes(user="alice", n=1000)
|
|
with pytest.raises(QuotaExceededError):
|
|
q.check_daily_budget(user="alice")
|
|
|
|
def test_per_user_isolation(self):
|
|
q = make_tracker(max_daily_bytes=100)
|
|
q.record_bytes(user="alice", n=80)
|
|
q.record_bytes(user="bob", n=80) # bob's bucket independent
|
|
# alice's check fails when over cap; bob's check still passes.
|
|
q.record_bytes(user="alice", n=30) # alice now at 110
|
|
with pytest.raises(QuotaExceededError):
|
|
q.check_daily_budget(user="alice")
|
|
q.check_daily_budget(user="bob") # bob still under
|
|
|
|
def test_reset_on_utc_midnight(self, monkeypatch):
|
|
q = make_tracker(max_daily_bytes=100)
|
|
d1 = datetime(2026, 4, 27, 23, 0, 0, tzinfo=timezone.utc)
|
|
monkeypatch.setattr("app.api.v2_quota._utcnow", lambda: d1)
|
|
q.record_bytes(user="alice", n=80)
|
|
assert q.bytes_used_today(user="alice") == 80
|
|
|
|
d2 = d1 + timedelta(hours=2) # crosses UTC midnight
|
|
monkeypatch.setattr("app.api.v2_quota._utcnow", lambda: d2)
|
|
assert q.bytes_used_today(user="alice") == 0
|
|
q.record_bytes(user="alice", n=80) # ok, fresh bucket
|