diff --git a/tests/test_cookie_custom_order.py b/tests/test_cookie_custom_order.py new file mode 100644 index 0000000..5327f8b --- /dev/null +++ b/tests/test_cookie_custom_order.py @@ -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() diff --git a/web/backend/routers/cookies.py b/web/backend/routers/cookies.py index 547b990..576ef7d 100644 --- a/web/backend/routers/cookies.py +++ b/web/backend/routers/cookies.py @@ -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]