mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
perf(pool): 添加性能计时日志,优化批量删除分批与模型解除关联查询
- 前端批量操作对话框添加 loadAllKeys/executeAction 计时日志 - 后端池账号列表接口和批量删除接口添加分阶段耗时日志 - 批量删除按数据库类型自动选择分批大小,统一前端 batch size 为 2000 - 优化 auto_disassociate 查询:先检查 unlimited key 提前返回,使用 load_only 减少字段加载 - 新增批量操作路由和自动解除关联的单元测试
This commit is contained in:
94
tests/api/test_admin_pool_batch_action_routes.py
Normal file
94
tests/api/test_admin_pool_batch_action_routes.py
Normal file
@@ -0,0 +1,94 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from src.api.admin.pool.routes import AdminBatchActionKeysAdapter
|
||||
from src.api.admin.pool.schemas import BatchActionRequest
|
||||
|
||||
|
||||
def _build_context(db: MagicMock) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
db=db,
|
||||
user=SimpleNamespace(username="admin-1"),
|
||||
add_audit_metadata=lambda **_: None,
|
||||
)
|
||||
|
||||
|
||||
def _mock_provider_lookup(db: MagicMock, provider_id: str) -> None:
|
||||
db.query.return_value.filter.return_value.first.return_value = SimpleNamespace(id=provider_id)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_delete_uses_single_statement_for_non_sqlite(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
db = MagicMock()
|
||||
provider_id = "provider-1"
|
||||
_mock_provider_lookup(db, provider_id)
|
||||
db.get_bind.return_value = SimpleNamespace(dialect=SimpleNamespace(name="postgresql"))
|
||||
db.execute.return_value = SimpleNamespace(rowcount=1200)
|
||||
side_effect = AsyncMock()
|
||||
monkeypatch.setattr(
|
||||
"src.services.provider_keys.key_side_effects.run_delete_key_side_effects",
|
||||
side_effect,
|
||||
)
|
||||
|
||||
adapter = AdminBatchActionKeysAdapter(
|
||||
provider_id=provider_id,
|
||||
body=BatchActionRequest(
|
||||
key_ids=[f"key-{idx}" for idx in range(1200)],
|
||||
action="delete",
|
||||
),
|
||||
)
|
||||
|
||||
result = await adapter.handle(_build_context(db))
|
||||
|
||||
assert result.affected == 1200
|
||||
assert db.execute.call_count == 1
|
||||
db.commit.assert_called_once()
|
||||
side_effect.assert_awaited_once_with(
|
||||
db=db,
|
||||
provider_id=provider_id,
|
||||
deleted_key_allowed_models=None,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_delete_chunks_sqlite_and_runs_side_effect_once(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
db = MagicMock()
|
||||
provider_id = "provider-2"
|
||||
_mock_provider_lookup(db, provider_id)
|
||||
db.get_bind.return_value = SimpleNamespace(dialect=SimpleNamespace(name="sqlite"))
|
||||
db.execute.side_effect = [
|
||||
SimpleNamespace(rowcount=900),
|
||||
SimpleNamespace(rowcount=300),
|
||||
]
|
||||
side_effect = AsyncMock()
|
||||
monkeypatch.setattr(
|
||||
"src.services.provider_keys.key_side_effects.run_delete_key_side_effects",
|
||||
side_effect,
|
||||
)
|
||||
|
||||
adapter = AdminBatchActionKeysAdapter(
|
||||
provider_id=provider_id,
|
||||
body=BatchActionRequest(
|
||||
key_ids=[f"key-{idx}" for idx in range(1200)],
|
||||
action="delete",
|
||||
),
|
||||
)
|
||||
|
||||
result = await adapter.handle(_build_context(db))
|
||||
|
||||
assert result.affected == 1200
|
||||
assert db.execute.call_count == 2
|
||||
db.commit.assert_called_once()
|
||||
side_effect.assert_awaited_once_with(
|
||||
db=db,
|
||||
provider_id=provider_id,
|
||||
deleted_key_allowed_models=None,
|
||||
)
|
||||
93
tests/services/test_global_model_auto_disassociate.py
Normal file
93
tests/services/test_global_model_auto_disassociate.py
Normal file
@@ -0,0 +1,93 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from src.models.database import Model, Provider, ProviderAPIKey
|
||||
from src.services.model.global_model import GlobalModelService
|
||||
|
||||
|
||||
def test_auto_disassociate_short_circuits_when_unlimited_key_exists() -> None:
|
||||
db = MagicMock()
|
||||
provider_query = MagicMock()
|
||||
provider_query.filter.return_value.first.return_value = SimpleNamespace(name="Provider A")
|
||||
|
||||
unlimited_query = MagicMock()
|
||||
unlimited_query.filter.return_value.limit.return_value.first.return_value = object()
|
||||
|
||||
def _query(*entities: object) -> MagicMock:
|
||||
entity = entities[0]
|
||||
if entity is Provider:
|
||||
return provider_query
|
||||
if entity is ProviderAPIKey.id:
|
||||
return unlimited_query
|
||||
if entity is Model:
|
||||
raise AssertionError("model query should not run when unlimited key exists")
|
||||
raise AssertionError(f"unexpected query: {entities}")
|
||||
|
||||
db.query.side_effect = _query
|
||||
|
||||
result = GlobalModelService.auto_disassociate_provider_by_key_whitelist(db, "provider-1")
|
||||
|
||||
assert result == {"success": [], "errors": []}
|
||||
db.delete.assert_not_called()
|
||||
db.commit.assert_not_called()
|
||||
|
||||
|
||||
def test_auto_disassociate_deletes_unmatched_auto_associated_models(
|
||||
monkeypatch,
|
||||
) -> None:
|
||||
db = MagicMock()
|
||||
provider_query = MagicMock()
|
||||
provider_query.filter.return_value.first.return_value = SimpleNamespace(name="Provider B")
|
||||
|
||||
unlimited_query = MagicMock()
|
||||
unlimited_query.filter.return_value.limit.return_value.first.return_value = None
|
||||
|
||||
allowed_models_query = MagicMock()
|
||||
allowed_models_query.filter.return_value.all.return_value = [
|
||||
SimpleNamespace(allowed_models=["gpt-4o"]),
|
||||
SimpleNamespace(allowed_models=[]),
|
||||
]
|
||||
|
||||
model = SimpleNamespace(
|
||||
id="model-1",
|
||||
global_model=SimpleNamespace(
|
||||
id="gm-1",
|
||||
name="claude-sonnet",
|
||||
config={"model_mappings": ["claude-*"]},
|
||||
),
|
||||
)
|
||||
models_query = MagicMock()
|
||||
models_query.options.return_value.filter.return_value.all.return_value = [model]
|
||||
|
||||
def _query(*entities: object) -> MagicMock:
|
||||
entity = entities[0]
|
||||
if entity is Provider:
|
||||
return provider_query
|
||||
if entity is ProviderAPIKey.id:
|
||||
return unlimited_query
|
||||
if entity is ProviderAPIKey.allowed_models:
|
||||
return allowed_models_query
|
||||
if entity is Model:
|
||||
return models_query
|
||||
raise AssertionError(f"unexpected query: {entities}")
|
||||
|
||||
db.query.side_effect = _query
|
||||
monkeypatch.setattr(
|
||||
"src.core.model_permissions.match_model_with_pattern",
|
||||
lambda pattern, allowed_model: pattern == allowed_model,
|
||||
)
|
||||
|
||||
result = GlobalModelService.auto_disassociate_provider_by_key_whitelist(db, "provider-2")
|
||||
|
||||
assert result["errors"] == []
|
||||
assert result["success"] == [
|
||||
{
|
||||
"model_id": "model-1",
|
||||
"global_model_id": "gm-1",
|
||||
"global_model_name": "claude-sonnet",
|
||||
}
|
||||
]
|
||||
db.delete.assert_called_once_with(model)
|
||||
db.commit.assert_called_once()
|
||||
Reference in New Issue
Block a user