mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +08:00
Merge remote-tracking branch 'origin/master' into dev
# Conflicts: # src/api/handlers/base/request_builder.py # src/api/handlers/openai_cli/adapter.py # src/services/system/maintenance_scheduler.py
This commit is contained in:
44
tests/services/billing/test_token_normalization.py
Normal file
44
tests/services/billing/test_token_normalization.py
Normal file
@@ -0,0 +1,44 @@
|
||||
from src.services.billing.token_normalization import normalize_input_tokens_for_billing
|
||||
from src.services.billing.usage_mapper import UsageMapper
|
||||
|
||||
|
||||
class TestNormalizeInputTokensForBilling:
|
||||
def test_openai_family_subtracts_cached_tokens(self) -> None:
|
||||
assert normalize_input_tokens_for_billing("openai:cli", 160_070, 81_664) == 78_406
|
||||
|
||||
def test_claude_family_does_not_change(self) -> None:
|
||||
assert normalize_input_tokens_for_billing("claude:cli", 160_070, 81_664) == 160_070
|
||||
|
||||
def test_gemini_family_subtracts_cached_tokens(self) -> None:
|
||||
# Gemini 的 promptTokenCount 包含 cachedContentTokenCount,需要扣除
|
||||
assert normalize_input_tokens_for_billing("gemini:chat", 323_392, 323_384) == 8
|
||||
assert normalize_input_tokens_for_billing("gemini:cli", 100, 20) == 80
|
||||
|
||||
def test_missing_format_does_not_change(self) -> None:
|
||||
assert normalize_input_tokens_for_billing(None, 100, 20) == 100
|
||||
assert normalize_input_tokens_for_billing("", 100, 20) == 100
|
||||
|
||||
def test_clamps_when_cached_tokens_exceed_input(self) -> None:
|
||||
assert normalize_input_tokens_for_billing("openai:cli", 10, 20) == 0
|
||||
assert normalize_input_tokens_for_billing("gemini:chat", 10, 20) == 0
|
||||
|
||||
|
||||
class TestUsageMapperOpenAICacheTokens:
|
||||
def test_openai_mapping_maps_cached_tokens_details(self) -> None:
|
||||
raw_usage = {
|
||||
"prompt_tokens": 100,
|
||||
"completion_tokens": 50,
|
||||
"prompt_tokens_details": {"cached_tokens": 20},
|
||||
}
|
||||
|
||||
usage = UsageMapper.map(raw_usage, api_format="openai:chat")
|
||||
assert usage.input_tokens == 100
|
||||
assert usage.output_tokens == 50
|
||||
assert usage.cache_read_tokens == 20
|
||||
|
||||
def test_openai_mapping_without_cached_tokens_is_unchanged(self) -> None:
|
||||
raw_usage = {"prompt_tokens": 100, "completion_tokens": 50}
|
||||
usage = UsageMapper.map(raw_usage, api_format="openai:chat")
|
||||
assert usage.input_tokens == 100
|
||||
assert usage.output_tokens == 50
|
||||
assert usage.cache_read_tokens == 0
|
||||
166
tests/services/test_maintenance_scheduler_user_quota_reset.py
Normal file
166
tests/services/test_maintenance_scheduler_user_quota_reset.py
Normal file
@@ -0,0 +1,166 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from src.services.system.maintenance_scheduler import MaintenanceScheduler
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_user_quota_reset_disabled(monkeypatch):
|
||||
scheduler = MaintenanceScheduler()
|
||||
|
||||
mock_db = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"src.services.system.maintenance_scheduler.create_session",
|
||||
lambda: mock_db,
|
||||
)
|
||||
|
||||
def fake_get_config(cls, db, key, default=None):
|
||||
if key == "enable_user_quota_reset":
|
||||
return False
|
||||
return default
|
||||
|
||||
mock_set_config = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"src.services.system.maintenance_scheduler.SystemConfigService.get_config",
|
||||
classmethod(fake_get_config),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"src.services.system.maintenance_scheduler.SystemConfigService.set_config",
|
||||
mock_set_config,
|
||||
)
|
||||
|
||||
await scheduler._perform_user_quota_reset()
|
||||
|
||||
assert not mock_db.query.called
|
||||
assert not mock_db.commit.called
|
||||
assert not mock_set_config.called
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_user_quota_reset_not_due_skips(monkeypatch):
|
||||
scheduler = MaintenanceScheduler()
|
||||
|
||||
mock_db = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"src.services.system.maintenance_scheduler.create_session",
|
||||
lambda: mock_db,
|
||||
)
|
||||
|
||||
last_reset_at = (datetime.now(timezone.utc) - timedelta(days=1)).isoformat()
|
||||
|
||||
def fake_get_config(cls, db, key, default=None):
|
||||
if key == "enable_user_quota_reset":
|
||||
return True
|
||||
if key == "user_quota_reset_interval_days":
|
||||
return 2
|
||||
if key == "user_quota_last_reset_at":
|
||||
return last_reset_at
|
||||
return default
|
||||
|
||||
mock_set_config = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"src.services.system.maintenance_scheduler.SystemConfigService.get_config",
|
||||
classmethod(fake_get_config),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"src.services.system.maintenance_scheduler.SystemConfigService.set_config",
|
||||
mock_set_config,
|
||||
)
|
||||
|
||||
await scheduler._perform_user_quota_reset()
|
||||
|
||||
assert not mock_db.query.called
|
||||
assert not mock_db.commit.called
|
||||
assert not mock_set_config.called
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_user_quota_reset_due_runs(monkeypatch):
|
||||
scheduler = MaintenanceScheduler()
|
||||
|
||||
mock_db = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"src.services.system.maintenance_scheduler.create_session",
|
||||
lambda: mock_db,
|
||||
)
|
||||
|
||||
last_reset_at = (datetime.now(timezone.utc) - timedelta(days=2)).isoformat()
|
||||
|
||||
def fake_get_config(cls, db, key, default=None):
|
||||
if key == "enable_user_quota_reset":
|
||||
return True
|
||||
if key == "user_quota_reset_interval_days":
|
||||
return 2
|
||||
if key == "user_quota_last_reset_at":
|
||||
return last_reset_at
|
||||
return default
|
||||
|
||||
mock_set_config = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"src.services.system.maintenance_scheduler.SystemConfigService.get_config",
|
||||
classmethod(fake_get_config),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"src.services.system.maintenance_scheduler.SystemConfigService.set_config",
|
||||
mock_set_config,
|
||||
)
|
||||
|
||||
mock_query = MagicMock()
|
||||
mock_filter = MagicMock()
|
||||
mock_filter.update.return_value = 7
|
||||
mock_query.filter.return_value = mock_filter
|
||||
mock_db.query.return_value = mock_query
|
||||
|
||||
await scheduler._perform_user_quota_reset()
|
||||
|
||||
mock_db.query.assert_called_once()
|
||||
mock_filter.update.assert_called_once()
|
||||
_, update_kwargs = mock_filter.update.call_args
|
||||
assert update_kwargs["synchronize_session"] is False
|
||||
mock_db.commit.assert_called_once()
|
||||
mock_set_config.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_user_quota_reset_invalid_interval_defaults_to_1(monkeypatch):
|
||||
scheduler = MaintenanceScheduler()
|
||||
|
||||
mock_db = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"src.services.system.maintenance_scheduler.create_session",
|
||||
lambda: mock_db,
|
||||
)
|
||||
|
||||
def fake_get_config(cls, db, key, default=None):
|
||||
if key == "enable_user_quota_reset":
|
||||
return True
|
||||
if key == "user_quota_reset_interval_days":
|
||||
return "abc"
|
||||
if key == "user_quota_last_reset_at":
|
||||
return None
|
||||
return default
|
||||
|
||||
mock_set_config = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"src.services.system.maintenance_scheduler.SystemConfigService.get_config",
|
||||
classmethod(fake_get_config),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"src.services.system.maintenance_scheduler.SystemConfigService.set_config",
|
||||
mock_set_config,
|
||||
)
|
||||
|
||||
mock_query = MagicMock()
|
||||
mock_filter = MagicMock()
|
||||
mock_filter.update.return_value = 1
|
||||
mock_query.filter.return_value = mock_filter
|
||||
mock_db.query.return_value = mock_query
|
||||
|
||||
await scheduler._perform_user_quota_reset()
|
||||
|
||||
mock_db.commit.assert_called_once()
|
||||
mock_set_config.assert_called_once()
|
||||
Reference in New Issue
Block a user