spb/company-atlas

1"""Rate limiting (token bucket, tiers, headers), API keys and streamed exports."""2from __future__ import annotations34import csv5import hashlib6import io7import json89import pytest10import test_api_support as support1112from companyatlas.api.ratelimit import RateLimiter, flush_usage, resolve_api_key13from companyatlas.db import fetch_one, transaction1415client = support.client16fixture_data = support.fixture_data17ADMIN = support.ADMIN18RAW_API_KEY = support.RAW_API_KEY1920pytestmark = pytest.mark.asyncio(loop_scope="session")21V = "/api/v1"222324async def test_token_bucket_semantics() -> None:25    rl = RateLimiter(limits={"anonymous": 3, "paid": 6, "internal": None})26    t = 1000.027    results = [rl.take("ip:1", "anonymous", now=t) for _ in range(4)]28    assert [r[0] for r in results] == [True, True, True, False]29    assert results[0][1] == 3 and results[0][2] == 2 and results[3][2] == 0 and results[3][3] > 030    allowed, *_ = rl.take("ip:1", "anonymous", now=t + 20)       # 3/min → one token back after 20 s31    assert allowed is True32    assert rl.take("ip:2", "anonymous", now=t)[0] is True         # independent bucket33    assert rl.take("key:x", "internal", now=t) == (True, 0, 0, 0.0)34    assert rl.take("key:y", "unknown-tier", now=t)[1] == 3         # unknown tier falls back to anonymous353637async def test_rate_limit_headers_and_tiers(client):  # type: ignore[no-untyped-def]38    from companyatlas.api import ratelimit3940    r = await client.get(f"{V}/methodology")41    assert r.headers["x-ratelimit-tier"] == "internal"                # loopback without X-Forwarded-For = our own SSR42    public = {"X-Forwarded-For": "203.0.113.7"}43    r = await client.get(f"{V}/methodology", headers=public)44    anon = ratelimit.limiter.limits["anonymous"]45    assert r.headers["x-ratelimit-tier"] == "anonymous" and r.headers["x-ratelimit-limit"] == str(anon) and int(r.headers["x-ratelimit-remaining"]) < anon46    assert ratelimit.TIER_LIMITS_PER_MIN == {"anonymous": 120, "authenticated": 600, "paid": 3000, "internal": None}47    r = await client.get(f"{V}/methodology", headers={"X-CA-API-Key": RAW_API_KEY})48    assert r.headers["x-ratelimit-tier"] == "paid" and r.headers["x-ratelimit-limit"] == "3000"49    r = await client.get(f"{V}/methodology", headers={"X-CA-API-Key": "ca_paid_not_a_real_key_000000", **public})50    assert r.headers["x-ratelimit-tier"] == "anonymous"51    r = await client.get(f"{V}/methodology", headers=ADMIN)52    assert r.headers["x-ratelimit-tier"] == "admin" and "x-ratelimit-remaining" not in r.headers53    r = await client.get("/health")54    assert "x-ratelimit-tier" not in r.headers                       # bypassed path55    info = await resolve_api_key(RAW_API_KEY)56    assert info and info["tier"] == "paid"57    await flush_usage()58    async with transaction() as conn:59        row = await fetch_one(conn, "select request_count, last_used_at from api_keys where key_hash = :h", h=hashlib.sha256(RAW_API_KEY.encode()).hexdigest())60    assert row and row["request_count"] >= 1 and row["last_used_at"] is not None616263async def test_429_from_exhausted_bucket(client, monkeypatch):  # type: ignore[no-untyped-def]64    from companyatlas.api import ratelimit6566    tiny = RateLimiter(limits={"anonymous": 2, "authenticated": 600, "paid": 3000, "internal": None})67    monkeypatch.setattr(ratelimit, "limiter", tiny)68    headers = {"X-Forwarded-For": "203.0.113.9, 10.0.0.1"}69    codes = [(await client.get(f"{V}/methodology", headers=headers)).status_code for _ in range(3)]70    assert codes == [200, 200, 429]71    r = await client.get(f"{V}/methodology", headers=headers)72    assert r.status_code == 429 and int(r.headers["retry-after"]) >= 1 and r.headers["x-ratelimit-remaining"] == "0" and r.json() == {"detail": "rate limit exceeded"}73    assert (await client.get(f"{V}/methodology", headers={"X-Forwarded-For": "203.0.113.10"})).status_code == 200747576async def test_export_events_formats(client, fixture_data):  # type: ignore[no-untyped-def]77    r = await client.get(f"{V}/export/events.csv", params={"country": "ZZ"})78    assert r.status_code == 200 and r.headers["content-type"].startswith("text/csv") and "attachment" in r.headers["content-disposition"]79    rows = list(csv.reader(io.StringIO(r.text)))80    assert rows[0][:4] == ["id", "detected_at", "company_slug", "company_name"] and len(rows) == 581    r = await client.get(f"{V}/export/events.ndjson", params={"country": "ZZ", "event_type": "PRICING"})82    lines = [json.loads(x) for x in r.text.splitlines() if x]83    assert r.headers["content-type"].startswith("application/x-ndjson") and len(lines) == 1 and lines[0]["id"] == fixture_data["ev_pricing"]84    r = await client.get(f"{V}/export/events.json", params={"company": fixture_data["alpha_slug"], "limit": 2})85    data = r.json()86    assert isinstance(data, list) and len(data) == 2 and data[0]["company"]["slug"] == fixture_data["alpha_slug"]87    assert (await client.get(f"{V}/export/events.xml")).status_code == 40488    assert (await client.get(f"{V}/export/events.json", params={"company": "nope"})).status_code == 404899091async def test_export_companies_and_jobs(client, fixture_data):  # type: ignore[no-untyped-def]92    r = await client.get(f"{V}/export/companies.csv", params={"country": "ZZ"})93    rows = list(csv.reader(io.StringIO(r.text)))94    assert rows[0][0] == "id" and "activity_score" in rows[0] and len(rows) == 395    r = await client.get(f"{V}/export/companies.json", params={"industry": fixture_data["industry"]})96    assert {c["slug"] for c in r.json()} == {fixture_data["alpha_slug"], fixture_data["beta_slug"]}97    r = await client.get(f"{V}/export/jobs.ndjson", params={"company": fixture_data["alpha_slug"], "status": "all"})98    lines = [json.loads(x) for x in r.text.splitlines() if x]99    assert len(lines) == 3 and all(j["company_slug"] == fixture_data["alpha_slug"] for j in lines)100    r = await client.get(f"{V}/export/jobs.csv", params={"company": fixture_data["alpha_slug"], "ai": 1})101    rows = list(csv.reader(io.StringIO(r.text)))102    assert len(rows) == 2 and rows[1][2] == "Senior ML Engineer"103