Add multi-plan subscriptions, quotas, and token usage APIs
Implements #16/#17/#36: Founders/Standard/Pro/Business/Backer catalog, Backer email whitelist, prompt-window + token-period gates, and tokens_in/out on conversation/prompt + subscription usage APIs.
This commit is contained in:
@@ -0,0 +1,261 @@
|
||||
"""Tests for multi-plan catalog, Backer whitelist, quotas, and subscription API."""
|
||||
|
||||
from datetime import timedelta
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from django.test import TestCase, override_settings
|
||||
from django.urls import reverse
|
||||
from django.utils import timezone
|
||||
from rest_framework import status
|
||||
from rest_framework.test import APITestCase
|
||||
|
||||
from chat_backend.models import PromptMetric
|
||||
from chat_backend.tests.factories import make_company, make_conversation, make_user
|
||||
from finance.models import BackerEmail, SubscriptionPlan, UserSubscription
|
||||
from finance.services.plans import (
|
||||
assign_plan,
|
||||
needs_checkout,
|
||||
seed_subscription_plans,
|
||||
try_redeem_backer_email,
|
||||
)
|
||||
from finance.services.quotas import (
|
||||
FeatureNotAllowed,
|
||||
QuotaExceeded,
|
||||
assert_feature_allowed,
|
||||
assert_within_quotas,
|
||||
get_usage_snapshot,
|
||||
)
|
||||
|
||||
|
||||
class PlanCatalogTestCase(TestCase):
|
||||
def test_seed_creates_expected_plans(self):
|
||||
plans = {p.slug: p for p in seed_subscription_plans()}
|
||||
self.assertEqual(
|
||||
set(plans),
|
||||
{"founders", "standard", "pro", "business", "backer"},
|
||||
)
|
||||
self.assertTrue(plans["founders"].is_public)
|
||||
self.assertTrue(plans["founders"].is_selectable)
|
||||
self.assertEqual(plans["founders"].price_cents, 1000)
|
||||
self.assertEqual(plans["founders"].prompt_quota_per_window, 300)
|
||||
|
||||
self.assertFalse(plans["standard"].is_public)
|
||||
self.assertEqual(plans["standard"].price_cents, 1500)
|
||||
self.assertEqual(plans["standard"].prompt_quota_per_window, 100)
|
||||
self.assertFalse(plans["standard"].allows_image_generation)
|
||||
|
||||
self.assertEqual(plans["pro"].price_cents, 4000)
|
||||
self.assertEqual(plans["pro"].prompt_quota_per_window, 200)
|
||||
self.assertTrue(plans["pro"].allows_image_generation)
|
||||
|
||||
self.assertEqual(plans["business"].price_cents, 9900)
|
||||
self.assertEqual(plans["business"].prompt_quota_per_window, 300)
|
||||
|
||||
self.assertEqual(plans["backer"].price_cents, 0)
|
||||
self.assertFalse(plans["backer"].is_selectable)
|
||||
self.assertTrue(plans["backer"].allows_all_future_features)
|
||||
|
||||
|
||||
class BackerRedeemTestCase(TestCase):
|
||||
def setUp(self):
|
||||
seed_subscription_plans()
|
||||
self.company = make_company()
|
||||
|
||||
def test_redeem_assigns_backer_and_skips_checkout(self):
|
||||
BackerEmail.objects.create(email="backer@example.com")
|
||||
user = make_user(email="backer@example.com", company=self.company)
|
||||
sub = try_redeem_backer_email(user)
|
||||
self.assertIsNotNone(sub)
|
||||
self.assertEqual(sub.plan.slug, "backer")
|
||||
self.assertEqual(sub.source, UserSubscription.Source.BACKER)
|
||||
self.assertFalse(needs_checkout(user))
|
||||
entry = BackerEmail.objects.get(email="backer@example.com")
|
||||
self.assertIsNotNone(entry.redeemed_at)
|
||||
self.assertEqual(entry.redeemed_user_id, user.pk)
|
||||
|
||||
def test_redeem_is_one_shot(self):
|
||||
BackerEmail.objects.create(email="once@example.com")
|
||||
user = make_user(email="once@example.com", company=self.company)
|
||||
self.assertIsNotNone(try_redeem_backer_email(user))
|
||||
self.assertIsNone(try_redeem_backer_email(user))
|
||||
|
||||
|
||||
@override_settings(ENFORCE_SUBSCRIPTION_GATES=True)
|
||||
class QuotaGateTestCase(TestCase):
|
||||
def setUp(self):
|
||||
seed_subscription_plans()
|
||||
self.company = make_company()
|
||||
self.user = make_user(company=self.company)
|
||||
self.plan = SubscriptionPlan.objects.get(slug="standard")
|
||||
assign_plan(
|
||||
self.user,
|
||||
plan=self.plan,
|
||||
source=UserSubscription.Source.ADMIN,
|
||||
)
|
||||
self.conversation = make_conversation(user=self.user)
|
||||
|
||||
def _add_metrics(self, count, *, tokens_in=None, tokens_out=None):
|
||||
now = timezone.now()
|
||||
for i in range(count):
|
||||
PromptMetric.objects.create(
|
||||
prompt_id=1000 + i,
|
||||
conversation_id=self.conversation.id,
|
||||
start_time=now,
|
||||
prompt_length=10,
|
||||
tokens_in=tokens_in,
|
||||
tokens_out=tokens_out,
|
||||
has_file=False,
|
||||
model_name="test",
|
||||
)
|
||||
|
||||
def test_prompt_quota_blocks(self):
|
||||
self._add_metrics(100)
|
||||
with self.assertRaises(QuotaExceeded) as ctx:
|
||||
assert_within_quotas(self.user)
|
||||
self.assertEqual(ctx.exception.code, "prompt_quota_exceeded")
|
||||
|
||||
def test_feature_gate_blocks_image_on_standard(self):
|
||||
with self.assertRaises(FeatureNotAllowed) as ctx:
|
||||
assert_feature_allowed(self.user, "image_generation")
|
||||
self.assertEqual(ctx.exception.code, "feature_not_allowed")
|
||||
|
||||
def test_pro_allows_image(self):
|
||||
pro = SubscriptionPlan.objects.get(slug="pro")
|
||||
assign_plan(
|
||||
self.user, plan=pro, source=UserSubscription.Source.ADMIN
|
||||
)
|
||||
assert_feature_allowed(self.user, "image_generation")
|
||||
|
||||
def test_token_quota_blocks_when_reported(self):
|
||||
self.plan.monthly_token_quota = 50
|
||||
self.plan.prompt_quota_per_window = 1000
|
||||
self.plan.save()
|
||||
self._add_metrics(1, tokens_in=30, tokens_out=30)
|
||||
with self.assertRaises(QuotaExceeded) as ctx:
|
||||
assert_within_quotas(self.user)
|
||||
self.assertEqual(ctx.exception.code, "token_quota_exceeded")
|
||||
|
||||
def test_null_tokens_do_not_fabricate_zero_usage(self):
|
||||
self._add_metrics(3, tokens_in=None, tokens_out=None)
|
||||
usage = get_usage_snapshot(self.user)
|
||||
self.assertIsNone(usage.tokens_in_period)
|
||||
self.assertIsNone(usage.tokens_out_period)
|
||||
self.assertIsNone(usage.tokens_total_period)
|
||||
self.assertGreaterEqual(usage.turns_missing_token_usage, 3)
|
||||
|
||||
|
||||
class PlanListAndSubscriptionApiTestCase(APITestCase):
|
||||
def setUp(self):
|
||||
seed_subscription_plans()
|
||||
self.company = make_company()
|
||||
self.user = make_user(company=self.company)
|
||||
|
||||
def test_public_plans_only_founders(self):
|
||||
url = reverse("finance_plans")
|
||||
response = self.client.get(url)
|
||||
self.assertEqual(response.status_code, status.HTTP_200_OK)
|
||||
slugs = [row["slug"] for row in response.data]
|
||||
self.assertEqual(slugs, ["founders"])
|
||||
|
||||
def test_subscription_me_includes_usage_nulls(self):
|
||||
self.client.force_authenticate(user=self.user)
|
||||
founders = SubscriptionPlan.objects.get(slug="founders")
|
||||
assign_plan(
|
||||
self.user,
|
||||
plan=founders,
|
||||
source=UserSubscription.Source.STRIPE,
|
||||
)
|
||||
url = reverse("finance_subscription_me")
|
||||
response = self.client.get(url)
|
||||
self.assertEqual(response.status_code, status.HTTP_200_OK)
|
||||
self.assertEqual(response.data["plan"]["slug"], "founders")
|
||||
self.assertFalse(response.data["needs_checkout"])
|
||||
self.assertIsNone(response.data["usage"]["tokens_in_period"])
|
||||
self.assertEqual(response.data["usage"]["prompt_quota"], 300)
|
||||
|
||||
|
||||
class CheckoutUsesFoundersPlanTestCase(APITestCase):
|
||||
def setUp(self):
|
||||
seed_subscription_plans()
|
||||
self.company = make_company()
|
||||
self.user = make_user(company=self.company)
|
||||
self.client.force_authenticate(user=self.user)
|
||||
self.url = reverse("finance_checkout")
|
||||
|
||||
@override_settings(
|
||||
STRIPE_SECRET_KEY="sk_test_fake",
|
||||
SUBSCRIPTION_PRICE_AMOUNT_CENTS=1000,
|
||||
SUBSCRIPTION_PRICE_CURRENCY="usd",
|
||||
SUBSCRIPTION_PRICE_INTERVAL="month",
|
||||
SUBSCRIPTION_PRODUCT_NAME="Founders",
|
||||
STRIPE_PRICE_ID="",
|
||||
STRIPE_CHECKOUT_SUCCESS_URL="http://localhost:3000/ok",
|
||||
STRIPE_CHECKOUT_CANCEL_URL="http://localhost:3000/cancel",
|
||||
)
|
||||
@patch("finance.services.stripe_service.stripe.checkout.Session.create")
|
||||
def test_checkout_defaults_to_founders(self, mock_create):
|
||||
mock_session = MagicMock()
|
||||
mock_session.id = "cs_test_founders"
|
||||
mock_session.url = "https://checkout.stripe.com/c/pay/cs_test_founders"
|
||||
mock_session.customer = None
|
||||
mock_create.return_value = mock_session
|
||||
|
||||
response = self.client.post(self.url, {}, format="json")
|
||||
self.assertEqual(response.status_code, status.HTTP_201_CREATED)
|
||||
self.assertEqual(response.data["plan_slug"], "founders")
|
||||
kwargs = mock_create.call_args.kwargs
|
||||
self.assertEqual(kwargs["metadata"]["plan_slug"], "founders")
|
||||
self.assertEqual(
|
||||
kwargs["line_items"][0]["price_data"]["unit_amount"], 1000
|
||||
)
|
||||
|
||||
def test_backer_cannot_checkout(self):
|
||||
backer = SubscriptionPlan.objects.get(slug="backer")
|
||||
assign_plan(
|
||||
self.user, plan=backer, source=UserSubscription.Source.BACKER
|
||||
)
|
||||
response = self.client.post(self.url, {}, format="json")
|
||||
self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
|
||||
self.assertFalse(response.data["needs_checkout"])
|
||||
|
||||
|
||||
class TokenSerializerApiTestCase(APITestCase):
|
||||
def setUp(self):
|
||||
self.company = make_company()
|
||||
self.user = make_user(company=self.company)
|
||||
self.client.force_authenticate(user=self.user)
|
||||
self.conversation = make_conversation(user=self.user, title="Tok")
|
||||
|
||||
def test_conversation_tokens_null_when_unreported(self):
|
||||
PromptMetric.objects.create(
|
||||
prompt_id=1,
|
||||
conversation_id=self.conversation.id,
|
||||
start_time=timezone.now(),
|
||||
prompt_length=5,
|
||||
tokens_in=None,
|
||||
tokens_out=None,
|
||||
has_file=False,
|
||||
model_name="t",
|
||||
)
|
||||
response = self.client.get(reverse("conversations"))
|
||||
self.assertEqual(response.status_code, status.HTTP_200_OK)
|
||||
row = next(r for r in response.data if r["id"] == self.conversation.id)
|
||||
self.assertIsNone(row["tokens_in"])
|
||||
self.assertIsNone(row["tokens_out"])
|
||||
|
||||
def test_conversation_tokens_sum_when_reported(self):
|
||||
for tin, tout, pid in ((10, 20, 1), (5, None, 2), (None, 7, 3)):
|
||||
PromptMetric.objects.create(
|
||||
prompt_id=pid,
|
||||
conversation_id=self.conversation.id,
|
||||
start_time=timezone.now(),
|
||||
prompt_length=5,
|
||||
tokens_in=tin,
|
||||
tokens_out=tout,
|
||||
has_file=False,
|
||||
model_name="t",
|
||||
)
|
||||
response = self.client.get(reverse("conversations"))
|
||||
row = next(r for r in response.data if r["id"] == self.conversation.id)
|
||||
self.assertEqual(row["tokens_in"], 15)
|
||||
self.assertEqual(row["tokens_out"], 27)
|
||||
Reference in New Issue
Block a user