fix: preserve custom cookie input order
This commit is contained in:
@@ -0,0 +1,70 @@
|
||||
import os
|
||||
import unittest
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
os.environ.setdefault("DATABASE_URL", "sqlite://")
|
||||
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from web.backend.database import Base
|
||||
from web.backend.models import Account, LoginTask, User
|
||||
from web.backend.routers.cookies import _order_cookie_tasks
|
||||
|
||||
|
||||
class CustomCookieOrderTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.engine = create_engine("sqlite://")
|
||||
Base.metadata.create_all(self.engine)
|
||||
self.db = Session(self.engine)
|
||||
self.user = User(username="admin", password_hash="hash", role="super_admin")
|
||||
self.db.add(self.user)
|
||||
self.db.flush()
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
accounts = {
|
||||
name: Account(username=name, password="p", email="e", email_password="ep")
|
||||
for name in ("account-a", "account-b", "account-c")
|
||||
}
|
||||
self.db.add_all(accounts.values())
|
||||
self.db.flush()
|
||||
# 完成时间特意与用户输入顺序相反。
|
||||
for index, name in enumerate(("account-a", "account-b", "account-c")):
|
||||
self.db.add(LoginTask(
|
||||
batch_id=f"batch-{name}",
|
||||
account_id=accounts[name].id,
|
||||
status="success",
|
||||
created_by=self.user.id,
|
||||
finished_at=now + timedelta(minutes=index),
|
||||
))
|
||||
self.db.commit()
|
||||
|
||||
def tearDown(self):
|
||||
self.db.close()
|
||||
Base.metadata.drop_all(self.engine)
|
||||
self.engine.dispose()
|
||||
|
||||
def test_custom_cookie_order_follows_input_not_finished_time(self):
|
||||
selected_names = ["account-c", "account-a", "account-b"]
|
||||
tasks = (
|
||||
_order_cookie_tasks(
|
||||
self.db.query(LoginTask)
|
||||
.join(Account, LoginTask.account_id == Account.id)
|
||||
.filter(Account.username.in_(selected_names)),
|
||||
selected_names,
|
||||
)
|
||||
.all()
|
||||
)
|
||||
|
||||
account_names = {
|
||||
account.id: account.username
|
||||
for account in self.db.query(Account).all()
|
||||
}
|
||||
self.assertEqual(
|
||||
[account_names[task.account_id] for task in tasks],
|
||||
selected_names,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -6,7 +6,7 @@ from datetime import datetime, timezone
|
||||
import requests
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query
|
||||
from fastapi.responses import StreamingResponse
|
||||
from sqlalchemy import or_
|
||||
from sqlalchemy import case, or_
|
||||
from sqlalchemy.orm import Session, defer, joinedload
|
||||
import io
|
||||
import csv
|
||||
@@ -41,6 +41,19 @@ def _parse_account_names(raw_names: str) -> list[str]:
|
||||
seen.add(name)
|
||||
return names
|
||||
|
||||
|
||||
def _order_cookie_tasks(query, selected_names: list[str]):
|
||||
"""按自定义账号名的输入顺序排序,未指定时保持默认的最新优先。"""
|
||||
if selected_names:
|
||||
input_order = case(
|
||||
{name: index for index, name in enumerate(selected_names)},
|
||||
value=Account.username,
|
||||
else_=len(selected_names),
|
||||
)
|
||||
return query.order_by(input_order, LoginTask.finished_at.desc(), LoginTask.id.desc())
|
||||
return query.order_by(LoginTask.finished_at.desc(), LoginTask.id.desc())
|
||||
|
||||
|
||||
# 斗鱼 Cookie 有效性检测接口(抓包参考:钱包中心鱼丸余额 + 用户等级详情)
|
||||
CHECK_UA = (
|
||||
"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 "
|
||||
@@ -157,16 +170,18 @@ def list_cookies(
|
||||
)
|
||||
search_text = (search or "").strip()
|
||||
selected_names = _parse_account_names(account_names)
|
||||
can_view_all = user_has_permission(current, "login:view_all")
|
||||
account_joined = not can_view_all
|
||||
if selected_names:
|
||||
query = query.filter(
|
||||
LoginTask.account_id.in_(
|
||||
db.query(Account.id).filter(Account.username.in_(selected_names))
|
||||
)
|
||||
)
|
||||
if not account_joined:
|
||||
query = query.join(Account, LoginTask.account_id == Account.id)
|
||||
account_joined = True
|
||||
query = query.filter(Account.username.in_(selected_names))
|
||||
if search_text:
|
||||
pattern = f"%{search_text}%"
|
||||
if user_has_permission(current, "login:view_all"):
|
||||
if not account_joined:
|
||||
query = query.join(Account, LoginTask.account_id == Account.id)
|
||||
account_joined = True
|
||||
query = query.outerjoin(User, Account.assigned_to == User.id).filter(or_(
|
||||
Account.username.ilike(pattern),
|
||||
User.username.ilike(pattern),
|
||||
@@ -175,7 +190,7 @@ def list_cookies(
|
||||
total = None
|
||||
if page is not None:
|
||||
total = query.order_by(None).count()
|
||||
query = query.order_by(LoginTask.finished_at.desc())
|
||||
query = _order_cookie_tasks(query, selected_names)
|
||||
if page is not None:
|
||||
query = query.offset((page - 1) * page_size).limit(page_size)
|
||||
|
||||
@@ -331,12 +346,10 @@ def export_cookies(
|
||||
query = _visible_cookie_tasks_query(db, current)
|
||||
selected_names = _parse_account_names(account_names)
|
||||
if selected_names:
|
||||
query = query.filter(
|
||||
LoginTask.account_id.in_(
|
||||
db.query(Account.id).filter(Account.username.in_(selected_names))
|
||||
)
|
||||
)
|
||||
tasks = query.order_by(LoginTask.finished_at.desc()).all()
|
||||
if user_has_permission(current, "login:view_all"):
|
||||
query = query.join(Account, LoginTask.account_id == Account.id)
|
||||
query = query.filter(Account.username.in_(selected_names))
|
||||
tasks = _order_cookie_tasks(query, selected_names).all()
|
||||
|
||||
# 批量查账号,避免 N+1
|
||||
account_ids = [t.account_id for t in tasks]
|
||||
|
||||
Reference in New Issue
Block a user