from unittest.mock import MagicMock from src.models.database import GlobalModel, Model from src.services.billing.default_rules import DefaultBillingRuleGenerator from src.services.billing.formula_engine import FormulaEngine from src.services.billing.rule_service import BillingRuleService class TestDefaultBillingRuleGenerator: def test_default_rule_basic_chat_cost(self) -> None: global_model = GlobalModel( name="test-model", display_name="Test Model", is_active=True, default_price_per_request=0.01, default_tiered_pricing={ "tiers": [ { "up_to": None, "input_price_per_1m": 3.0, "output_price_per_1m": 15.0, # cache prices intentionally omitted (legacy derives from input price) } ] }, ) rule = DefaultBillingRuleGenerator.generate_for_model( global_model=global_model, model=None, task_type="chat", ) engine = FormulaEngine() result = engine.evaluate( expression=rule.expression, variables=rule.variables, dimensions={ "input_tokens": 1000, "output_tokens": 500, "cache_creation_tokens": 200, "cache_read_tokens": 300, "request_count": 1, # tier key "total_input_context": 1000 + 300, }, dimension_mappings=rule.dimension_mappings, strict_mode=True, ) assert result.status == "complete" assert abs(float(result.cost) - 0.02134) < 1e-9 def test_default_rule_tiered_pricing_uses_total_input_context(self) -> None: global_model = GlobalModel( name="tiered-model", display_name="Tiered Model", is_active=True, default_price_per_request=None, default_tiered_pricing={ "tiers": [ {"up_to": 200000, "input_price_per_1m": 3.0, "output_price_per_1m": 15.0}, {"up_to": None, "input_price_per_1m": 1.5, "output_price_per_1m": 7.5}, ] }, ) rule = DefaultBillingRuleGenerator.generate_for_model( global_model=global_model, model=None, task_type="chat", ) engine = FormulaEngine() result = engine.evaluate( expression=rule.expression, variables=rule.variables, dimensions={ "input_tokens": 250000, "output_tokens": 10000, "cache_creation_tokens": 0, "cache_read_tokens": 0, "request_count": 1, "total_input_context": 250000, }, dimension_mappings=rule.dimension_mappings, strict_mode=True, ) # 250000 * 1.5 / 1M = 0.375 # 10000 * 7.5 / 1M = 0.075 assert result.status == "complete" assert abs(float(result.cost) - 0.45) < 1e-9 def test_default_rule_cache_ttl_pricing_overrides_cache_read_price(self) -> None: global_model = GlobalModel( name="ttl-model", display_name="TTL Model", is_active=True, default_price_per_request=0.0, default_tiered_pricing={ "tiers": [ { "up_to": None, "input_price_per_1m": 3.0, "output_price_per_1m": 15.0, "cache_read_price_per_1m": 0.3, "cache_ttl_pricing": [ {"ttl_minutes": 5, "cache_read_price_per_1m": 0.3}, {"ttl_minutes": 60, "cache_read_price_per_1m": 0.5}, ], } ] }, ) rule = DefaultBillingRuleGenerator.generate_for_model( global_model=global_model, model=None, task_type="chat", ) engine = FormulaEngine() result = engine.evaluate( expression=rule.expression, variables=rule.variables, dimensions={ "input_tokens": 0, "output_tokens": 0, "cache_creation_tokens": 0, "cache_read_tokens": 1000, "cache_ttl_minutes": 60, "request_count": 1, "total_input_context": 0 + 1000, }, dimension_mappings=rule.dimension_mappings, strict_mode=True, ) # TTL=60 should use cache_read_price_per_1m=0.5 # 1000 * 0.5 / 1M = 0.0005 assert result.status == "complete" assert abs(float(result.cost) - 0.0005) < 1e-9 def test_default_rule_cache_ttl_pricing_overrides_cache_creation_price(self) -> None: global_model = GlobalModel( name="ttl-creation-model", display_name="TTL Creation Model", is_active=True, default_price_per_request=0.0, default_tiered_pricing={ "tiers": [ { "up_to": None, "input_price_per_1m": 3.0, "output_price_per_1m": 15.0, "cache_creation_price_per_1m": 3.75, "cache_ttl_pricing": [ {"ttl_minutes": 5, "cache_creation_price_per_1m": 3.75}, {"ttl_minutes": 60, "cache_creation_price_per_1m": 6.0}, ], } ] }, ) rule = DefaultBillingRuleGenerator.generate_for_model( global_model=global_model, model=None, task_type="chat", ) engine = FormulaEngine() # TTL=5: cache_creation_price_per_1m=3.75 # 1_000_000 * 3.75 / 1M = 3.75 result_5m = engine.evaluate( expression=rule.expression, variables=rule.variables, dimensions={ "input_tokens": 0, "output_tokens": 0, "cache_creation_tokens": 1_000_000, "cache_read_tokens": 0, "cache_ttl_minutes": 5, "request_count": 1, "total_input_context": 0, }, dimension_mappings=rule.dimension_mappings, strict_mode=True, ) assert result_5m.status == "complete" assert abs(float(result_5m.cost) - 3.75) < 1e-9 # TTL=60: cache_creation_price_per_1m=6.0 # 1_000_000 * 6.0 / 1M = 6.0 result_1h = engine.evaluate( expression=rule.expression, variables=rule.variables, dimensions={ "input_tokens": 0, "output_tokens": 0, "cache_creation_tokens": 1_000_000, "cache_read_tokens": 0, "cache_ttl_minutes": 60, "request_count": 1, "total_input_context": 0, }, dimension_mappings=rule.dimension_mappings, strict_mode=True, ) assert result_1h.status == "complete" assert abs(float(result_1h.cost) - 6.0) < 1e-9 class TestBillingRuleServiceDefaultFallback: def test_find_rule_returns_default_for_chat_when_no_db_rule(self) -> None: from src.services.billing.cache import BillingCache BillingCache.invalidate_all() global_model = GlobalModel( id="gm-1", name="test-model", display_name="Test Model", is_active=True, default_price_per_request=0.0, default_tiered_pricing={ "tiers": [{"up_to": None, "input_price_per_1m": 3.0, "output_price_per_1m": 15.0}] }, ) model_obj = Model( id="m-1", provider_id="p-1", global_model_id="gm-1", provider_model_name="provider-test-model", is_active=True, tiered_pricing=None, price_per_request=None, ) model_obj.global_model = global_model # Build a mock Session with deterministic query().filter().first() chain. q_global = MagicMock() q_global.filter.return_value.first.return_value = global_model q_model = MagicMock() q_model.filter.return_value.first.return_value = model_obj q_rule_model = MagicMock() q_rule_model.filter.return_value.first.return_value = None q_rule_global = MagicMock() q_rule_global.filter.return_value.first.return_value = None db = MagicMock() db.query.side_effect = [q_global, q_model, q_rule_model, q_rule_global] lookup = BillingRuleService.find_rule( db, provider_id="p-1", model_name="test-model", task_type="chat", ) assert lookup is not None assert lookup.scope == "default" assert lookup.rule.id == "__default__" assert lookup.effective_task_type == "chat" # Cached: second call should not touch db.query again. db.query.reset_mock() lookup2 = BillingRuleService.find_rule( db, provider_id="p-1", model_name="test-model", task_type="chat", ) assert lookup2 is not None assert lookup2.scope == "default" assert db.query.call_count == 0 def test_find_rule_returns_default_for_video_when_require_rule_false(self) -> None: from src.config.settings import config from src.services.billing.cache import BillingCache BillingCache.invalidate_all() old_require = config.billing_require_rule config.billing_require_rule = False try: global_model = GlobalModel( id="gm-2", name="video-model", display_name="Video Model", is_active=True, default_price_per_request=0.0, default_tiered_pricing={ "tiers": [ {"up_to": None, "input_price_per_1m": 3.0, "output_price_per_1m": 15.0} ] }, ) q_global = MagicMock() q_global.filter.return_value.first.return_value = global_model q_rule_global = MagicMock() q_rule_global.filter.return_value.first.return_value = None db = MagicMock() # provider_id omitted -> only GlobalModel query + global BillingRule query db.query.side_effect = [q_global, q_rule_global] lookup = BillingRuleService.find_rule( db, provider_id=None, model_name="video-model", task_type="video", ) assert lookup is not None assert lookup.scope == "default" assert lookup.rule.id == "__default__" assert lookup.effective_task_type == "video" finally: config.billing_require_rule = old_require def test_find_rule_returns_template_for_video_when_require_rule_true(self) -> None: from src.config.settings import config from src.services.billing.cache import BillingCache BillingCache.invalidate_all() old_require = config.billing_require_rule config.billing_require_rule = True try: global_model = GlobalModel( id="gm-3", name="video-model-2", display_name="Video Model 2", is_active=True, default_price_per_request=0.0, default_tiered_pricing={ "tiers": [ {"up_to": None, "input_price_per_1m": 3.0, "output_price_per_1m": 15.0} ] }, ) q_global = MagicMock() q_global.filter.return_value.first.return_value = global_model q_rule_global = MagicMock() q_rule_global.filter.return_value.first.return_value = None db = MagicMock() db.query.side_effect = [q_global, q_rule_global] lookup = BillingRuleService.find_rule( db, provider_id=None, model_name="video-model-2", task_type="video", ) assert lookup is not None assert lookup.scope == "default" assert lookup.rule.id == "__default__" # Universal Billing Rule is now used for all task types assert lookup.rule.name == "Universal Billing Rule" finally: config.billing_require_rule = old_require