307 lines
11 KiB
Python
307 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import sys
|
|
import unittest
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
|
|
from sqlalchemy.orm import sessionmaker
|
|
|
|
|
|
BACKEND_DIR = Path(__file__).resolve().parents[1]
|
|
os.environ["KEFU_DB_TYPE"] = "sqlite"
|
|
os.environ["KEFU_DATABASE_URL"] = ""
|
|
os.environ["KEFU_DB_PATH"] = str(BACKEND_DIR / "kefu.db")
|
|
if str(BACKEND_DIR) not in sys.path:
|
|
sys.path.insert(0, str(BACKEND_DIR))
|
|
|
|
import main
|
|
from models.database import Base
|
|
from models.models import Account, MessageLog
|
|
|
|
|
|
class _CountResult:
|
|
def __init__(self, count: int):
|
|
self.count = count
|
|
|
|
def scalar_one(self):
|
|
return self.count
|
|
|
|
|
|
class _RowsResult:
|
|
def __init__(self, rows):
|
|
self.rows = list(rows)
|
|
|
|
def scalars(self):
|
|
return self
|
|
|
|
def all(self):
|
|
return list(self.rows)
|
|
|
|
|
|
class AccountPaginationTests(unittest.IsolatedAsyncioTestCase):
|
|
async def test_legacy_cookie_sync_only_reads_accounts_missing_db_cookie(self):
|
|
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
|
|
async with engine.begin() as connection:
|
|
await connection.run_sync(Base.metadata.create_all)
|
|
session_factory = sessionmaker(
|
|
engine,
|
|
class_=AsyncSession,
|
|
expire_on_commit=False,
|
|
)
|
|
try:
|
|
async with session_factory() as db:
|
|
db.add_all(
|
|
[
|
|
Account(
|
|
id=2301,
|
|
username="already-in-db",
|
|
cookie_data='{"cookies": [{"value": "large"}]}',
|
|
im_session_data="x" * 100_000,
|
|
),
|
|
Account(id=2302, username="legacy-file", cookie_data=None),
|
|
]
|
|
)
|
|
await db.commit()
|
|
|
|
def read_cookie(account_id: int):
|
|
self.assertEqual(account_id, 2302)
|
|
return '{"cookies": [{"value": "migrated"}]}'
|
|
|
|
with (
|
|
patch.object(main, "AsyncSessionLocal", session_factory),
|
|
patch.object(main, "read_cookie_file", side_effect=read_cookie) as read_file,
|
|
patch.object(main, "get_cookie_path", return_value="legacy-2302.json"),
|
|
):
|
|
await main._sync_legacy_cookie_files()
|
|
|
|
read_file.assert_called_once_with(2302)
|
|
async with session_factory() as db:
|
|
migrated = await db.get(Account, 2302)
|
|
self.assertEqual(
|
|
migrated.cookie_data,
|
|
'{"cookies": [{"value": "migrated"}]}',
|
|
)
|
|
self.assertEqual(migrated.cookie_path, "legacy-2302.json")
|
|
finally:
|
|
await engine.dispose()
|
|
|
|
async def test_account_edit_invalidates_running_follow_config_cache(self):
|
|
account = SimpleNamespace(id=2201)
|
|
db = SimpleNamespace(commit=AsyncMock(), refresh=AsyncMock())
|
|
worker = SimpleNamespace(invalidate_follow_welcome_config=MagicMock())
|
|
original_workers = main.manager.workers
|
|
main.manager.workers = {2201: worker}
|
|
try:
|
|
with (
|
|
patch.object(main, "get_owned_account", AsyncMock(return_value=account)),
|
|
patch.object(main, "_build_account_response", return_value={"id": 2201}),
|
|
):
|
|
response = await main.update_account(
|
|
account_id=2201,
|
|
body=main.AccountUpdate(follow_welcome_enabled=True),
|
|
db=db,
|
|
user=SimpleNamespace(id=7, role="operator"),
|
|
)
|
|
|
|
self.assertEqual(response, {"id": 2201})
|
|
worker.invalidate_follow_welcome_config.assert_called_once_with()
|
|
finally:
|
|
main.manager.workers = original_workers
|
|
|
|
async def test_log_stats_uses_one_aggregate_and_respects_ownership(self):
|
|
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
|
|
async with engine.begin() as connection:
|
|
await connection.run_sync(Base.metadata.create_all)
|
|
session_factory = sessionmaker(
|
|
engine,
|
|
class_=AsyncSession,
|
|
expire_on_commit=False,
|
|
)
|
|
try:
|
|
async with session_factory() as db:
|
|
db.add_all(
|
|
[
|
|
Account(id=2101, owner_id=7, status="offline"),
|
|
Account(id=2102, owner_id=8, status="offline"),
|
|
MessageLog(account_id=2101, status="received"),
|
|
MessageLog(account_id=2101, status="replied"),
|
|
MessageLog(account_id=2102, status="replied"),
|
|
]
|
|
)
|
|
await db.commit()
|
|
|
|
with patch.object(db, "execute", wraps=db.execute) as execute:
|
|
stats = await main.get_logs_stats(
|
|
account_id=None,
|
|
db=db,
|
|
user=SimpleNamespace(id=7, role="user"),
|
|
)
|
|
|
|
self.assertEqual(stats, {"total": 2, "replied": 1})
|
|
self.assertEqual(execute.await_count, 1)
|
|
finally:
|
|
await engine.dispose()
|
|
|
|
async def test_account_options_returns_lightweight_runtime_fields(self):
|
|
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
|
|
async with engine.begin() as connection:
|
|
await connection.run_sync(Base.metadata.create_all)
|
|
session_factory = sessionmaker(
|
|
engine,
|
|
class_=AsyncSession,
|
|
expire_on_commit=False,
|
|
)
|
|
original_workers = main.manager.workers
|
|
main.manager.workers = {2001: SimpleNamespace(is_running=True)}
|
|
try:
|
|
async with session_factory() as db:
|
|
db.add_all(
|
|
[
|
|
Account(
|
|
id=2001,
|
|
owner_id=7,
|
|
username="owned",
|
|
status="offline",
|
|
cookie_data='{"cookies": []}',
|
|
im_session_data="x" * 100_000,
|
|
reply_cooldown_seconds=12,
|
|
),
|
|
Account(
|
|
id=2002,
|
|
owner_id=8,
|
|
username="other",
|
|
status="online",
|
|
cookie_data='{"cookies": []}',
|
|
),
|
|
]
|
|
)
|
|
await db.commit()
|
|
|
|
options = await main.get_account_options(
|
|
db=db,
|
|
user=SimpleNamespace(id=7, role="user"),
|
|
)
|
|
|
|
self.assertEqual(len(options), 1)
|
|
self.assertEqual(options[0].id, 2001)
|
|
self.assertEqual(options[0].status, "online")
|
|
self.assertTrue(options[0].has_cookie)
|
|
self.assertEqual(options[0].reply_cooldown_seconds, 12)
|
|
self.assertEqual(options[0].reply_cooldown_effective, 12)
|
|
self.assertFalse(hasattr(options[0], "im_session_data"))
|
|
finally:
|
|
main.manager.workers = original_workers
|
|
await engine.dispose()
|
|
|
|
async def test_paginated_list_counts_then_loads_only_current_page(self):
|
|
page_rows = [
|
|
SimpleNamespace(id=10, status="offline"),
|
|
SimpleNamespace(id=11, status="online"),
|
|
]
|
|
db = SimpleNamespace(
|
|
execute=AsyncMock(
|
|
side_effect=[
|
|
_CountResult(392),
|
|
_RowsResult(page_rows),
|
|
]
|
|
)
|
|
)
|
|
|
|
with (
|
|
patch.object(main.manager, "is_running", side_effect=[False, True]),
|
|
patch.object(
|
|
main,
|
|
"_build_account_response",
|
|
side_effect=lambda account: {"id": account.id, "status": account.status},
|
|
) as build_response,
|
|
):
|
|
response = await main.get_accounts(
|
|
page=20,
|
|
page_size=20,
|
|
q=None,
|
|
status=None,
|
|
db=db,
|
|
user=SimpleNamespace(id=1, role="admin"),
|
|
)
|
|
|
|
self.assertEqual(db.execute.await_count, 2)
|
|
self.assertEqual(response["total"], 392)
|
|
self.assertEqual(response["page"], 20)
|
|
self.assertEqual(response["page_size"], 20)
|
|
self.assertEqual([item["id"] for item in response["items"]], [10, 11])
|
|
self.assertEqual(build_response.call_count, 2)
|
|
|
|
count_sql = str(db.execute.await_args_list[0].args[0]).upper()
|
|
page_sql = str(db.execute.await_args_list[1].args[0]).upper()
|
|
self.assertIn("COUNT", count_sql)
|
|
self.assertNotIn(" LIMIT ", count_sql)
|
|
self.assertIn(" LIMIT ", page_sql)
|
|
self.assertIn(" OFFSET ", page_sql)
|
|
|
|
async def test_status_filter_uses_effective_runtime_worker_state(self):
|
|
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
|
|
async with engine.begin() as connection:
|
|
await connection.run_sync(Base.metadata.create_all)
|
|
session_factory = sessionmaker(
|
|
engine,
|
|
class_=AsyncSession,
|
|
expire_on_commit=False,
|
|
)
|
|
original_workers = main.manager.workers
|
|
main.manager.workers = {
|
|
1001: SimpleNamespace(is_running=True),
|
|
1002: SimpleNamespace(is_running=False),
|
|
}
|
|
try:
|
|
async with session_factory() as db:
|
|
db.add_all(
|
|
[
|
|
Account(id=1001, status="offline"),
|
|
Account(id=1002, status="online"),
|
|
]
|
|
)
|
|
await db.commit()
|
|
|
|
with patch.object(
|
|
main,
|
|
"_build_account_response",
|
|
side_effect=lambda account: {
|
|
"id": account.id,
|
|
"status": account.status,
|
|
},
|
|
):
|
|
online = await main.get_accounts(
|
|
page=1,
|
|
page_size=20,
|
|
q=None,
|
|
status="online",
|
|
db=db,
|
|
user=SimpleNamespace(id=1, role="admin"),
|
|
)
|
|
await db.rollback()
|
|
db.expire_all()
|
|
offline = await main.get_accounts(
|
|
page=1,
|
|
page_size=20,
|
|
q=None,
|
|
status="offline",
|
|
db=db,
|
|
user=SimpleNamespace(id=1, role="admin"),
|
|
)
|
|
|
|
self.assertEqual(online["total"], 1)
|
|
self.assertEqual(online["items"], [{"id": 1001, "status": "online"}])
|
|
self.assertEqual(offline["total"], 1)
|
|
self.assertEqual(offline["items"], [{"id": 1002, "status": "offline"}])
|
|
finally:
|
|
main.manager.workers = original_workers
|
|
await engine.dispose()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|