feat: 修正优先级自动计算算法为分目标分组独立排序,并优化一键分组整理完成后仅触发一次网站级优先级重排

This commit is contained in:
liumangmang
2026-07-01 20:19:54 +08:00
parent fd46492012
commit dd3649a6ad
4 changed files with 536 additions and 280 deletions
+28 -10
View File
@@ -44,7 +44,8 @@ from app.services.website_client import Sub2ApiWebsiteClient, _extract_id
from app.services.website_sync import (
binding_sources,
sync_binding,
build_rate_priority_map,
build_target_group_priority_map,
sync_account_priorities_for_website,
reconcile_upstream_keys_full,
sync_account_priorities_for_upstream,
latest_rate_map,
@@ -646,12 +647,16 @@ def import_upstream_keys_as_accounts(
upstream_base_url = _u.base_url
# 按倍率自动分配优先级
rate_priority_map: dict[str, int] = {}
rate_priority_map: dict[tuple[str, int, str], int] = {}
if body.auto_priority_by_rate:
upstream_ids = {row.upstream_id for row in rows}
try:
rate_priority_map = _build_rate_priority_map(db, upstream_ids)
except HTTPException:
target_group_sources = {}
for row in rows:
target_group_id = body.target_group_map.get(row.group_id)
if target_group_id:
target_group_sources.setdefault(str(target_group_id), []).append((row.upstream_id, row.group_id))
rate_priority_map = build_target_group_priority_map(db, target_group_sources)
except Exception:
# 没有快照时忽略,后续 fallback 到 body.priority
pass
@@ -763,7 +768,7 @@ def import_upstream_keys_as_accounts(
"group_ids": group_ids,
"rate_multiplier": 1,
"concurrency": body.concurrency,
"priority": rate_priority_map.get(f"{row.upstream_id}:{row.group_id}", body.priority) if body.auto_priority_by_rate else body.priority,
"priority": rate_priority_map.get((str(target_group_id), row.upstream_id, row.group_id), body.priority) if body.auto_priority_by_rate else body.priority,
"notes": f"Imported by SmartUp from upstream key #{row.id}",
}
try:
@@ -863,10 +868,17 @@ def organize_website_groups(
except Exception as exc:
logger.warning("organize reconcile failed for upstream %s: %s", uid, exc)
# 3. 按倍率自动分配优先级
rate_priority_map: dict[str, int] = {}
# 3. 按倍率自动分配优先级 (分目标分组进行优先级计算)
rate_priority_map: dict[tuple[str, int, str], int] = {}
try:
rate_priority_map = _build_rate_priority_map(db, upstream_ids)
target_group_sources = {}
for b in bindings:
for src in binding_sources(b):
uid = src.get("upstream_id")
gid = src.get("group_id")
if uid and gid:
target_group_sources.setdefault(str(b.target_group_id), []).append((int(uid), str(gid)))
rate_priority_map = build_target_group_priority_map(db, target_group_sources)
except Exception:
pass
@@ -1110,7 +1122,7 @@ def organize_website_groups(
group_ids.append(target_group_id)
account_name = f"SmartUp-{row.group_name or row.group_id}-{row.id}"
priority_val = rate_priority_map.get(f"{row.upstream_id}:{row.group_id}", 1)
priority_val = rate_priority_map.get((str(target_group_id), row.upstream_id, row.group_id), 1)
account_body = {
"name": account_name,
"platform": platform,
@@ -1190,6 +1202,12 @@ def organize_website_groups(
if failed_count:
parts.append(f"失败 {failed_count}")
# 整理完成后,调用一次网站级 priority 重排以修正历史和当前账号优先级
try:
sync_account_priorities_for_website(db, wid)
except Exception as exc:
logger.warning("failed to sync priorities after organize for website %s: %s", wid, exc)
message = "整理完成:" + " / ".join(parts) if parts else "整理完成:无任何绑定或数据"
return OrganizeGroupsResponse(
success=failed_count == 0,
+336 -269
View File
@@ -432,12 +432,61 @@ def _try_send_priority_webhook(
logger.warning("account_priority_changed webhook failed for website %s: %s", wid, exc)
def sync_account_priorities_for_upstream(
def build_target_group_priority_map(
db: Session,
upstream_id: int,
website_id: int | None = None,
target_group_sources: dict[str, list[tuple[int, str]]],
) -> dict[tuple[str, int, str], int]:
"""根据目标分组内的上游分组倍率构建 (target_group_id, upstream_id, group_id) → priority 映射。
每个目标分组独立排序,倍率低到高对应 1, 11, 21, 31...
相同倍率共享同一 priority。
找不到倍率快照的来源不参与排序。
"""
priority_map: dict[tuple[str, int, str], int] = {}
# 提取所有涉及的 upstream_id
all_upstream_ids = set()
for sources in target_group_sources.values():
for upstream_id, _ in sources:
all_upstream_ids.add(upstream_id)
# 预载最新倍率快照
rate_maps = {}
for uid in all_upstream_ids:
try:
rate_maps[uid] = latest_rate_map(db, uid)
except Exception:
rate_maps[uid] = {}
for tg_id, sources in target_group_sources.items():
# 收集有有效倍率的来源
rated_sources: list[tuple[tuple[int, str], float]] = []
for upstream_id, group_id in sources:
groups = rate_maps.get(upstream_id) or {}
g = groups.get(group_id)
if isinstance(g, dict):
rate = _snapshot_group_rate(g)
rated_sources.append(((upstream_id, group_id), rate))
if not rated_sources:
continue
# 独立排序并赋予 priority (间隔 10)
unique_rates = sorted(set(r for _, r in rated_sources))
rate_to_priority = {rate: priority_for_rate_rank(idx) for idx, rate in enumerate(unique_rates)}
for (upstream_id, group_id), rate in rated_sources:
priority_map[(str(tg_id), upstream_id, str(group_id))] = rate_to_priority[rate]
return priority_map
def sync_account_priorities_for_website(
db: Session,
website_id: int,
trigger_upstream_id: int | None = None,
) -> list[dict]:
"""上游倍率变化后,自动更新已导入下游账号的 priority
"""计算并同步某个网站上所有账号的优先级
只处理同一目标分组内有多个账号(存在竞争)的情况:
- 竞争分组键:imported_target_group_id(老数据 fallback 到 group_id
@@ -447,6 +496,284 @@ def sync_account_priorities_for_upstream(
"""
from collections import defaultdict
website = db.query(Website).filter(Website.id == website_id).first()
if not website or not website.enabled:
logger.info("skip account priority sync: website %s not found or disabled", website_id)
return []
if trigger_upstream_id is not None:
upstream_name = db.query(Upstream.name).filter(Upstream.id == trigger_upstream_id).scalar() or f"#{trigger_upstream_id}"
else:
upstream_name = "分组整理/一键整理"
# ── 1. 建立远端 client 连接并加载账号与分组列表 ────────────────────────
account_info_map: dict[str, dict[str, Any]] = {}
target_group_names: dict[str, str] = {}
remote_account_ids: set[str] | None = None
site_results: list[dict] = []
competitive_buckets: dict[str, list] = {}
priority_assignment: dict[str, int] = {}
try:
with Sub2ApiWebsiteClient(
base_url=website.base_url,
api_prefix=website.api_prefix,
auth_type=website.auth_type,
auth_config=json.loads(website.auth_config_json or "{}"),
timeout=float(website.timeout_seconds),
) as client:
try:
remote_accounts = client.list_accounts()
if remote_accounts is not None:
remote_account_ids = set()
for acc in remote_accounts:
acc_id = _extract_id(acc)
if acc_id:
acc_id_str = str(acc_id)
remote_account_ids.add(acc_id_str)
account_info_map[acc_id_str] = {
"name": acc.get("name") or acc.get("account_name") or acc_id_str,
"priority": acc.get("priority"),
}
except Exception as e:
logger.warning("failed to fetch remote accounts for website %s: %s", website_id, e)
try:
grps = client.get_groups(endpoint=website.groups_endpoint) or []
for g in grps:
g_id = g.get("id")
if g_id is not None:
target_group_names[str(g_id)] = g.get("name") or str(g_id)
except Exception as e:
logger.warning("failed to fetch groups for website %s at sync: %s", website_id, e)
# ── 2. 查询该网站所有已导入 Key(跨上游),用于倍率查询 ────────────
all_website_keys = (
db.query(UpstreamGeneratedKey)
.filter(
UpstreamGeneratedKey.imported_website_id == website_id,
UpstreamGeneratedKey.imported_account_id.isnot(None),
UpstreamGeneratedKey.status != "orphaned",
)
.all()
)
_backfill_missing_target_groups_from_remote_accounts(db, website, all_website_keys)
# ── 3. 自愈校验:清除远端已不存在的账号,并从后续计算中排除 ──────────
active_website_keys: list[UpstreamGeneratedKey] = []
for row in all_website_keys:
account_id = row.imported_account_id
if remote_account_ids is not None and str(account_id) not in remote_account_ids:
tg_id_str = str(row.imported_target_group_id or row.group_id)
tg_name = target_group_names.get(tg_id_str) or row.imported_target_group_name
site_results.append(
_priority_result(
row,
None,
"stale",
"目标账号已不存在,已清除导入标记;请重新导入账号",
account_name=account_id,
old_priority=None,
source_rate=None,
target_group_name=tg_name,
)
)
# 清理本地导入标记并恢复为 created
row.imported_website_id = None
row.imported_account_id = None
row.imported_at = None
row.status = "created"
db.commit()
else:
active_website_keys.append(row)
# ── 4. 预取快照倍率 ────────────────────────────────────────────────
raw_rate_map: dict[str, float] = {}
if active_website_keys:
all_upstream_ids = {k.upstream_id for k in active_website_keys}
try:
# 构建 "{upstream_id}:{group_id}" → rate 查询表
for uid in all_upstream_ids:
groups = latest_rate_map(db, uid)
for gid, g in groups.items():
if isinstance(g, dict):
raw_rate_map[f"{uid}:{gid}"] = _snapshot_group_rate(g)
except Exception as exc:
logger.warning("build rate map failed for website %s: %s", website_id, exc)
# 补填 stale 账号结果的倍率
for res in site_results:
if res["status"] == "stale":
res["source_rate"] = raw_rate_map.get(f"{res['upstream_id']}:{res['group_id']}")
# ── 5. 基于仍存在的账号分桶并校验竞争 ──────────────────────────
buckets: dict[str, list[UpstreamGeneratedKey]] = defaultdict(list)
for row in active_website_keys:
comp_key = row.imported_target_group_id or row.group_id
buckets[comp_key].append(row)
# 只保留账号数 > 1 的分组(有竞争才需要排序)
competitive_buckets = {k: v for k, v in buckets.items() if len(v) > 1}
if not competitive_buckets:
logger.info(
"skip account priority sync for website %s: no competitive groups (all single-account)",
website_id,
)
else:
# ── 6. 每个竞争分组内独立计算 priority ──────────────────────
priority_assignment = {}
for comp_key, comp_rows in competitive_buckets.items():
# 只保留快照中能查到倍率的账号;无数据的账号不参与排序
rated = [
(row, raw_rate_map[f"{row.upstream_id}:{row.group_id}"])
for row in comp_rows
if f"{row.upstream_id}:{row.group_id}" in raw_rate_map
]
# 过滤后有效账号不足 2 个 → 此分组无竞争意义,整组跳过
if len(rated) < 2:
logger.info(
"skip competitive bucket %s for website %s: only %d account(s) have rate data",
comp_key, website_id, len(rated),
)
continue
# 组内按倍率升序排序(倍率低 → priority 小 → 优先)
unique_rates = sorted(set(r for _, r in rated))
rate_to_prio = {rate: priority_for_rate_rank(idx) for idx, rate in enumerate(unique_rates)}
for row, rate in rated:
priority_assignment[row.imported_account_id] = rate_to_prio[rate]
# ── 7. 更新竞争账号 the priority ──────────────────────────────
for comp_rows in competitive_buckets.values():
for row in comp_rows:
account_id = row.imported_account_id
new_priority = priority_assignment.get(account_id)
if new_priority is None:
continue
acc_info = account_info_map.get(str(account_id)) or {}
acc_name = acc_info.get("name") or str(account_id)
old_prio = acc_info.get("priority")
src_rate = raw_rate_map.get(f"{row.upstream_id}:{row.group_id}")
tg_id_str = str(row.imported_target_group_id or row.group_id)
tg_name = target_group_names.get(tg_id_str)
if tg_name and row.imported_target_group_name != tg_name:
row.imported_target_group_name = tg_name
db.commit()
try:
client.update_account(account_id, {"priority": new_priority})
logger.info(
"updated priority for account %s (website=%s, upstream=%s, group=%s"
", comp_group=%s): %s",
account_id, website_id, row.upstream_id, row.group_id,
row.imported_target_group_id or row.group_id, new_priority,
)
site_results.append(
_priority_result(
row,
new_priority,
"success",
f"优先级已更新为 {new_priority}",
account_name=acc_name,
old_priority=old_prio,
source_rate=src_rate,
target_group_name=tg_name or row.imported_target_group_name,
)
)
except Exception as exc:
# 兜底清理:如果 PUT 404
import httpx
is_404 = False
if isinstance(exc, WebsiteError) and "不存在或已被删除" in str(exc):
is_404 = True
elif isinstance(exc, WebsiteError) and isinstance(exc.__cause__, httpx.HTTPStatusError) and exc.__cause__.response.status_code == 404:
is_404 = True
if is_404:
site_results.append(
_priority_result(
row,
new_priority,
"stale",
"目标账号已不存在,已清除导入标记;请重新导入账号",
account_name=acc_name,
old_priority=old_prio,
source_rate=src_rate,
target_group_name=tg_name or row.imported_target_group_name,
)
)
row.imported_website_id = None
row.imported_account_id = None
row.imported_at = None
row.status = "created"
db.commit()
continue
logger.warning(
"failed to update priority for account %s (website=%s): %s",
account_id, website_id, exc,
)
site_results.append(
_priority_result(
row,
new_priority,
"failed",
str(exc),
account_name=acc_name,
old_priority=old_prio,
source_rate=src_rate,
target_group_name=tg_name or row.imported_target_group_name,
)
)
except Exception as exc:
logger.warning("failed to connect website %s for account priority sync: %s", website_id, exc)
key_query = db.query(UpstreamGeneratedKey).filter(
UpstreamGeneratedKey.imported_website_id == website_id,
UpstreamGeneratedKey.imported_account_id.isnot(None),
UpstreamGeneratedKey.status != "orphaned",
)
if trigger_upstream_id is not None:
key_query = key_query.filter(UpstreamGeneratedKey.upstream_id == trigger_upstream_id)
rows_to_fail = key_query.all()
for row in rows_to_fail:
site_results.append(
_priority_result(
row,
None,
"failed",
f"连接网站失败: {exc}",
)
)
# ── 8. 写入日志与钉钉通知 ──────────────────────────────────────────
notify_updates = [r for r in site_results if r["status"] in ("success", "failed")]
if notify_updates or any(r["status"] == "stale" for r in site_results):
priority_map_snapshot = {
f"{row.upstream_id}:{row.group_id}": priority_assignment.get(row.imported_account_id)
for comp_rows in competitive_buckets.values()
for row in comp_rows
if priority_assignment.get(row.imported_account_id) is not None
}
_write_priority_sync_log_with_map(db, website_id, upstream_name, site_results, priority_map_snapshot)
_try_send_priority_webhook(
db, website_id, website.name,
trigger_upstream_id or 0, upstream_name,
notify_updates,
)
return site_results
def sync_account_priorities_for_upstream(
db: Session,
upstream_id: int,
website_id: int | None = None,
) -> list[dict]:
"""上游倍率变化后,自动更新已导入下游账号的 priority。
"""
key_query = db.query(UpstreamGeneratedKey).filter(
UpstreamGeneratedKey.upstream_id == upstream_id,
UpstreamGeneratedKey.imported_website_id.isnot(None),
@@ -459,271 +786,11 @@ def sync_account_priorities_for_upstream(
if not key_rows:
return []
upstream_name = db.query(Upstream.name).filter(Upstream.id == upstream_id).scalar() or f"#{upstream_id}"
# 按 imported_website_id 分组
website_groups: dict[int, list[UpstreamGeneratedKey]] = {}
for row in key_rows:
wid = row.imported_website_id
website_groups.setdefault(wid, []).append(row)
all_results: list[dict] = []
for wid, rows in website_groups.items():
website = db.query(Website).filter(Website.id == wid).first()
if not website or not website.enabled:
logger.info("skip account priority sync: website %s not found or disabled", wid)
# 不写日志、不发通知:网站不可用时账号无法更新,沉默跳过
continue
# ── 1. 建立远端 client 连接并加载账号与分组列表 ────────────────────────
account_info_map: dict[str, dict[str, Any]] = {}
target_group_names: dict[str, str] = {}
remote_account_ids: set[str] | None = None
site_results: list[dict] = []
competitive_buckets: dict[str, list] = {}
priority_assignment: dict[str, int] = {}
try:
with Sub2ApiWebsiteClient(
base_url=website.base_url,
api_prefix=website.api_prefix,
auth_type=website.auth_type,
auth_config=json.loads(website.auth_config_json or "{}"),
timeout=float(website.timeout_seconds),
) as client:
try:
remote_accounts = client.list_accounts()
if remote_accounts is not None:
remote_account_ids = set()
for acc in remote_accounts:
acc_id = _extract_id(acc)
if acc_id:
acc_id_str = str(acc_id)
remote_account_ids.add(acc_id_str)
account_info_map[acc_id_str] = {
"name": acc.get("name") or acc.get("account_name") or acc_id_str,
"priority": acc.get("priority"),
}
except Exception as e:
logger.warning("failed to fetch remote accounts for website %s: %s", wid, e)
try:
grps = client.get_groups(endpoint=website.groups_endpoint) or []
for g in grps:
g_id = g.get("id")
if g_id is not None:
target_group_names[str(g_id)] = g.get("name") or str(g_id)
except Exception as e:
logger.warning("failed to fetch groups for website %s at sync: %s", wid, e)
# ── 2. 查询该网站所有已导入 Key(跨上游),用于倍率查询 ────────────
all_website_keys = (
db.query(UpstreamGeneratedKey)
.filter(
UpstreamGeneratedKey.imported_website_id == wid,
UpstreamGeneratedKey.imported_account_id.isnot(None),
UpstreamGeneratedKey.status != "orphaned",
)
.all()
)
_backfill_missing_target_groups_from_remote_accounts(db, website, all_website_keys)
# ── 3. 自愈校验:清除远端已不存在的账号,并从后续计算中排除 ──────────
active_website_keys: list[UpstreamGeneratedKey] = []
for row in all_website_keys:
account_id = row.imported_account_id
if remote_account_ids is not None and str(account_id) not in remote_account_ids:
tg_id_str = str(row.imported_target_group_id or row.group_id)
tg_name = target_group_names.get(tg_id_str) or row.imported_target_group_name
site_results.append(
_priority_result(
row,
None,
"stale",
"目标账号已不存在,已清除导入标记;请重新导入账号",
account_name=account_id,
old_priority=None,
source_rate=None,
target_group_name=tg_name,
)
)
# 清理本地导入标记并恢复为 created
row.imported_website_id = None
row.imported_account_id = None
row.imported_at = None
row.status = "created"
db.commit()
else:
active_website_keys.append(row)
# ── 4. 预取快照倍率 ────────────────────────────────────────────────
raw_rate_map: dict[str, float] = {}
if active_website_keys:
all_upstream_ids = {k.upstream_id for k in active_website_keys}
try:
# 构建 "{upstream_id}:{group_id}" → rate 查询表
for uid in all_upstream_ids:
groups = latest_rate_map(db, uid)
for gid, g in groups.items():
if isinstance(g, dict):
raw_rate_map[f"{uid}:{gid}"] = _snapshot_group_rate(g)
except Exception as exc:
logger.warning("build rate map failed for website %s: %s", wid, exc)
# 补填 stale 账号结果的倍率
for res in site_results:
if res["status"] == "stale":
res["source_rate"] = raw_rate_map.get(f"{res['upstream_id']}:{res['group_id']}")
# ── 5. 基于仍存在的账号分桶并校验竞争 ──────────────────────────
buckets: dict[str, list[UpstreamGeneratedKey]] = defaultdict(list)
for row in active_website_keys:
comp_key = row.imported_target_group_id or row.group_id
buckets[comp_key].append(row)
# 只保留账号数 > 1 的分组(有竞争才需要排序)
competitive_buckets = {k: v for k, v in buckets.items() if len(v) > 1}
if not competitive_buckets:
logger.info(
"skip account priority sync for website %s: no competitive groups (all single-account)",
wid,
)
else:
# ── 6. 每个竞争分组内独立计算 priority ──────────────────────
priority_assignment: dict[str, int] = {}
for comp_key, comp_rows in competitive_buckets.items():
# 只保留快照中能查到倍率的账号;无数据的账号不参与排序
rated = [
(row, raw_rate_map[f"{row.upstream_id}:{row.group_id}"])
for row in comp_rows
if f"{row.upstream_id}:{row.group_id}" in raw_rate_map
]
# 过滤后有效账号不足 2 个 → 此分组无竞争意义,整组跳过
if len(rated) < 2:
logger.info(
"skip competitive bucket %s for website %s: only %d account(s) have rate data",
comp_key, wid, len(rated),
)
continue
# 组内按倍率升序排序(倍率低 → priority 小 → 优先)
unique_rates = sorted(set(r for _, r in rated))
rate_to_prio = {rate: priority_for_rate_rank(idx) for idx, rate in enumerate(unique_rates)}
for row, rate in rated:
priority_assignment[row.imported_account_id] = rate_to_prio[rate]
# ── 7. 更新竞争账号的 priority ──────────────────────────────
for comp_rows in competitive_buckets.values():
for row in comp_rows:
account_id = row.imported_account_id
new_priority = priority_assignment.get(account_id)
if new_priority is None:
continue
acc_info = account_info_map.get(str(account_id)) or {}
acc_name = acc_info.get("name") or str(account_id)
old_prio = acc_info.get("priority")
src_rate = raw_rate_map.get(f"{row.upstream_id}:{row.group_id}")
tg_id_str = str(row.imported_target_group_id or row.group_id)
tg_name = target_group_names.get(tg_id_str)
if tg_name and row.imported_target_group_name != tg_name:
row.imported_target_group_name = tg_name
db.commit()
try:
client.update_account(account_id, {"priority": new_priority})
logger.info(
"updated priority for account %s (website=%s, upstream=%s, group=%s"
", comp_group=%s): %s",
account_id, wid, row.upstream_id, row.group_id,
row.imported_target_group_id or row.group_id, new_priority,
)
site_results.append(
_priority_result(
row,
new_priority,
"success",
f"优先级已更新为 {new_priority}",
account_name=acc_name,
old_priority=old_prio,
source_rate=src_rate,
target_group_name=tg_name or row.imported_target_group_name,
)
)
except Exception as exc:
# 兜底清理:如果 PUT 404
import httpx
is_404 = False
if isinstance(exc, WebsiteError) and "不存在或已被删除" in str(exc):
is_404 = True
elif isinstance(exc, WebsiteError) and isinstance(exc.__cause__, httpx.HTTPStatusError) and exc.__cause__.response.status_code == 404:
is_404 = True
if is_404:
site_results.append(
_priority_result(
row,
new_priority,
"stale",
"目标账号已不存在,已清除导入标记;请重新导入账号",
account_name=acc_name,
old_priority=old_prio,
source_rate=src_rate,
target_group_name=tg_name or row.imported_target_group_name,
)
)
row.imported_website_id = None
row.imported_account_id = None
row.imported_at = None
row.status = "created"
db.commit()
continue
logger.warning(
"failed to update priority for account %s (website=%s): %s",
account_id, wid, exc,
)
site_results.append(
_priority_result(
row,
new_priority,
"failed",
str(exc),
account_name=acc_name,
old_priority=old_prio,
source_rate=src_rate,
target_group_name=tg_name or row.imported_target_group_name,
)
)
except Exception as exc:
logger.warning("failed to connect website %s for account priority sync: %s", wid, exc)
for row in rows:
site_results.append(
_priority_result(
row,
None,
"failed",
f"连接网站失败: {exc}",
)
)
# ── 8. 写入日志与钉钉通知 ──────────────────────────────────────────
notify_updates = [r for r in site_results if r["status"] in ("success", "failed")]
if notify_updates or any(r["status"] == "stale" for r in site_results):
# 构建简化的 priority_map 供日志参考(只包含竞争分组的账号)
priority_map_snapshot = {
f"{row.upstream_id}:{row.group_id}": priority_assignment.get(row.imported_account_id)
for comp_rows in competitive_buckets.values()
for row in comp_rows
}
_write_priority_sync_log_with_map(db, wid, upstream_name, site_results, priority_map_snapshot)
if notify_updates:
_try_send_priority_webhook(db, wid, website.name, upstream_id, upstream_name, notify_updates)
all_results.extend(site_results)
return all_results
wids = sorted(list({row.imported_website_id for row in key_rows}))
results = []
for wid in wids:
results.extend(sync_account_priorities_for_website(db, wid, trigger_upstream_id=upstream_id))
return results
@@ -0,0 +1,171 @@
import pytest
import json
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.services.website_sync import (
build_target_group_priority_map,
)
from app.routers.websites import (
organize_website_groups,
import_upstream_keys_as_accounts,
)
from app.schemas.website import ImportAccountsRequest
from app.models.snapshot import UpstreamRateSnapshot
def _make_snapshot(db, upstream_id, rates):
snapshot_json = json.dumps({
"groups": {
gid: {"id": gid, "name": gid, "rate_multiplier": str(rate)}
for gid, rate in rates.items()
}
})
snap = UpstreamRateSnapshot(upstream_id=upstream_id, snapshot_json=snapshot_json)
db.add(snap)
db.commit()
@pytest.fixture()
def db_session():
engine = create_engine(
"sqlite://",
connect_args={"check_same_thread": False},
poolclass=StaticPool,
)
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_target_group_priority_calculation(db_session):
u1 = Upstream(name="U1", base_url="http://u1")
db_session.add(u1)
db_session.commit()
db_session.refresh(u1)
# Same target group, rates: 0.05, 0.065, 0.07, 0.08
_make_snapshot(db_session, u1.id, {
"G1": 0.05,
"G2": 0.065,
"G3": 0.07,
"G4": 0.08,
"G5": 0.05, # identical rate as G1
})
target_group_sources = {
"TG1": [
(u1.id, "G1"),
(u1.id, "G2"),
(u1.id, "G3"),
(u1.id, "G4"),
(u1.id, "G5"),
],
"TG2": [
(u1.id, "G2"), # rates 0.065 -> priority 1
(u1.id, "G3"), # rates 0.07 -> priority 11
]
}
priority_map = build_target_group_priority_map(db_session, target_group_sources)
# Assert priority values in TG1 (rates: 0.05, 0.065, 0.07, 0.08)
# 0.05 -> rank 0 -> priority 1
# 0.065 -> rank 1 -> priority 11
# 0.07 -> rank 2 -> priority 21
# 0.08 -> rank 3 -> priority 31
assert priority_map[("TG1", u1.id, "G1")] == 1
assert priority_map[("TG1", u1.id, "G5")] == 1 # identical rate shares priority
assert priority_map[("TG1", u1.id, "G2")] == 11
assert priority_map[("TG1", u1.id, "G3")] == 21
assert priority_map[("TG1", u1.id, "G4")] == 31
# Assert priority values in TG2 (starts from 1 independently)
# 0.065 -> rank 0 -> priority 1
# 0.07 -> rank 1 -> priority 11
assert priority_map[("TG2", u1.id, "G2")] == 1
assert priority_map[("TG2", u1.id, "G3")] == 11
def test_organize_website_groups_priority(db_session, monkeypatch):
w = Website(name="W1", base_url="http://w1", enabled=True, site_type="sub2api", auth_config_json="{}", timeout_seconds=30)
u = Upstream(name="U1", base_url="http://u1")
db_session.add_all([w, u])
db_session.commit()
db_session.refresh(w)
db_session.refresh(u)
_make_snapshot(db_session, u.id, {"G1": 0.05, "G2": 0.07})
b1 = WebsiteGroupBinding(
website_id=w.id,
target_group_id="TG1",
target_group_name="TG1",
source_groups_json=json.dumps([
{"upstream_id": u.id, "group_id": "G1"},
{"upstream_id": u.id, "group_id": "G2"},
]),
enabled=True
)
db_session.add(b1)
db_session.commit()
k1 = UpstreamGeneratedKey(upstream_id=u.id, group_id="G1", group_name="G1", key_name="K1", key_value="V1")
k2 = UpstreamGeneratedKey(upstream_id=u.id, group_id="G2", group_name="G2", key_name="K2", key_value="V2")
db_session.add_all([k1, k2])
db_session.commit()
created_accounts = []
updated_accounts = []
sync_called = []
class FakeClient:
def __init__(self, *args, **kwargs): pass
def __enter__(self): return self
def __exit__(self, *args): pass
def get_groups(self, *args, **kwargs):
return [{"id": "TG1", "name": "TG1"}]
def list_accounts(self):
return []
def create_account(self, body):
created_accounts.append(body)
return {"id": f"remote-{len(created_accounts)}", "name": body["name"]}
def update_account(self, account_id, body):
updated_accounts.append((account_id, body))
return {"id": account_id, "name": "acc"}
def extract_id(self, data):
return data.get("id")
monkeypatch.setattr("app.routers.websites._client", lambda website: FakeClient())
monkeypatch.setattr("app.services.website_client.Sub2ApiWebsiteClient", FakeClient)
# Mock website-level priority sync to assert it is called once
def mock_sync_website(db, website_id, trigger_upstream_id=None):
sync_called.append((website_id, trigger_upstream_id))
return []
monkeypatch.setattr("app.routers.websites.sync_account_priorities_for_website", mock_sync_website)
organize_website_groups(w.id, db_session)
# Assert k1 got priority 1, k2 got priority 11 (based on G1 rate 0.05 < G2 rate 0.07)
assert len(created_accounts) == 2
acc1 = next(a for a in created_accounts if "G1" in a["name"])
acc2 = next(a for a in created_accounts if "G2" in a["name"])
assert acc1["priority"] == 1
assert acc2["priority"] == 11
# Assert website level priority sync was triggered once at the end
assert len(sync_called) == 1
assert sync_called[0][0] == w.id
+1 -1
View File
@@ -480,7 +480,7 @@ def test_import_auto_priority_by_rate(db_session, monkeypatch):
req = ImportAccountsRequest(
upstream_key_ids=[k1.id, k2.id],
target_group_map={},
target_group_map={"G1": "TG1", "G2": "TG1"},
auto_priority_by_rate=True,
priority=10,
account_name_prefix="test",