from datetime import timedelta from decimal import Decimal from cryptography.fernet import Fernet from django.contrib import admin from django.contrib.auth import get_user_model from django.core.exceptions import ValidationError from django.db import IntegrityError, transaction from django.test import TestCase, override_settings from django.utils import timezone from apps.ai.models import AiModel, ModelAlias from apps.billing.models import CallRecord, ExchangeRate, PointsLedger, PricingRule from apps.billing.pricing import ( NoPricingRuleError, calculate_points_cost, calculate_points_granted, get_pricing_rule, quote_recharge_points, ) from apps.users.models import ApiKey, UserWallet TEST_ENCRYPTION_KEY = Fernet.generate_key().decode("ascii") class BillingCoreModelTests(TestCase): def setUp(self): self.user = get_user_model().objects.create_user( username="client", email="client@example.com", password="password", ) def test_user_wallet_defaults_to_zero_and_rejects_negative_balance(self): wallet = UserWallet.objects.create(user=self.user) self.assertEqual(wallet.points_balance, 0) other_user = get_user_model().objects.create_user( username="negative", email="negative@example.com", password="password", ) with self.assertRaises(IntegrityError): with transaction.atomic(): UserWallet.objects.create(user=other_user, points_balance=-1) def test_api_key_hashes_plaintext_and_matches_only_raw_key(self): api_key, raw_key = ApiKey.create_for_user(self.user, name="desktop") self.assertTrue(raw_key.startswith("sk_cmhub_")) self.assertEqual(api_key.key_prefix, raw_key[: ApiKey.KEY_PREFIX_LENGTH]) self.assertEqual(len(api_key.key_hash), 64) self.assertNotEqual(api_key.key_hash, raw_key) self.assertNotIn(raw_key, str(api_key.__dict__)) self.assertTrue(api_key.matches_key(raw_key)) self.assertFalse(api_key.matches_key(raw_key + "-wrong")) self.assertEqual(api_key.status, ApiKey.Status.ACTIVE) def test_call_record_stores_summary_reference_but_no_provider_raw_field(self): api_key, _raw_key = ApiKey.create_for_user(self.user) call = CallRecord.objects.create( user=self.user, api_key=api_key, operation_type=CallRecord.OperationType.IMAGE, alias="image-hd", model_used="gpt-image-2", resolution="1K", prompt="Generate an image", points_cost=10, status=CallRecord.Status.SUCCESS, upstream_latency_ms=1234, result_ref="https://cdn.example.test/result.png", result_summary="stored image result", ) field_names = {field.name for field in CallRecord._meta.fields} self.assertNotIn("raw", field_names) self.assertNotIn("provider_raw", field_names) self.assertEqual(call.user, self.user) self.assertEqual(call.api_key, api_key) self.assertEqual(call.alias, "image-hd") self.assertEqual(call.model_used, "gpt-image-2") self.assertEqual(call.result_ref, "https://cdn.example.test/result.png") self.assertEqual(call.result_summary, "stored image result") def test_points_ledger_records_balance_and_requires_adjust_reason(self): call = CallRecord.objects.create( user=self.user, operation_type=CallRecord.OperationType.TITLE, alias="title-standard", model_used="gpt-5.5", points_cost=2, status=CallRecord.Status.SUCCESS, ) ledger = PointsLedger.objects.create( user=self.user, change_type=PointsLedger.ChangeType.CONSUME, points_delta=-2, balance_after=98, ref_call=call, ) self.assertEqual(ledger.ref_call, call) self.assertEqual(ledger.balance_after, 98) adjustment = PointsLedger( user=self.user, change_type=PointsLedger.ChangeType.ADJUST, points_delta=10, balance_after=108, ) with self.assertRaises(ValidationError): adjustment.full_clean() def test_points_ledger_rejects_zero_delta_and_negative_balance_after(self): with self.assertRaises(IntegrityError): with transaction.atomic(): PointsLedger.objects.create( user=self.user, change_type=PointsLedger.ChangeType.RECHARGE, points_delta=0, balance_after=100, ) with self.assertRaises(IntegrityError): with transaction.atomic(): PointsLedger.objects.create( user=self.user, change_type=PointsLedger.ChangeType.CONSUME, points_delta=-1, balance_after=-1, ) def test_billing_models_are_registered_in_admin(self): self.assertIn(UserWallet, admin.site._registry) self.assertIn(ApiKey, admin.site._registry) self.assertIn(PointsLedger, admin.site._registry) self.assertIn(CallRecord, admin.site._registry) @override_settings(AI_KEY_ENCRYPTION_KEY=TEST_ENCRYPTION_KEY) class PricingCalculationTests(TestCase): def setUp(self): self.text_model = self.create_ai_model( name="Text model A", model="gpt-5.5-a", capabilities=["text"], ) self.replacement_text_model = self.create_ai_model( name="Text model B", model="gpt-5.5-b", capabilities=["text"], ) self.alias = ModelAlias.objects.create( operation_type=ModelAlias.OperationType.TITLE, alias="title-standard", ai_model=self.text_model, ) def create_ai_model(self, *, name, model, capabilities): ai_model = AiModel( name=name, url="https://api.vectorengine.ai/v1", model=model, api_type=AiModel.ApiType.CHAT, capabilities=capabilities, ) ai_model.set_api_key("sk-test-secret") ai_model.save() return ai_model def test_exact_resolution_rule_overrides_alias_default(self): PricingRule.objects.create( operation_type=CallRecord.OperationType.IMAGE, alias="image-hd", resolution="", points_cost=10, ) PricingRule.objects.create( operation_type=CallRecord.OperationType.IMAGE, alias="image-hd", resolution="1k", points_cost=15, ) self.assertEqual( calculate_points_cost( CallRecord.OperationType.IMAGE, "image-hd", "1K", ), 15, ) self.assertEqual( calculate_points_cost( CallRecord.OperationType.IMAGE, "image-hd", "2K", ), 10, ) exact_rule = get_pricing_rule( CallRecord.OperationType.IMAGE, "image-hd", "1k", ) self.assertEqual(exact_rule.resolution, "1K") def test_pricing_is_bound_to_alias_not_underlying_model(self): PricingRule.objects.create( operation_type=CallRecord.OperationType.TITLE, alias="title-standard", points_cost=2, ) self.assertEqual( calculate_points_cost( CallRecord.OperationType.TITLE, "title-standard", "1K", ), 2, ) self.alias.ai_model = self.replacement_text_model self.alias.save() self.assertEqual( calculate_points_cost( CallRecord.OperationType.TITLE, "title-standard", "1K", ), 2, ) def test_missing_pricing_rule_raises_no_pricing_rule_code(self): with self.assertRaises(NoPricingRuleError) as context: calculate_points_cost( CallRecord.OperationType.TITLE, "missing-alias", "1K", ) self.assertEqual(context.exception.code, "no_pricing_rule") def test_exchange_rate_uses_latest_active_effective_rate_and_floors_points(self): now = timezone.now() ExchangeRate.objects.create( currency="cny", points_per_unit=Decimal("10.0000"), effective_from=now - timedelta(days=2), ) current = ExchangeRate.objects.create( currency="CNY", points_per_unit=Decimal("12.5000"), effective_from=now - timedelta(days=1), ) ExchangeRate.objects.create( currency="CNY", points_per_unit=Decimal("20.0000"), effective_from=now, is_active=False, ) ExchangeRate.objects.create( currency="CNY", points_per_unit=Decimal("99.0000"), effective_from=now + timedelta(days=1), ) quote = quote_recharge_points(Decimal("8.88"), currency="cny", at=now) self.assertEqual(quote.exchange_rate, current) self.assertEqual(quote.currency, "CNY") self.assertEqual(quote.points_per_unit, Decimal("12.5000")) self.assertEqual(quote.points_granted, 111) self.assertEqual(calculate_points_granted("8.88", currency="CNY", at=now), 111)