mirror of
https://github.com/hanzoai/insights.git
synced 2026-10-30 16:15:31 +00:00
zeekay
19fdf55ceb
A key now opens with what it is allowed to touch. `pk-` is publishable and safe in a browser; `sk-` is secret and never leaves a server. The five marks it replaces all began `ph` -- the upstream's initials -- and they were the last of its brand a customer read back to us, in every SDK snippet. No word rule could reach them, because `ph` is not a word. The marks were also not telling the truth. The constant said `phc_` while every live project token said `hi_`, so the middleware that routes on the constant matched no real token at all. Two spellings existed and one was fiction. Minting now happens in exactly one function and reading in exactly one, so a second spelling has nowhere to live. `phx_` (personal) and `phs_` (project secret) both become `sk-`. The mark now says only that a value is secret; which store holds it is settled by asking. That was already how personal keys worked -- keys minted before the marks carry none, so the store always decided -- and both secret resolvers already returned None on a miss and fell through, so one shared extractor replaces the two that could drift apart. OAuth's access and refresh tokens stay distinguishable as `at-` and `rt-`, because the parser tells them apart before either is looked up, and the domain-verification token becomes `vt-` rather than keeping a sixth mint of its own. `bin/debrand --check` now refuses the marks as a shape, so a seventh cannot be coined. Two rules that were about to rename things that are not ours to spell are bounded: `Lemon` no longer eats `LemonSqueezy`, and the `hog_flow` rules no longer reach a Postgres table, a Datastore table or a route. Keys are stored hashed, so a mark cannot be rewritten in place: the three live project tokens have to be reminted.
323 lines
16 KiB
Python
323 lines
16 KiB
Python
from datetime import UTC, datetime, timedelta
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from llm_gateway.services.plan_resolver import (
|
|
PlanInfo,
|
|
PlanResolver,
|
|
get_billing_period_number,
|
|
is_pro_plan,
|
|
resolve_plan_info,
|
|
)
|
|
|
|
|
|
class TestIsProPlan:
|
|
@pytest.mark.parametrize(
|
|
"plan_key,expected",
|
|
[
|
|
("insights-code-200-20260301", True),
|
|
("insights-code-200-20260501", True),
|
|
("insights-code-pro-200-20260301", True),
|
|
("insights-code-pro-0-20260422", True),
|
|
("insights-code-free-20260301", False),
|
|
("insights-code-free-20260501", False),
|
|
("some-other-plan", False),
|
|
(None, False),
|
|
("", False),
|
|
],
|
|
)
|
|
def test_is_pro_plan(self, plan_key: str | None, expected: bool) -> None:
|
|
assert is_pro_plan(plan_key) is expected
|
|
|
|
|
|
class TestGetBillingPeriodNumber:
|
|
def test_none_returns_zero(self) -> None:
|
|
assert get_billing_period_number(None) == 0
|
|
|
|
def test_invalid_date_returns_zero(self) -> None:
|
|
assert get_billing_period_number("not-a-date") == 0
|
|
|
|
def test_today_returns_zero(self) -> None:
|
|
now = datetime.now(tz=UTC).isoformat()
|
|
assert get_billing_period_number(now) == 0
|
|
|
|
def test_five_days_ago_returns_zero(self) -> None:
|
|
created = (datetime.now(tz=UTC) - timedelta(days=5)).isoformat()
|
|
assert get_billing_period_number(created) == 0
|
|
|
|
def test_thirty_days_returns_one(self) -> None:
|
|
created = (datetime.now(tz=UTC) - timedelta(days=30)).isoformat()
|
|
assert get_billing_period_number(created) == 1
|
|
|
|
def test_sixty_days_returns_two(self) -> None:
|
|
created = (datetime.now(tz=UTC) - timedelta(days=60)).isoformat()
|
|
assert get_billing_period_number(created) == 2
|
|
|
|
def test_custom_period_days(self) -> None:
|
|
created = (datetime.now(tz=UTC) - timedelta(days=15)).isoformat()
|
|
assert get_billing_period_number(created, period_days=7) == 2
|
|
|
|
def test_just_under_boundary(self) -> None:
|
|
created = (datetime.now(tz=UTC) - timedelta(days=29, hours=23)).isoformat()
|
|
assert get_billing_period_number(created) == 0
|
|
|
|
def test_billing_period_start_overrides_seat_created_at(self) -> None:
|
|
old_seat = (datetime.now(tz=UTC) - timedelta(days=35)).isoformat()
|
|
recent_period = (datetime.now(tz=UTC) - timedelta(days=5)).isoformat()
|
|
assert get_billing_period_number(old_seat, billing_period_start=recent_period) == 0
|
|
|
|
def test_billing_period_start_none_falls_back_to_seat_created_at(self) -> None:
|
|
old_seat = (datetime.now(tz=UTC) - timedelta(days=35)).isoformat()
|
|
assert get_billing_period_number(old_seat, billing_period_start=None) == 1
|
|
|
|
|
|
class TestPlanResolver:
|
|
@pytest.fixture
|
|
def resolver(self) -> PlanResolver:
|
|
return PlanResolver(redis=None, http_client=AsyncMock())
|
|
|
|
@pytest.fixture
|
|
def resolver_with_redis(self) -> PlanResolver:
|
|
redis = AsyncMock()
|
|
redis.get = AsyncMock(return_value=None)
|
|
redis.set = AsyncMock()
|
|
return PlanResolver(redis=redis, http_client=AsyncMock())
|
|
|
|
async def test_returns_none_plan_when_no_api_url(self, resolver: PlanResolver) -> None:
|
|
with patch("llm_gateway.services.plan_resolver.get_settings") as mock_settings:
|
|
mock_settings.return_value.insights_api_base_url = ""
|
|
mock_settings.return_value.plan_cache_ttl = 300
|
|
result = await resolver.get_plan(user_id=1, auth_header="Bearer sk-test")
|
|
assert result.plan_key is None
|
|
assert result.seat_missing is False
|
|
|
|
async def test_returns_cached_plan(self, resolver_with_redis: PlanResolver) -> None:
|
|
import json
|
|
|
|
cached = json.dumps({"plan_key": "insights-code-200-20260301", "created_at": None})
|
|
assert resolver_with_redis._redis is not None
|
|
resolver_with_redis._redis.get = AsyncMock(return_value=cached.encode()) # type: ignore[method-assign] # ty: ignore[invalid-assignment]
|
|
result = await resolver_with_redis.get_plan(user_id=1, auth_header="Bearer sk-test")
|
|
assert result.plan_key == "insights-code-200-20260301"
|
|
|
|
async def test_cached_null_plan_returns_none_plan_key(self, resolver_with_redis: PlanResolver) -> None:
|
|
import json
|
|
|
|
cached = json.dumps({"plan_key": None, "created_at": None})
|
|
assert resolver_with_redis._redis is not None
|
|
resolver_with_redis._redis.get = AsyncMock(return_value=cached.encode()) # type: ignore[method-assign] # ty: ignore[invalid-assignment]
|
|
result = await resolver_with_redis.get_plan(user_id=1, auth_header="Bearer sk-test")
|
|
assert result.plan_key is None
|
|
# Entries cached before seat_missing existed must fail loose.
|
|
assert result.seat_missing is False
|
|
|
|
async def test_cached_old_seat_returns_plan(self, resolver_with_redis: PlanResolver) -> None:
|
|
import json
|
|
|
|
old_date = (datetime.now(tz=UTC) - timedelta(days=60)).isoformat()
|
|
cached = json.dumps({"plan_key": "insights-code-free-20260301", "created_at": old_date})
|
|
assert resolver_with_redis._redis is not None
|
|
resolver_with_redis._redis.get = AsyncMock(return_value=cached.encode()) # type: ignore[method-assign] # ty: ignore[invalid-assignment]
|
|
result = await resolver_with_redis.get_plan(user_id=1, auth_header="Bearer sk-test")
|
|
assert result.plan_key == "insights-code-free-20260301"
|
|
assert result.seat_created_at == old_date
|
|
|
|
async def test_fetches_from_api_and_caches(self, resolver_with_redis: PlanResolver) -> None:
|
|
created_at = datetime.now(tz=UTC).isoformat()
|
|
mock_resp = MagicMock()
|
|
mock_resp.status_code = 200
|
|
mock_resp.json.return_value = {
|
|
"plan_key": "insights-code-pro-200-20260301",
|
|
"created_at": created_at,
|
|
"billing_period": {
|
|
"current_period_start": "2026-04-01T00:00:00+00:00",
|
|
"current_period_end": "2026-05-01T00:00:00+00:00",
|
|
"interval": "month",
|
|
},
|
|
}
|
|
mock_resp.raise_for_status = MagicMock()
|
|
resolver_with_redis._http.get = AsyncMock(return_value=mock_resp) # type: ignore[method-assign] # ty: ignore[invalid-assignment]
|
|
|
|
with patch("llm_gateway.services.plan_resolver.get_settings") as mock_settings:
|
|
mock_settings.return_value.insights_api_base_url = "https://app.hanzo.ai"
|
|
mock_settings.return_value.plan_cache_ttl = 300
|
|
result = await resolver_with_redis.get_plan(user_id=1, auth_header="Bearer sk-test")
|
|
|
|
assert result.plan_key == "insights-code-pro-200-20260301"
|
|
assert result.billing_period is not None
|
|
assert result.billing_period.current_period_start == "2026-04-01T00:00:00+00:00"
|
|
assert result.billing_period.interval == "month"
|
|
assert resolver_with_redis._redis is not None
|
|
resolver_with_redis._redis.set.assert_called_once() # type: ignore[attr-defined]
|
|
|
|
resolver_with_redis._http.get.assert_called_once_with(
|
|
"https://app.hanzo.ai/api/seats/me/",
|
|
params={"product_key": "insights_code", "best": "true"},
|
|
headers={"Authorization": "Bearer sk-test"},
|
|
timeout=2.0,
|
|
)
|
|
|
|
async def test_fetches_from_api_without_billing_period(self, resolver_with_redis: PlanResolver) -> None:
|
|
mock_resp = MagicMock()
|
|
mock_resp.status_code = 200
|
|
mock_resp.json.return_value = {"plan_key": "insights-code-200-20260301", "created_at": None}
|
|
mock_resp.raise_for_status = MagicMock()
|
|
resolver_with_redis._http.get = AsyncMock(return_value=mock_resp) # type: ignore[method-assign] # ty: ignore[invalid-assignment]
|
|
|
|
with patch("llm_gateway.services.plan_resolver.get_settings") as mock_settings:
|
|
mock_settings.return_value.insights_api_base_url = "https://app.hanzo.ai"
|
|
mock_settings.return_value.plan_cache_ttl = 300
|
|
result = await resolver_with_redis.get_plan(user_id=1, auth_header="Bearer sk-test")
|
|
|
|
assert result.plan_key == "insights-code-200-20260301"
|
|
assert result.billing_period is None
|
|
|
|
async def test_forwards_auth_header(self, resolver: PlanResolver) -> None:
|
|
mock_resp = MagicMock()
|
|
mock_resp.status_code = 200
|
|
mock_resp.json.return_value = {"plan_key": None, "created_at": None}
|
|
mock_resp.raise_for_status = MagicMock()
|
|
resolver._http.get = AsyncMock(return_value=mock_resp) # type: ignore[method-assign] # ty: ignore[invalid-assignment]
|
|
|
|
with patch("llm_gateway.services.plan_resolver.get_settings") as mock_settings:
|
|
mock_settings.return_value.insights_api_base_url = "https://app.hanzo.ai"
|
|
mock_settings.return_value.plan_cache_ttl = 300
|
|
await resolver.get_plan(user_id=1, auth_header="Bearer sk-mysecretkey")
|
|
|
|
resolver._http.get.assert_called_once()
|
|
call_kwargs = resolver._http.get.call_args
|
|
assert call_kwargs.kwargs["headers"]["Authorization"] == "Bearer sk-mysecretkey"
|
|
|
|
async def test_404_returns_none_plan(self, resolver: PlanResolver) -> None:
|
|
mock_resp = MagicMock()
|
|
mock_resp.status_code = 404
|
|
resolver._http.get = AsyncMock(return_value=mock_resp) # type: ignore[method-assign] # ty: ignore[invalid-assignment]
|
|
|
|
with patch("llm_gateway.services.plan_resolver.get_settings") as mock_settings:
|
|
mock_settings.return_value.insights_api_base_url = "https://app.hanzo.ai"
|
|
mock_settings.return_value.plan_cache_ttl = 300
|
|
result = await resolver.get_plan(user_id=1, auth_header="Bearer sk-test")
|
|
|
|
assert result.plan_key is None
|
|
assert result.seat_missing is True
|
|
|
|
async def test_404_seat_missing_round_trips_through_cache(self, resolver_with_redis: PlanResolver) -> None:
|
|
import json
|
|
|
|
mock_resp = MagicMock()
|
|
mock_resp.status_code = 404
|
|
resolver_with_redis._http.get = AsyncMock(return_value=mock_resp) # type: ignore[method-assign] # ty: ignore[invalid-assignment]
|
|
|
|
with patch("llm_gateway.services.plan_resolver.get_settings") as mock_settings:
|
|
mock_settings.return_value.insights_api_base_url = "https://app.hanzo.ai"
|
|
mock_settings.return_value.plan_cache_ttl = 300
|
|
await resolver_with_redis.get_plan(user_id=1, auth_header="Bearer sk-test")
|
|
|
|
assert resolver_with_redis._redis is not None
|
|
cached_payload = resolver_with_redis._redis.set.call_args.args[1] # type: ignore[attr-defined]
|
|
resolver_with_redis._redis.get = AsyncMock(return_value=cached_payload.encode()) # type: ignore[method-assign] # ty: ignore[invalid-assignment]
|
|
cached_result = await resolver_with_redis.get_plan(user_id=1, auth_header="Bearer sk-test")
|
|
assert json.loads(cached_payload)["seat_missing"] is True
|
|
assert cached_result.seat_missing is True
|
|
|
|
async def test_api_error_returns_none_plan(self, resolver: PlanResolver) -> None:
|
|
resolver._http.get = AsyncMock(side_effect=Exception("connection refused")) # type: ignore[method-assign] # ty: ignore[invalid-assignment]
|
|
|
|
with patch("llm_gateway.services.plan_resolver.get_settings") as mock_settings:
|
|
mock_settings.return_value.insights_api_base_url = "https://app.hanzo.ai"
|
|
mock_settings.return_value.plan_cache_ttl = 300
|
|
result = await resolver.get_plan(user_id=1, auth_header="Bearer sk-test")
|
|
|
|
assert result.plan_key is None
|
|
assert result.seat_missing is False
|
|
|
|
async def test_empty_auth_header_skips_fetch(self, resolver: PlanResolver) -> None:
|
|
result = await resolver.get_plan(user_id=1, auth_header="")
|
|
assert result.plan_key is None
|
|
resolver._http.get.assert_not_called() # type: ignore[attr-defined]
|
|
|
|
async def test_cached_plan_with_billing_period(self, resolver_with_redis: PlanResolver) -> None:
|
|
import json
|
|
|
|
cached = json.dumps(
|
|
{
|
|
"plan_key": "insights-code-pro-200-20260301",
|
|
"created_at": None,
|
|
"billing_period": {
|
|
"current_period_start": "2026-04-01T00:00:00+00:00",
|
|
"current_period_end": "2026-05-01T00:00:00+00:00",
|
|
"interval": "month",
|
|
},
|
|
}
|
|
)
|
|
assert resolver_with_redis._redis is not None
|
|
resolver_with_redis._redis.get = AsyncMock(return_value=cached.encode()) # type: ignore[method-assign] # ty: ignore[invalid-assignment]
|
|
result = await resolver_with_redis.get_plan(user_id=1, auth_header="Bearer sk-test")
|
|
assert result.billing_period is not None
|
|
assert result.billing_period.current_period_start == "2026-04-01T00:00:00+00:00"
|
|
assert result.billing_period.interval == "month"
|
|
|
|
async def test_cached_plan_without_billing_period(self, resolver_with_redis: PlanResolver) -> None:
|
|
import json
|
|
|
|
cached = json.dumps({"plan_key": "insights-code-200-20260301", "created_at": None})
|
|
assert resolver_with_redis._redis is not None
|
|
resolver_with_redis._redis.get = AsyncMock(return_value=cached.encode()) # type: ignore[method-assign] # ty: ignore[invalid-assignment]
|
|
result = await resolver_with_redis.get_plan(user_id=1, auth_header="Bearer sk-test")
|
|
assert result.billing_period is None
|
|
|
|
async def test_api_error_not_cached(self, resolver_with_redis: PlanResolver) -> None:
|
|
resolver_with_redis._http.get = AsyncMock(side_effect=Exception("connection refused")) # type: ignore[method-assign] # ty: ignore[invalid-assignment]
|
|
|
|
with patch("llm_gateway.services.plan_resolver.get_settings") as mock_settings:
|
|
mock_settings.return_value.insights_api_base_url = "https://app.hanzo.ai"
|
|
mock_settings.return_value.plan_cache_ttl = 300
|
|
result = await resolver_with_redis.get_plan(user_id=1, auth_header="Bearer sk-test")
|
|
|
|
assert result.plan_key is None
|
|
assert resolver_with_redis._redis is not None
|
|
resolver_with_redis._redis.set.assert_not_called() # type: ignore[attr-defined]
|
|
|
|
|
|
class TestResolvePlanInfoCredentialForwarding:
|
|
"""Credential forwarding must mirror extract_token: either header
|
|
authenticates, so either header must reach plan resolution — otherwise
|
|
x-api-key callers resolve no plan, silently disabling every
|
|
plan-conditioned control (free vs pro cost caps, the seat-scoped
|
|
credit-bucket exemption)."""
|
|
|
|
def _make_request(self, headers: dict[str, str]) -> MagicMock:
|
|
from starlette.datastructures import Headers
|
|
|
|
request = MagicMock()
|
|
request.headers = Headers(headers)
|
|
resolver = MagicMock()
|
|
resolver.get_plan = AsyncMock(
|
|
return_value=PlanInfo(plan_key="insights-code-pro-200-20260301", seat_created_at=None)
|
|
)
|
|
request.app.state.plan_resolver = resolver
|
|
return request
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_forwards_authorization_header(self) -> None:
|
|
request = self._make_request({"Authorization": "Bearer sk-test"})
|
|
result = await resolve_plan_info(request, user_id=1, product="insights_code")
|
|
assert result.plan_key == "insights-code-pro-200-20260301"
|
|
assert request.app.state.plan_resolver.get_plan.await_args.kwargs["auth_header"] == "Bearer sk-test"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_forwards_x_api_key_as_bearer(self) -> None:
|
|
request = self._make_request({"x-api-key": " sk-test "})
|
|
result = await resolve_plan_info(request, user_id=1, product="insights_code")
|
|
assert result.plan_key == "insights-code-pro-200-20260301"
|
|
assert request.app.state.plan_resolver.get_plan.await_args.kwargs["auth_header"] == "Bearer sk-test"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_x_api_key_takes_precedence_over_authorization(self) -> None:
|
|
# Matches extract_token — the token that authenticated the request is
|
|
# the one whose plan gets resolved.
|
|
request = self._make_request({"x-api-key": "sk-from_api_key", "Authorization": "Bearer sk-from_auth"})
|
|
await resolve_plan_info(request, user_id=1, product="insights_code")
|
|
assert request.app.state.plan_resolver.get_plan.await_args.kwargs["auth_header"] == "Bearer sk-from_api_key"
|