Files
SmartUp/backend/test_organize_groups.py
T

706 lines
25 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
import json
import pytest
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.pool import StaticPool
from app.database import Base
from app.models.upstream import Upstream
from app.models.upstream_key import UpstreamGeneratedKey
from app.models.website import Website, WebsiteGroupBinding
from app.routers.websites import organize_website_groups
@pytest.fixture()
def db_session():
engine = create_engine(
"sqlite://",
connect_args={"check_same_thread": False},
poolclass=StaticPool,
)
# Ensure all models are created
from app.models import admin_user, upstream, snapshot, webhook_config, notification_log, custom_page, website, revoked_token, upstream_key
Base.metadata.create_all(bind=engine)
TestingSessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
db = TestingSessionLocal()
try:
yield db
finally:
db.close()
Base.metadata.drop_all(bind=engine)
def test_organize_groups_full_scenarios(db_session, monkeypatch):
"""测试一键整理分组的各种核心场景:
1. 已有 Key 未导入时,创建目标账号并绑定目标分组
2. 已导入且目标账号存在时,不重复创建(若未绑定目标分组则补齐绑定)
3. 已导入但目标账号不存在时,清理旧标记并重建
4. 绑定关系存在但没有对应 Key 时,返回 missing_key
5. 多个目标分组、多上游分组绑定时,正确分类
"""
# 1. 建立测试 Website 和 Upstream
w = Website(
name="W1",
site_type="sub2api",
base_url="http://w1",
enabled=True,
auth_config_json="{}",
timeout_seconds=30
)
u1 = Upstream(name="U1", base_url="http://u1")
db_session.add_all([w, u1])
db_session.commit()
db_session.refresh(w)
db_session.refresh(u1)
# 2. 建立绑定关系
# 绑定 1:目标分组 TG1 ↔ 上游分组 G1(有 Key 未导入),G2(没有 Key — 对应 missing_key
b1 = WebsiteGroupBinding(
website_id=w.id,
target_group_id="TG1",
target_group_name="TG1-Group",
source_groups_json=json.dumps([
{"upstream_id": u1.id, "group_id": "G1"},
{"upstream_id": u1.id, "group_id": "G2"},
]),
enabled=True
)
# 绑定 2:目标分组 TG2 ↔ 上游分组 G3(已导入且账号仍存在),G4(已导入但账号已删除)
b2 = WebsiteGroupBinding(
website_id=w.id,
target_group_id="TG2",
target_group_name="TG2-Group",
source_groups_json=json.dumps([
{"upstream_id": u1.id, "group_id": "G3"},
{"upstream_id": u1.id, "group_id": "G4"},
]),
enabled=True
)
db_session.add_all([b1, b2])
db_session.commit()
# 3. 建立上游 Key 记录
# G1: 未导入 Key
k1 = UpstreamGeneratedKey(
upstream_id=u1.id,
group_id="G1",
group_name="G1-Name",
key_name="Key-G1",
key_value="sk-g1-secret",
status="created",
)
# G3: 已导入且账号存在
k3 = UpstreamGeneratedKey(
upstream_id=u1.id,
group_id="G3",
group_name="G3-Name",
key_name="Key-G3",
key_value="sk-g3-secret",
status="imported",
imported_website_id=w.id,
imported_account_id="ACC-G3",
imported_target_group_id="TG2",
)
# G4: 已导入但账号已删除
k4 = UpstreamGeneratedKey(
upstream_id=u1.id,
group_id="G4",
group_name="G4-Name",
key_name="Key-G4",
key_value="sk-g4-secret",
status="imported",
imported_website_id=w.id,
imported_account_id="ACC-G4-DELETED",
imported_target_group_id="TG2",
)
db_session.add_all([k1, k3, k4])
db_session.commit()
# 4. Mock Sub2ApiWebsiteClient 交互
created_accounts = []
updated_accounts = []
class MockClient:
def __init__(self, **kwargs):
pass
def __enter__(self):
return self
def __exit__(self, *a):
pass
def get_groups(self, *a, **kw):
return [
{"id": "TG1", "name": "TG1-Group"},
{"id": "TG2", "name": "TG2-Group"},
]
def list_accounts(self):
# ACC-G3 仍存在,ACC-G4-DELETED 不在此列表中代表已被删除
# 另外,我们假设 ACC-G3 目前绑定的 group_ids 是 [999],即不包含 TG2 (或整型数值)
return [
{
"id": "ACC-G3",
"name": "SmartUp-G3-ACC",
"group_ids": [999],
}
]
def extract_id(self, val):
return val.get("id") if isinstance(val, dict) else str(val)
def create_account(self, body):
# 模拟创建账号
acc_id = f"NEW-{body['name']}"
new_acc = {"id": acc_id, "name": body["name"], "group_ids": body["group_ids"]}
created_accounts.append(new_acc)
return new_acc
def update_account(self, account_id, body):
# 模拟更新账号的分组绑定
updated_accounts.append((account_id, body))
return {"id": account_id, "group_ids": body.get("group_ids")}
monkeypatch.setattr("app.routers.websites.Sub2ApiWebsiteClient", MockClient)
monkeypatch.setattr("app.routers.websites.sync_account_priorities_for_website", lambda db, wid: [])
# 5. 执行一键整理
response = organize_website_groups(wid=w.id, db=db_session)
# 6. 断言结果
assert response.success is True
# 整理完成:已创建 1 / 已重建 1 / 已存在已迁移 1 / 缺少 Key 1
assert "已创建 1" in response.message
assert "已重建 1" in response.message
assert "已存在已迁移 1" in response.message
assert "缺少 Key 1" in response.message
items = response.items
assert len(items) == 4
# 按 target_group_id 分类校验
# TG1 分组下的 G1 (已创建)
item_g1 = next(item for item in items if item.target_group_id == "TG1" and item.source_group_id == "G1")
assert item_g1.status == "created"
assert item_g1.key_name == "Key-G1"
assert item_g1.account_id.startswith("NEW-U1-")
# TG1 分组下的 G2 (缺少 Key)
item_g2 = next(item for item in items if item.target_group_id == "TG1" and item.source_group_id == "G2")
assert item_g2.status == "missing_key"
assert item_g2.message == "请先生成上游 Key"
# TG2 分组下的 G3 (已存在并补齐)
item_g3 = next(item for item in items if item.target_group_id == "TG2" and item.source_group_id == "G3")
assert item_g3.status == "exists"
assert "已迁移到目标分组" in item_g3.message
# 校验是否触发了 update_account 补齐了 TG2
assert len(updated_accounts) == 1
assert updated_accounts[0][0] == "ACC-G3"
# TG2 (或者其数值) 应该在更新后的 group_ids 中
assert "TG2" in updated_accounts[0][1]["group_ids"] or 999 in updated_accounts[0][1]["group_ids"]
# TG2 分组下的 G4 (清理后重建)
item_g4 = next(item for item in items if item.target_group_id == "TG2" and item.source_group_id == "G4")
assert item_g4.status == "recreated"
assert item_g4.message == "清理后重建账号"
assert item_g4.account_id.startswith("NEW-U1-")
# 7. 检查数据库中 Key 记录的状态变化
db_session.refresh(k1)
db_session.refresh(k3)
db_session.refresh(k4)
assert k1.imported_website_id == w.id
assert k1.imported_account_id is not None
assert k1.imported_target_group_id == "TG1"
assert k1.status == "imported"
assert k3.imported_target_group_id == "TG2"
assert k4.imported_website_id == w.id
assert k4.imported_account_id is not None
assert k4.imported_target_group_id == "TG2"
assert k4.status == "imported"
def test_organize_groups_list_accounts_none_conservatively_skips(db_session, monkeypatch):
"""测试一键整理分组时,若远端账号列表拉取失败返回 None:
- 已导入的账号不得被清除或重建,而应该保守跳过(状态记为 failed,提示无法校验)。
- 未导入的账号仍能正常导入创建。
"""
w = Website(
name="W1",
site_type="sub2api",
base_url="http://w1",
enabled=True,
auth_config_json="{}",
timeout_seconds=30
)
u1 = Upstream(name="U1", base_url="http://u1")
db_session.add_all([w, u1])
db_session.commit()
db_session.refresh(w)
db_session.refresh(u1)
# 绑定关系:目标分组 TG1 ↔ 上游分组 G1(未导入),G3(已导入)
b1 = WebsiteGroupBinding(
website_id=w.id,
target_group_id="TG1",
target_group_name="TG1-Group",
source_groups_json=json.dumps([
{"upstream_id": u1.id, "group_id": "G1", "group_name": "G1-SourceGroup"},
{"upstream_id": u1.id, "group_id": "G3", "group_name": "G3-SourceGroup"},
]),
enabled=True
)
db_session.add_all([b1])
db_session.commit()
# k1: 未导入 Key
k1 = UpstreamGeneratedKey(
upstream_id=u1.id,
group_id="G1",
group_name="G1-SourceGroup",
key_name="Key-G1",
key_value="sk-g1-secret",
status="created",
)
# k3: 已导入 Key
k3 = UpstreamGeneratedKey(
upstream_id=u1.id,
group_id="G3",
group_name="G3-SourceGroup",
key_name="Key-G3",
key_value="sk-g3-secret",
status="imported",
imported_website_id=w.id,
imported_account_id="ACC-G3-EXISTING",
imported_target_group_id="TG1",
)
db_session.add_all([k1, k3])
db_session.commit()
created_accounts = []
class MockClient:
def __init__(self, **kwargs): pass
def __enter__(self): return self
def __exit__(self, *a): pass
def get_groups(self, *a, **kw):
return [{"id": "TG1", "name": "TG1-Group"}]
def list_accounts(self):
# 模拟获取失败返回 None
return None
def extract_id(self, val):
return val.get("id") if isinstance(val, dict) else str(val)
def create_account(self, body):
acc_id = f"NEW-{body['name']}"
new_acc = {"id": acc_id, "name": body["name"], "group_ids": body["group_ids"]}
created_accounts.append(new_acc)
return new_acc
monkeypatch.setattr("app.routers.websites.Sub2ApiWebsiteClient", MockClient)
monkeypatch.setattr("app.routers.websites.sync_account_priorities_for_website", lambda db, wid: [])
# 执行一键整理
response = organize_website_groups(wid=w.id, db=db_session)
# 断言结果:由于其中有 1 个已导入账号校验失败,response.success 应为 False(有失败项)
assert response.success is False
# 统计信息中应该有“已创建 1”与“失败 1”
assert "已创建 1" in response.message
assert "失败 1" in response.message
items = response.items
assert len(items) == 2
# G1 (未导入) 仍应该成功创建账号
item_g1 = next(item for item in items if item.source_group_id == "G1")
assert item_g1.status == "created"
assert item_g1.account_id.startswith("NEW-U1-")
assert item_g1.source_group_name == "G1-SourceGroup" # 校验 group_name 优先从绑定中获取!
# G3 (已导入) 应该保守跳过
item_g3 = next(item for item in items if item.source_group_id == "G3")
assert item_g3.status == "failed"
assert item_g3.message == "无法校验目标账号状态,已保守跳过"
# 校验 DB 中的 k3 导入标记和状态绝对不能被清除/变动!
db_session.refresh(k3)
assert k3.imported_website_id == w.id
assert k3.imported_account_id == "ACC-G3-EXISTING"
assert k3.status == "imported"
def test_organize_groups_force_alignment(db_session, monkeypatch):
"""测试强对齐场景:
1. 已存在账号在旧 SmartUp 分组,整理后移除旧分组并加入新分组。
2. 已存在账号同时有非 SmartUp 分组,整理后保留非 SmartUp 分组。
3. 已存在账号已经完全一致,整理不重复调用更新接口。
4. 多目标分组、多上游来源时,每个账号按自己的当前绑定关系对齐。
"""
w = Website(
name="W1",
site_type="sub2api",
base_url="http://w1",
enabled=True,
auth_config_json="{}",
timeout_seconds=30
)
u1 = Upstream(name="U1", base_url="http://u1")
db_session.add_all([w, u1])
db_session.commit()
db_session.refresh(w)
db_session.refresh(u1)
# 两个绑定关系:
# 绑定 1:目标分组 TG1 ↔ G1
# 绑定 2:目标分组 TG2 ↔ G2
b1 = WebsiteGroupBinding(
website_id=w.id,
target_group_id="TG1",
target_group_name="TG1-Group",
source_groups_json=json.dumps([{"upstream_id": u1.id, "group_id": "G1"}]),
enabled=True
)
b2 = WebsiteGroupBinding(
website_id=w.id,
target_group_id="TG2",
target_group_name="TG2-Group",
source_groups_json=json.dumps([{"upstream_id": u1.id, "group_id": "G2"}]),
enabled=True
)
db_session.add_all([b1, b2])
db_session.commit()
# 3. 建立 3 个 Key,对应 3 个不同的远端账号状态
# k1: 原本在 TG_old (旧 SmartUp 分组),但根据绑定 1 它现在应该在 TG1
k1 = UpstreamGeneratedKey(
upstream_id=u1.id,
group_id="G1",
group_name="G1-Group",
key_name="Key-G1",
key_value="sk-g1-secret",
status="imported",
imported_website_id=w.id,
imported_account_id="ACC-G1",
imported_target_group_id="TG_old", # 旧 SmartUp 分组
)
# k2: 已经在 TG2 (正确),且同时有非 SmartUp 分组 (999) 和另一个旧 SmartUp 分组 (TG1 - 应该被移除)
k2 = UpstreamGeneratedKey(
upstream_id=u1.id,
group_id="G2",
group_name="G2-Group",
key_name="Key-G2",
key_value="sk-g2-secret",
status="imported",
imported_website_id=w.id,
imported_account_id="ACC-G2",
imported_target_group_id="TG2",
)
db_session.add_all([k1, k2])
db_session.commit()
updated_accounts = []
class MockClient:
def __init__(self, **kwargs): pass
def __enter__(self): return self
def __exit__(self, *a): pass
def get_groups(self, *a, **kw):
return [
{"id": "TG1", "name": "TG1-Group"},
{"id": "TG2", "name": "TG2-Group"},
]
def list_accounts(self):
return [
{
"id": "ACC-G1",
"name": "SmartUp-G1",
"group_ids": ["TG_old"], # 包含旧 SmartUp 分组 TG_old,没有 TG1
},
{
"id": "ACC-G2",
"name": "SmartUp-G2",
"group_ids": ["TG2", 999, "TG1"], # TG2 是当前正确分组,999 是非托管分组,TG1 是旧托管分组(应移除)
}
]
def extract_id(self, val):
return val.get("id") if isinstance(val, dict) else str(val)
def update_account(self, account_id, body):
updated_accounts.append((account_id, body))
return {"id": account_id, "group_ids": body.get("group_ids")}
monkeypatch.setattr("app.routers.websites.Sub2ApiWebsiteClient", MockClient)
monkeypatch.setattr("app.routers.websites.sync_account_priorities_for_website", lambda db, wid: [])
# 执行一键整理
response = organize_website_groups(wid=w.id, db=db_session)
assert response.success is True
# ACC-G1 (G1) 的 group_ids 应变更为 ["TG1"] (TG2 被移除了)
# ACC-G2 (G2) 的 group_ids 应变更为 [999, "TG2"] (TG1 被移除了,999 保留)
assert len(updated_accounts) == 2
# 验证 ACC-G1 变更
up_g1 = next(item for item in updated_accounts if item[0] == "ACC-G1")
assert set(str(g) for g in up_g1[1]["group_ids"]) == {"TG1"}
# 验证 ACC-G2 变更
up_g2 = next(item for item in updated_accounts if item[0] == "ACC-G2")
assert set(str(g) for g in up_g2[1]["group_ids"]) == {"999", "TG2"}
# 验证本地 DB 的标记已被同步更新为正确的 TG1 / TG2
db_session.refresh(k1)
db_session.refresh(k2)
assert k1.imported_target_group_id == "TG1"
assert k2.imported_target_group_id == "TG2"
# 验证结果消息汇总包含 aligned 0, migrated 2 (因为两个都发生了迁移/更新)
assert "已存在已迁移 2" in response.message
# 让我们跑一次完全一致的再次整理,验证不触发 update_account (已对齐)
updated_accounts.clear()
class MockClientAligned(MockClient):
def list_accounts(self):
return [
{
"id": "ACC-G1",
"name": "SmartUp-G1",
"group_ids": ["TG1"],
},
{
"id": "ACC-G2",
"name": "SmartUp-G2",
"group_ids": [999, "TG2"],
}
]
monkeypatch.setattr("app.routers.websites.Sub2ApiWebsiteClient", MockClientAligned)
response_aligned = organize_website_groups(wid=w.id, db=db_session)
assert response_aligned.success is True
# 既然已经一致,不会触发 update_account
assert len(updated_accounts) == 0
assert "已存在且已对齐 2" in response_aligned.message
def test_organize_groups_corrects_mismatched_platform_instead_of_recreating(db_session, monkeypatch):
"""测试一键整理遇到已存在账号平台错配时,只调用 update_account() 修正平台,不创建新账号。"""
from app.models.snapshot import UpstreamRateSnapshot
w = Website(
name="W1",
site_type="sub2api",
base_url="http://w1",
enabled=True,
auth_config_json="{}",
timeout_seconds=30
)
u1 = Upstream(name="U1", base_url="http://u1")
db_session.add_all([w, u1])
db_session.commit()
db_session.refresh(w)
db_session.refresh(u1)
# 绑定关系:目标分组 TG1 ↔ G1
b1 = WebsiteGroupBinding(
website_id=w.id,
target_group_id="TG1",
target_group_name="TG1-Group",
source_groups_json=json.dumps([{"upstream_id": u1.id, "group_id": "G1"}]),
enabled=True
)
db_session.add(b1)
# G1: 已导入,导入的目标分组为 TG1,账号为 ACC-G1
k1 = UpstreamGeneratedKey(
upstream_id=u1.id,
group_id="G1",
group_name="G1-Group",
key_name="Key-G1",
key_value="sk-g1-secret",
status="imported",
imported_website_id=w.id,
imported_account_id="ACC-G1",
imported_target_group_id="TG1",
)
db_session.add(k1)
# 写入快照指定平台为 anthropic
snapshot = UpstreamRateSnapshot(
upstream_id=u1.id,
snapshot_json=json.dumps({
"groups": {
"G1": {
"group_name": "G1-Group",
"rate": 0.1,
"platform": "anthropic"
}
}
})
)
db_session.add(snapshot)
db_session.commit()
updated_accounts = []
create_called = False
mock_platform = "openai"
class MockClient:
def __init__(self, **kwargs): pass
def __enter__(self): return self
def __exit__(self, *a): pass
def get_groups(self, *a, **kw):
return [{"id": "TG1", "name": "TG1-Group"}]
def list_accounts(self):
return [
{
"id": "ACC-G1",
"name": "SmartUp-G1",
"group_ids": ["TG1"],
"platform": mock_platform,
}
]
def extract_id(self, val):
return val.get("id") if isinstance(val, dict) else str(val)
def create_account(self, body):
nonlocal create_called
create_called = True
return {"id": "NEW-ACC", "name": body["name"], "group_ids": body["group_ids"]}
def update_account(self, account_id, body):
nonlocal mock_platform
updated_accounts.append((account_id, body))
if "platform" in body:
mock_platform = body["platform"]
return {"id": account_id, "platform": body.get("platform")}
monkeypatch.setattr("app.routers.websites.Sub2ApiWebsiteClient", MockClient)
monkeypatch.setattr("app.routers.websites.sync_account_priorities_for_website", lambda db, wid: [])
response = organize_website_groups(wid=w.id, db=db_session)
assert response.success is True
assert not create_called
assert len(updated_accounts) == 1
assert updated_accounts[0][0] == "ACC-G1"
assert updated_accounts[0][1]["platform"] == "anthropic"
assert "平台从 openai 修正为 anthropic" in response.items[0].message
def _seed_platform_mismatch(db_session, monkeypatch, list_accounts_fn):
"""公用工厂:搭建平台错配场景并注入可定制的 list_accounts。"""
from app.models.snapshot import UpstreamRateSnapshot
w = Website(
name="W-PlatChk",
site_type="sub2api",
base_url="http://wp",
enabled=True,
auth_config_json="{}",
timeout_seconds=30,
)
u1 = Upstream(name="U-PlatChk", base_url="http://up")
db_session.add_all([w, u1])
db_session.commit()
db_session.refresh(w)
db_session.refresh(u1)
b1 = WebsiteGroupBinding(
website_id=w.id,
target_group_id="TG1",
target_group_name="TG1-Group",
source_groups_json=json.dumps([{"upstream_id": u1.id, "group_id": "G1"}]),
enabled=True,
)
db_session.add(b1)
k1 = UpstreamGeneratedKey(
upstream_id=u1.id,
group_id="G1",
group_name="G1-Group",
key_name="Key-G1",
key_value="sk-g1-secret",
status="imported",
imported_website_id=w.id,
imported_account_id="ACC-G1",
imported_target_group_id="TG1",
)
db_session.add(k1)
snapshot = UpstreamRateSnapshot(
upstream_id=u1.id,
snapshot_json=json.dumps({
"groups": {
"G1": {"group_name": "G1-Group", "rate": 0.1, "platform": "anthropic"}
}
}),
)
db_session.add(snapshot)
db_session.commit()
call_count = {"n": 0}
class MockClient:
def __init__(self, **kwargs): pass
def __enter__(self): return self
def __exit__(self, *a): pass
def get_groups(self, *a, **kw):
return [{"id": "TG1", "name": "TG1-Group"}]
def list_accounts(self):
call_count["n"] += 1
return list_accounts_fn(call_count["n"])
def extract_id(self, val):
return val.get("id") if isinstance(val, dict) else str(val)
def update_account(self, account_id, body):
return {"id": account_id}
monkeypatch.setattr("app.routers.websites.Sub2ApiWebsiteClient", MockClient)
monkeypatch.setattr(
"app.routers.websites.sync_account_priorities_for_website", lambda db, wid: []
)
return w
def test_organize_groups_platform_update_remote_silently_ignores(db_session, monkeypatch):
"""复查时远端仍返回旧平台值 → 必须报错,不能静默认为成功。"""
def list_accounts_fn(call_n):
# 第 1 次初始列表;第 2 次复查:仍是 openai(忽略了 platform 写入)
return [{"id": "ACC-G1", "name": "SmartUp-G1", "group_ids": ["TG1"], "platform": "openai"}]
w = _seed_platform_mismatch(db_session, monkeypatch, list_accounts_fn)
response = organize_website_groups(wid=w.id, db=db_session)
failed_items = [item for item in response.items if item.status == "failed"]
assert len(failed_items) == 1, f"期望 1 个 failed 项,实际:{response.items}"
assert "不支持修改平台属性" in failed_items[0].message
def test_organize_groups_platform_update_list_returns_none(db_session, monkeypatch):
"""复查时 list_accounts() 返回 None → 必须报错,不能静默成功。"""
def list_accounts_fn(call_n):
if call_n == 1:
return [{"id": "ACC-G1", "name": "SmartUp-G1", "group_ids": ["TG1"], "platform": "openai"}]
return None # 复查时列表拉取失败
w = _seed_platform_mismatch(db_session, monkeypatch, list_accounts_fn)
response = organize_website_groups(wid=w.id, db=db_session)
failed_items = [item for item in response.items if item.status == "failed"]
assert len(failed_items) == 1, f"期望 1 个 failed 项,实际:{response.items}"
assert "复查失败" in failed_items[0].message
def test_organize_groups_platform_update_account_missing_from_list(db_session, monkeypatch):
"""复查时账号从列表中消失 → 必须报错,不能静默成功。"""
def list_accounts_fn(call_n):
if call_n == 1:
return [{"id": "ACC-G1", "name": "SmartUp-G1", "group_ids": ["TG1"], "platform": "openai"}]
return [] # 复查时账号消失
w = _seed_platform_mismatch(db_session, monkeypatch, list_accounts_fn)
response = organize_website_groups(wid=w.id, db=db_session)
failed_items = [item for item in response.items if item.status == "failed"]
assert len(failed_items) == 1, f"期望 1 个 failed 项,实际:{response.items}"
assert "复查失败" in failed_items[0].message