yawan-workbench/server/tests/test_billing.py

import unittest

from sqlmodel import SQLModel, Session, create_engine, select

from app.billing import (amount_minor, check_quota, consume_points, grant_points,
                         point_state, record_usage_event, require_ai_points, usage_summary)
from app.llm import _log, chat
from app.models import (AppUser, ModelCatalog, ModelProvider, PointTransaction,
                        UsageEvent, UsageQuota)


class BillingTests(unittest.TestCase):
    def setUp(self):
        self.engine = create_engine("sqlite://", connect_args={"check_same_thread": False})
        SQLModel.metadata.create_all(self.engine)
        self.session = Session(self.engine)
        provider = ModelProvider(name="测试 Provider", base_url="https://example.test/v1",
                                 models=["test-model"], default_model="test-model")
        self.session.add(provider)
        self.session.commit()
        self.session.refresh(provider)
        self.provider = provider
        self.session.add(ModelCatalog(provider_id=provider.id, model_name="test-model",
                                      input_rate=2.0, output_rate=4.0, rate_version="2026-test"))
        self.session.commit()

    def tearDown(self):
        self.session.close()

    def test_idempotency_and_amount_are_immutable(self):
        self.assertEqual(amount_minor(1_000_000, 500_000, 2.0, 4.0), 400)
        first = record_usage_event(self.session, idempotency_key="same-key", user_id=7,
                                   organization_id=3, model="test-model", input_quantity=1_000_000,
                                   output_quantity=500_000, input_rate=2.0, output_rate=4.0)
        second = record_usage_event(self.session, idempotency_key="same-key", user_id=99,
                                    organization_id=99, model="other", quantity=900)
        self.session.commit()
        self.assertEqual(first.id, second.id)
        self.assertEqual(second.user_id, 7)
        self.assertEqual(second.amount_minor, 400)
        self.assertEqual(len(self.session.exec(select(UsageEvent)).all()), 1)

    def test_llm_log_links_user_org_and_usage_event(self):
        _log(self.session, 11, 12, 13, 14, "生产 Agent", self.provider.name, "test-model",
             "ok", "", 1000, 500, 20, "trace-1", "span-1", 1, "prompt", "output",
             provider_id=self.provider.id, user_id=7, organization_id=3)
        self.session.commit()
        call = self.session.exec(select(__import__("app.models", fromlist=["LLMCall"]).LLMCall)).one()
        event = self.session.get(UsageEvent, call.usage_event_id)
        self.assertIsNotNone(event)
        self.assertEqual((call.user_id, call.organization_id), (7, 3))
        self.assertEqual(event.pricing_version, "2026-test")
        self.assertEqual((event.input_quantity, event.output_quantity), (1000, 500))

    def test_quota_blocks_only_configured_overage(self):
        self.session.add(UsageQuota(user_id=7, organization_id=3, token_limit=10,
                                    overage_policy="hard_limit"))
        record_usage_event(self.session, idempotency_key="quota-key", user_id=7,
                           organization_id=3, quantity=10, input_quantity=10, status="measured")
        self.session.commit()
        with self.assertRaises(Exception):
            check_quota(self.session, 7, 3)
        self.assertTrue(check_quota(self.session, 88, 3)["allowed"])

    def test_points_are_idempotent_and_block_empty_accounts(self):
        first = grant_points(self.session, user_id=7, organization_id=3, points=100,
                             reason="测试充值", operator_user_id=1,
                             idempotency_key="credit-1")
        second = grant_points(self.session, user_id=7, organization_id=3, points=100,
                              reason="重复提交", operator_user_id=1,
                              idempotency_key="credit-1")
        self.session.commit()
        self.assertEqual(first.id, second.id)
        self.assertEqual(point_state(self.session, 7, 3)["token_equivalent"], 10)
        self.assertTrue(consume_points(self.session, user_id=7, organization_id=3,
                                       points=40, reason="测试调用", idempotency_key="debit-1"))
        self.session.commit()
        self.assertEqual(point_state(self.session, 7, 3)["points"], 60)
        self.assertEqual(len(self.session.exec(select(PointTransaction)).all()), 2)
        consume_points(self.session, user_id=7, organization_id=3,
                       points=60, reason="耗尽", idempotency_key="debit-2")
        self.session.commit()
        with self.assertRaises(ValueError):
            require_ai_points(self.session, user_id=7, organization_id=3)

    def test_chat_requires_points_before_provider_request(self):
        result = chat(self.session, [{"role": "user", "content": "测试"}],
                      self.provider.id, "test-model", user_id=7, organization_id=3,
                      max_retries=0)
        self.assertEqual(result["status"], "failed")
        self.assertIn("点数余额不足", result["error"])

    def test_successful_llm_log_debits_ten_points_per_token(self):
        grant_points(self.session, user_id=8, organization_id=3, points=100,
                     reason="测试充值", operator_user_id=1,
                     idempotency_key="credit-llm")
        _log(self.session, 0, 0, 0, 0, "测试 Agent", self.provider.name, "test-model",
             "ok", "", 2, 3, 20, "trace-points", "span-points", 1, "prompt", "output",
             provider_id=self.provider.id, user_id=8, organization_id=3)
        self.session.commit()
        self.assertEqual(point_state(self.session, 8, 3)["points"], 50)


if __name__ == "__main__":
    unittest.main()