安全:增加测试批次原子清理

This commit is contained in:
焦龙言
2026-09-23 14:52:07 +08:00
parent b11c457e61
commit 1aaed57825
6 changed files with 966 additions and 1 deletions
+17
View File
@@ -144,6 +144,15 @@ class TransportCryptoSettings(BaseSettings):
) )
class TestBatchSettings(BaseSettings):
"""
测试批次安全配置
"""
test_batch_preview_ttl_seconds: int = 600
test_batch_online_purge_max_rows: int = 0
class GenSettings: class GenSettings:
""" """
代码生成配置 代码生成配置
@@ -265,6 +274,12 @@ class GetConfig:
""" """
return TransportCryptoSettings() return TransportCryptoSettings()
def get_test_batch_config(self) -> TestBatchSettings:
"""
获取测试批次安全配置
"""
return TestBatchSettings()
def get_gen_config(self) -> GenSettings: def get_gen_config(self) -> GenSettings:
""" """
获取代码生成配置 获取代码生成配置
@@ -328,6 +343,8 @@ RedisConfig = get_config.get_redis_config()
LogConfig = get_config.get_log_config() LogConfig = get_config.get_log_config()
# 传输层加解密配置 # 传输层加解密配置
TransportCryptoConfig = get_config.get_transport_crypto_config() TransportCryptoConfig = get_config.get_transport_crypto_config()
# 测试批次安全配置
TestBatchConfig = get_config.get_test_batch_config()
# 代码生成配置 # 代码生成配置
GenConfig = get_config.get_gen_config() GenConfig = get_config.get_gen_config()
# 上传配置 # 上传配置
@@ -145,6 +145,10 @@ class TestBatchDao:
] ]
) )
@classmethod
async def delete_batch_users(cls, db: AsyncSession, batch_id: int) -> None:
await db.execute(delete(BizTestBatchUser).where(BizTestBatchUser.batch_id == batch_id))
@classmethod @classmethod
async def count_owned_rows(cls, db: AsyncSession, batch_id: int) -> dict[str, int]: async def count_owned_rows(cls, db: AsyncSession, batch_id: int) -> dict[str, int]:
counts: dict[str, int] = {} counts: dict[str, int] = {}
@@ -0,0 +1,76 @@
from collections.abc import Collection
from sqlalchemy import delete, select
from sqlalchemy.ext.asyncio import AsyncSession
from exceptions.exception import ServiceException
from module_business.domain.test_batch_graph import BATCH_MODELS, BatchSelection
from module_business.entity.do.master_do import BizUserSite, BizUserWarehouse
class TestBatchPurgeDao:
"""测试批次清理范围查询与无提交删除。"""
@classmethod
async def load_selection(
cls,
db: AsyncSession,
batch_id: int,
*,
for_update: bool,
) -> BatchSelection:
selection: BatchSelection = {}
for table_name, (model, primary_key, _) in BATCH_MODELS.items():
statement = select(primary_key).where(model.test_batch_id == batch_id).order_by(primary_key)
if for_update:
statement = statement.with_for_update()
result = await db.execute(statement)
selection[table_name] = frozenset(row[0] for row in result.all())
return selection
@classmethod
async def list_display_keys(
cls,
db: AsyncSession,
selection: BatchSelection,
) -> dict[str, list[str]]:
keys: dict[str, list[str]] = {}
for table_name, (_, primary_key, display_column) in BATCH_MODELS.items():
ids = sorted(selection[table_name])
if not ids:
keys[table_name] = []
continue
result = await db.execute(select(display_column).where(primary_key.in_(ids)).order_by(display_column))
keys[table_name] = sorted(str(row[0]) for row in result.all())
return keys
@classmethod
async def delete_user_warehouse_links(
cls,
db: AsyncSession,
warehouse_ids: Collection[int],
) -> None:
if warehouse_ids:
await db.execute(delete(BizUserWarehouse).where(BizUserWarehouse.warehouse_id.in_(sorted(warehouse_ids))))
@classmethod
async def delete_user_site_links(
cls,
db: AsyncSession,
site_ids: Collection[int],
) -> None:
if site_ids:
await db.execute(delete(BizUserSite).where(BizUserSite.site_id.in_(sorted(site_ids))))
@classmethod
async def delete_table_rows(
cls,
db: AsyncSession,
table_name: str,
batch_id: int,
) -> int:
if table_name not in BATCH_MODELS:
raise ServiceException(message=f'不允许清理非安全表: {table_name}')
model = BATCH_MODELS[table_name][0]
result = await db.execute(delete(model).where(model.test_batch_id == batch_id))
return int(result.rowcount or 0)
@@ -1,10 +1,11 @@
import hashlib import hashlib
import heapq
import json import json
import secrets import secrets
from collections import defaultdict, deque from collections import defaultdict, deque
from dataclasses import dataclass from dataclasses import dataclass
from sqlalchemy import Column, select from sqlalchemy import Column, MetaData, select
from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.ext.asyncio import AsyncSession
from config.database import Base from config.database import Base
@@ -98,6 +99,25 @@ BATCH_LABELS = {
'biz_pallet_event': '托盘事件', 'biz_pallet_event': '托盘事件',
} }
_DELETE_TIE_BREAK_ORDER = (
'biz_pallet_inventory',
'biz_usage_record',
'biz_return_item',
'biz_shipment_item',
'biz_inbound_item',
'biz_return_order',
'biz_shipment_order',
'biz_inbound_order',
'biz_pallet_event',
'biz_filling_record',
'biz_production_line',
'biz_customer_site',
'biz_customer',
'biz_warehouse',
'biz_strain',
'biz_pallet',
)
ROOT_TABLES = { ROOT_TABLES = {
'strain': 'biz_strain', 'strain': 'biz_strain',
'customer': 'biz_customer', 'customer': 'biz_customer',
@@ -135,6 +155,42 @@ def issue_preview_token() -> tuple[str, str]:
return raw, hashlib.sha256(raw.encode()).hexdigest() return raw, hashlib.sha256(raw.encode()).hexdigest()
def derive_batch_delete_order(metadata: MetaData) -> list[str]:
"""从实际外键图生成子表优先的确定性删除顺序。"""
safe_tables = set(BATCH_MODELS)
metadata_batch_tables = {table.name for table in metadata.tables.values() if 'test_batch_id' in table.columns}
if metadata_batch_tables != safe_tables or set(_DELETE_TIE_BREAK_ORDER) != safe_tables:
raise RuntimeError('测试批次清理安全表集合不一致')
children_by_parent: dict[str, set[str]] = {table_name: set() for table_name in safe_tables}
parents_by_child: dict[str, set[str]] = {table_name: set() for table_name in safe_tables}
for child_name in sorted(safe_tables):
for foreign_key in metadata.tables[child_name].foreign_keys:
parent_name = foreign_key.target_fullname.rsplit('.', maxsplit=1)[0]
if parent_name not in safe_tables:
continue
children_by_parent[parent_name].add(child_name)
parents_by_child[child_name].add(parent_name)
priority = {table_name: index for index, table_name in enumerate(_DELETE_TIE_BREAK_ORDER)}
ready = [(priority[table_name], table_name) for table_name, children in children_by_parent.items() if not children]
heapq.heapify(ready)
order: list[str] = []
while ready:
_, child_name = heapq.heappop(ready)
order.append(child_name)
for parent_name in sorted(parents_by_child[child_name]):
children_by_parent[parent_name].remove(child_name)
if not children_by_parent[parent_name]:
heapq.heappush(ready, (priority[parent_name], parent_name))
if len(order) != len(safe_tables):
raise RuntimeError('测试批次清理外键图存在环')
if set(order) != safe_tables:
raise RuntimeError('测试批次清理安全表集合不一致')
return order
@dataclass(frozen=True) @dataclass(frozen=True)
class _GraphEdge: class _GraphEdge:
source_table: str source_table: str
@@ -0,0 +1,196 @@
import hashlib
import hmac
import json
from datetime import datetime, timedelta
from typing import Any
from sqlalchemy.ext.asyncio import AsyncSession
from common.vo import CrudResponseModel
from config.database import Base
from config.env import TestBatchConfig
from exceptions.exception import PermissionException, ServiceException, ServiceWarning
from module_admin.entity.vo.user_vo import CurrentUserModel
from module_business.dao.test_batch_dao import TestBatchDao
from module_business.dao.test_batch_purge_dao import TestBatchPurgeDao
from module_business.domain.enums import TestBatchStatus
from module_business.domain.test_batch_graph import (
BATCH_LABELS,
BATCH_MODELS,
BatchSelection,
derive_batch_delete_order,
issue_preview_token,
selection_hash,
)
from module_business.entity.vo.test_batch_vo import (
BatchImpactModel,
PurgeExecuteRequestModel,
PurgePreviewModel,
)
from module_business.service.test_batch_access_service import TestBatchAccessService
class TestBatchPurgeService:
"""预览并原子清理已关闭的测试批次。"""
@staticmethod
def _assert_admin(current_user: CurrentUserModel) -> None:
TestBatchAccessService.assert_super_admin(current_user)
@staticmethod
def _user_identity(current_user: CurrentUserModel) -> tuple[int, str]:
user = current_user.user
if user is None or user.user_id is None or user.user_name is None:
raise PermissionException('', '当前用户信息不可用')
return user.user_id, user.user_name
@staticmethod
def _assert_closed(batch: Any | None) -> None:
if batch is None:
raise ServiceWarning('', '测试批次不存在')
if batch.status == TestBatchStatus.PURGED.value:
raise ServiceWarning('', '测试批次已清理')
if batch.status == TestBatchStatus.ACTIVE.value:
raise ServiceWarning('', '测试批次必须先关闭')
if batch.status != TestBatchStatus.CLOSED.value:
raise ServiceWarning('', '测试批次状态异常')
@staticmethod
def _assert_token(batch: object, raw_token: str, now: datetime) -> None:
supplied_hash = hashlib.sha256(raw_token.encode()).hexdigest()
stored_hash = getattr(batch, 'preview_token_hash', None)
comparable_hash = stored_hash if isinstance(stored_hash, str) else '0' * 64
if not hmac.compare_digest(comparable_hash, supplied_hash):
raise ServiceWarning('', '预览令牌无效')
if getattr(batch, 'preview_action', None) != 'PURGE':
raise ServiceWarning('', '预览令牌无效')
expires_at = getattr(batch, 'preview_expires_at', None)
if expires_at is None or expires_at <= now:
raise ServiceWarning('', '预览令牌已过期')
@staticmethod
def _counts(selection: BatchSelection) -> dict[str, int]:
return {table_name: len(selection[table_name]) for table_name in BATCH_MODELS}
@staticmethod
def _blocking_reason(total_count: int, limit: int) -> str | None:
if limit <= 0:
return '尚未配置在线测试批次清理上限'
if total_count > limit:
return '批次数据量超过在线清理上限,请按运维流程处理'
return None
@classmethod
async def preview(
cls,
db: AsyncSession,
batch_id: int,
current_user: CurrentUserModel,
) -> PurgePreviewModel:
cls._assert_admin(current_user)
_, user_name = cls._user_identity(current_user)
try:
batch = await TestBatchDao.get_batch_for_update(db, batch_id)
cls._assert_closed(batch)
selection = await TestBatchPurgeDao.load_selection(db, batch_id, for_update=False)
display_keys = await TestBatchPurgeDao.list_display_keys(db, selection)
counts = cls._counts(selection)
total_count = sum(counts.values())
limit = TestBatchConfig.test_batch_online_purge_max_rows
blocking_reason = cls._blocking_reason(total_count, limit)
now = datetime.now()
expires_at = now + timedelta(seconds=TestBatchConfig.test_batch_preview_ttl_seconds)
raw_token, token_hash = issue_preview_token()
await TestBatchDao.store_preview(
db,
batch,
action='PURGE',
token_hash=token_hash,
scope_hash=selection_hash(selection),
expires_at=expires_at,
update_by=user_name,
update_time=now,
)
impacts = [
BatchImpactModel(
table_name=table_name,
label=BATCH_LABELS[table_name],
count=counts[table_name],
keys=display_keys[table_name],
)
for table_name in BATCH_MODELS
]
response = PurgePreviewModel(
batch_id=batch_id,
batch_name=batch.batch_name,
preview_token=raw_token,
expires_at=expires_at,
total_count=total_count,
impacts=impacts,
online_purge_allowed=blocking_reason is None,
online_purge_limit=limit,
blocking_reason=blocking_reason,
)
await db.commit()
return response
except Exception:
await db.rollback()
raise
@classmethod
async def purge(
cls,
db: AsyncSession,
batch_id: int,
command: PurgeExecuteRequestModel,
current_user: CurrentUserModel,
) -> CrudResponseModel:
cls._assert_admin(current_user)
user_id, user_name = cls._user_identity(current_user)
try:
batch = await TestBatchDao.get_batch_for_update(db, batch_id)
if batch is not None and batch.status == TestBatchStatus.PURGED.value:
response = CrudResponseModel(is_success=True, message='该测试批次已清理')
await db.commit()
return response
cls._assert_closed(batch)
if command.batch_name != batch.batch_name:
raise ServiceWarning('', '测试批次名称不匹配')
cls._assert_token(batch, command.preview_token, datetime.now())
selection = await TestBatchPurgeDao.load_selection(db, batch_id, for_update=True)
if selection_hash(selection) != batch.preview_scope_hash:
raise ServiceWarning('', '预览范围已变化')
counts = cls._counts(selection)
total_count = sum(counts.values())
blocking_reason = cls._blocking_reason(
total_count,
TestBatchConfig.test_batch_online_purge_max_rows,
)
if blocking_reason is not None:
raise ServiceWarning('', blocking_reason)
await TestBatchPurgeDao.delete_user_warehouse_links(db, selection['biz_warehouse'])
await TestBatchPurgeDao.delete_user_site_links(db, selection['biz_customer_site'])
for table_name in derive_batch_delete_order(Base.metadata):
deleted = await TestBatchPurgeDao.delete_table_rows(db, table_name, batch_id)
if deleted != counts[table_name]:
raise ServiceException(message=f'{table_name} 清理数量发生变化')
await TestBatchDao.delete_batch_users(db, batch_id)
now = datetime.now()
batch.status = TestBatchStatus.PURGED.value
batch.purged_by = user_id
batch.purged_at = now
batch.purge_summary = json.dumps(counts, ensure_ascii=False, sort_keys=True)
batch.update_by = user_name
batch.update_time = now
await TestBatchDao.clear_preview(db, batch)
response = CrudResponseModel(is_success=True, message='测试批次清理成功')
await db.commit()
return response
except Exception:
await db.rollback()
raise
@@ -0,0 +1,616 @@
import asyncio
import hashlib
import inspect
import json
import unittest
from contextlib import ExitStack, contextmanager
from datetime import datetime, timedelta
from types import SimpleNamespace
from typing import Any, get_type_hints
from unittest.mock import AsyncMock, patch
from sqlalchemy import Column, ForeignKeyConstraint, Integer, MetaData, Table
from sqlalchemy.ext.asyncio import AsyncSession
from common.vo import CrudResponseModel
from config.database import Base
from exceptions.exception import PermissionException, ServiceException, ServiceWarning
from module_admin.entity.vo.user_vo import CurrentUserModel
from module_business.dao.test_batch_dao import TestBatchDao
from module_business.dao.test_batch_purge_dao import TestBatchPurgeDao
from module_business.domain.test_batch_graph import (
BATCH_LABELS,
BATCH_MODELS,
BatchSelection,
derive_batch_delete_order,
selection_hash,
)
from module_business.entity.vo.test_batch_vo import PurgeExecuteRequestModel, PurgePreviewModel
from module_business.service.test_batch_purge_service import TestBatchPurgeService
from tests.module_business.test_batch_fixtures import make_admin, make_current_user
EXPECTED_DELETE_ORDER = [
'biz_pallet_inventory',
'biz_usage_record',
'biz_return_item',
'biz_shipment_item',
'biz_inbound_item',
'biz_return_order',
'biz_shipment_order',
'biz_inbound_order',
'biz_pallet_event',
'biz_filling_record',
'biz_production_line',
'biz_customer_site',
'biz_customer',
'biz_warehouse',
'biz_strain',
'biz_pallet',
]
class _RowsResult:
def __init__(self, rows: list[tuple[Any, ...]]) -> None:
self._rows = rows
def all(self) -> list[tuple[Any, ...]]:
return self._rows
class _SelectionSession:
def __init__(self) -> None:
self.statements: list[Any] = []
async def execute(self, statement: Any) -> _RowsResult:
self.statements.append(statement)
return _RowsResult([])
def assert_child_tables_precede_parents(
test_case: unittest.TestCase,
order: list[str],
metadata: MetaData,
) -> None:
positions = {table_name: index for index, table_name in enumerate(order)}
for child_name in BATCH_MODELS:
for foreign_key in metadata.tables[child_name].foreign_keys:
parent_name = foreign_key.target_fullname.rsplit('.', maxsplit=1)[0]
if parent_name in BATCH_MODELS:
test_case.assertLess(
positions[child_name],
positions[parent_name],
f'{child_name} 必须先于父表 {parent_name} 删除',
)
def make_batch(**overrides: Any) -> SimpleNamespace:
values = {
'batch_id': 7,
'batch_name': '待清理测试批次',
'status': 'CLOSED',
'preview_action': None,
'preview_token_hash': None,
'preview_scope_hash': None,
'preview_expires_at': None,
'purged_by': None,
'purged_at': None,
'purge_summary': None,
'update_by': 'admin',
'update_time': datetime(2026, 9, 23, 9, 0, 0),
}
values.update(overrides)
return SimpleNamespace(**values)
def make_selection(**table_ids: set[int]) -> BatchSelection:
selection = {table_name: frozenset() for table_name in BATCH_MODELS}
for table_name, ids in table_ids.items():
selection[table_name] = frozenset(ids)
return selection
def make_token_batch(raw_token: str, selection: BatchSelection, **overrides: Any) -> SimpleNamespace:
values = {
'preview_action': 'PURGE',
'preview_token_hash': hashlib.sha256(raw_token.encode()).hexdigest(),
'preview_scope_hash': selection_hash(selection),
'preview_expires_at': datetime.now() + timedelta(minutes=10),
}
values.update(overrides)
return make_batch(**values)
def purge_command(raw_token: str) -> PurgeExecuteRequestModel:
return PurgeExecuteRequestModel(
batchName='待清理测试批次',
previewToken=raw_token,
)
def purge_config(limit: int, ttl: int = 600) -> SimpleNamespace:
return SimpleNamespace(
test_batch_preview_ttl_seconds=ttl,
test_batch_online_purge_max_rows=limit,
)
class TestBatchDeleteTopologyTest(unittest.TestCase):
def test_delete_order_is_derived_from_real_foreign_keys(self) -> None:
order = derive_batch_delete_order(Base.metadata)
self.assertEqual(EXPECTED_DELETE_ORDER, order)
assert_child_tables_precede_parents(self, order, Base.metadata)
def test_delete_order_fails_fast_when_safe_table_is_missing(self) -> None:
metadata = MetaData()
for table_name in list(BATCH_MODELS)[:-1]:
Table(
table_name,
metadata,
Column('id', Integer, primary_key=True),
Column('test_batch_id', Integer),
)
with self.assertRaisesRegex(RuntimeError, '安全表集合'):
derive_batch_delete_order(metadata)
def test_delete_order_fails_fast_on_foreign_key_cycle(self) -> None:
metadata = MetaData()
for table_name in BATCH_MODELS:
Table(
table_name,
metadata,
Column('id', Integer, primary_key=True),
Column('test_batch_id', Integer),
Column('parent_id', Integer),
)
metadata.tables['biz_strain'].append_constraint(ForeignKeyConstraint(['parent_id'], ['biz_customer.id']))
metadata.tables['biz_customer'].append_constraint(ForeignKeyConstraint(['parent_id'], ['biz_strain.id']))
with self.assertRaisesRegex(RuntimeError, '环'):
derive_batch_delete_order(metadata)
class TestBatchPurgeDaoTest(unittest.IsolatedAsyncioTestCase):
async def test_load_selection_reads_all_safe_tables_with_locking_reads(self) -> None:
db = _SelectionSession()
selection = await TestBatchPurgeDao.load_selection(db, 7, for_update=True)
self.assertEqual(set(BATCH_MODELS), set(selection))
self.assertEqual(len(BATCH_MODELS), len(db.statements))
for statement in db.statements:
sql = str(statement.compile(compile_kwargs={'literal_binds': True})).upper()
with self.subTest(sql=sql):
self.assertIn('TEST_BATCH_ID = 7', sql)
self.assertIn('FOR UPDATE', sql)
async def test_delete_table_rows_rejects_tables_outside_safe_set(self) -> None:
db = AsyncMock(spec=AsyncSession)
with self.assertRaises(ServiceException) as caught:
await TestBatchPurgeDao.delete_table_rows(db, 'sys_user', 7)
self.assertIn('不允许清理', caught.exception.message)
db.execute.assert_not_awaited()
async def test_user_links_are_deleted_by_selected_warehouse_and_site_ids(self) -> None:
db = AsyncMock(spec=AsyncSession)
await TestBatchPurgeDao.delete_user_warehouse_links(db, frozenset({4, 2}))
await TestBatchPurgeDao.delete_user_site_links(db, frozenset({9, 3}))
statements = [call.args[0] for call in db.execute.await_args_list]
sql = [str(statement.compile(compile_kwargs={'literal_binds': True})).lower() for statement in statements]
self.assertIn('delete from biz_user_warehouse', sql[0])
self.assertIn('warehouse_id in (2, 4)', sql[0])
self.assertIn('delete from biz_user_site', sql[1])
self.assertIn('site_id in (3, 9)', sql[1])
db.commit.assert_not_awaited()
class TestBatchPurgeContractTest(unittest.TestCase):
def test_public_service_signatures_are_stable(self) -> None:
expected = {
'preview': {
'db': AsyncSession,
'batch_id': int,
'current_user': CurrentUserModel,
'return': PurgePreviewModel,
},
'purge': {
'db': AsyncSession,
'batch_id': int,
'command': PurgeExecuteRequestModel,
'current_user': CurrentUserModel,
'return': CrudResponseModel,
},
}
expected_parameters = {
'preview': ('db', 'batch_id', 'current_user'),
'purge': ('db', 'batch_id', 'command', 'current_user'),
}
for method_name, hints in expected.items():
method = getattr(TestBatchPurgeService, method_name)
with self.subTest(method=method_name):
self.assertEqual(expected_parameters[method_name], tuple(inspect.signature(method).parameters))
self.assertEqual(hints, get_type_hints(method))
class TestBatchPurgeServiceTest(unittest.IsolatedAsyncioTestCase):
async def asyncSetUp(self) -> None:
self.db = AsyncMock(spec=AsyncSession)
self.selection = make_selection(
biz_warehouse={2},
biz_customer_site={3},
biz_pallet={10},
biz_usage_record={50},
)
self.keys = {table_name: [] for table_name in BATCH_MODELS}
self.keys['biz_warehouse'] = ['W-2']
self.keys['biz_customer_site'] = ['SITE-3']
self.keys['biz_pallet'] = ['P-10']
self.keys['biz_usage_record'] = ['U-50']
@contextmanager
def _preview_stack(
self,
batch: SimpleNamespace,
*,
config: SimpleNamespace | None = None,
) -> Any:
contexts = (
patch.object(TestBatchDao, 'get_batch_for_update', AsyncMock(return_value=batch)),
patch.object(TestBatchPurgeDao, 'load_selection', AsyncMock(return_value=self.selection)),
patch.object(TestBatchPurgeDao, 'list_display_keys', AsyncMock(return_value=self.keys)),
patch(
'module_business.service.test_batch_purge_service.TestBatchConfig',
config or purge_config(10_000),
),
)
with ExitStack() as stack:
for context in contexts:
stack.enter_context(context)
yield
@contextmanager
def _purge_stack(
self,
batch: SimpleNamespace,
*,
selection: BatchSelection | None = None,
delete_rows: AsyncMock | None = None,
limit: int = 10_000,
) -> Any:
selected = selection or self.selection
row_delete = delete_rows or AsyncMock(side_effect=lambda _db, table_name, _batch_id: len(selected[table_name]))
mocks = {
'delete_warehouse_links': AsyncMock(),
'delete_site_links': AsyncMock(),
'delete_rows': row_delete,
'delete_batch_users': AsyncMock(),
}
contexts = (
patch.object(TestBatchDao, 'get_batch_for_update', AsyncMock(return_value=batch)),
patch.object(TestBatchPurgeDao, 'load_selection', AsyncMock(return_value=selected)),
patch.object(
TestBatchPurgeDao,
'delete_user_warehouse_links',
mocks['delete_warehouse_links'],
),
patch.object(TestBatchPurgeDao, 'delete_user_site_links', mocks['delete_site_links']),
patch.object(TestBatchPurgeDao, 'delete_table_rows', mocks['delete_rows']),
patch.object(TestBatchDao, 'delete_batch_users', mocks['delete_batch_users']),
patch(
'module_business.service.test_batch_purge_service.TestBatchConfig',
purge_config(limit),
),
)
with ExitStack() as stack:
for context in contexts:
stack.enter_context(context)
yield mocks
async def test_regular_user_with_purge_permission_is_rejected_before_locking(self) -> None:
lock_batch = AsyncMock()
with (
patch.object(TestBatchDao, 'get_batch_for_update', lock_batch),
self.assertRaisesRegex(PermissionException, '仅超级管理员'),
):
await TestBatchPurgeService.preview(
self.db,
7,
make_current_user(2, ['system:testbatch:purge']),
)
lock_batch.assert_not_awaited()
async def test_active_batch_cannot_be_purged(self) -> None:
with (
self._preview_stack(make_batch(status='ACTIVE')),
self.assertRaisesRegex(ServiceWarning, '必须先关闭'),
):
await TestBatchPurgeService.preview(self.db, 7, make_admin())
self.db.rollback.assert_awaited_once_with()
self.db.commit.assert_not_awaited()
async def test_preview_persists_purge_token_scope_and_all_impacts(self) -> None:
batch = make_batch()
before = datetime.now()
with self._preview_stack(batch, config=purge_config(100, ttl=600)):
result = await TestBatchPurgeService.preview(self.db, 7, make_admin())
after = datetime.now()
self.assertEqual(4, result.total_count)
self.assertEqual(list(BATCH_MODELS), [impact.table_name for impact in result.impacts])
self.assertEqual(
list(BATCH_LABELS.items()),
[(impact.table_name, impact.label) for impact in result.impacts],
)
self.assertEqual(
[('SITE-3',), ('W-2',), ('P-10',), ('U-50',)],
[tuple(impact.keys) for impact in result.impacts if impact.count],
)
self.assertTrue(result.online_purge_allowed)
self.assertEqual(100, result.online_purge_limit)
self.assertIsNone(result.blocking_reason)
self.assertEqual('PURGE', batch.preview_action)
self.assertNotEqual(result.preview_token, batch.preview_token_hash)
self.assertEqual(hashlib.sha256(result.preview_token.encode()).hexdigest(), batch.preview_token_hash)
self.assertEqual(selection_hash(self.selection), batch.preview_scope_hash)
self.assertGreaterEqual(result.expires_at, before + timedelta(seconds=599))
self.assertLessEqual(result.expires_at, after + timedelta(seconds=601))
self.db.commit.assert_awaited_once_with()
async def test_preview_remains_available_when_online_limit_is_unconfigured(self) -> None:
batch = make_batch()
with self._preview_stack(batch, config=purge_config(0)):
result = await TestBatchPurgeService.preview(self.db, 7, make_admin())
self.assertFalse(result.online_purge_allowed)
self.assertEqual(0, result.online_purge_limit)
self.assertEqual('尚未配置在线测试批次清理上限', result.blocking_reason)
self.assertIsNotNone(result.preview_token)
self.db.commit.assert_awaited_once_with()
async def test_preview_reports_rows_over_online_limit_without_blocking_inspection(self) -> None:
batch = make_batch()
with self._preview_stack(batch, config=purge_config(3)):
result = await TestBatchPurgeService.preview(self.db, 7, make_admin())
self.assertFalse(result.online_purge_allowed)
self.assertEqual(3, result.online_purge_limit)
self.assertEqual('批次数据量超过在线清理上限,请按运维流程处理', result.blocking_reason)
self.db.commit.assert_awaited_once_with()
async def test_purge_rejects_batch_name_mismatch_before_data_reads(self) -> None:
raw_token = 'n' * 43
batch = make_token_batch(raw_token, self.selection)
command = purge_command(raw_token)
command.batch_name = '错误批次名'
load_selection = AsyncMock()
with (
patch.object(TestBatchDao, 'get_batch_for_update', AsyncMock(return_value=batch)),
patch.object(TestBatchPurgeDao, 'load_selection', load_selection),
self.assertRaisesRegex(ServiceWarning, '测试批次名称不匹配'),
):
await TestBatchPurgeService.purge(self.db, 7, command, make_admin())
load_selection.assert_not_awaited()
self.db.rollback.assert_awaited_once_with()
async def test_purge_rejects_expired_token_before_data_reads(self) -> None:
raw_token = 'e' * 43
batch = make_token_batch(
raw_token,
self.selection,
preview_expires_at=datetime.now() - timedelta(seconds=1),
)
load_selection = AsyncMock()
with (
patch.object(TestBatchDao, 'get_batch_for_update', AsyncMock(return_value=batch)),
patch.object(TestBatchPurgeDao, 'load_selection', load_selection),
self.assertRaisesRegex(ServiceWarning, '预览令牌已过期'),
):
await TestBatchPurgeService.purge(self.db, 7, purge_command(raw_token), make_admin())
load_selection.assert_not_awaited()
self.db.rollback.assert_awaited_once_with()
async def test_purge_recomputes_locked_scope_and_rejects_preview_changes(self) -> None:
raw_token = 's' * 43
batch = make_token_batch(raw_token, self.selection)
changed = dict(self.selection)
changed['biz_pallet_event'] = frozenset({61})
delete_rows = AsyncMock()
with (
self._purge_stack(batch, selection=changed, delete_rows=delete_rows),
self.assertRaisesRegex(ServiceWarning, '预览范围已变化'),
):
await TestBatchPurgeService.purge(self.db, 7, purge_command(raw_token), make_admin())
delete_rows.assert_not_awaited()
self.db.rollback.assert_awaited_once_with()
async def test_purge_blocks_after_locked_recount_when_limit_is_unconfigured(self) -> None:
raw_token = 'z' * 43
batch = make_token_batch(raw_token, self.selection)
delete_rows = AsyncMock()
with (
self._purge_stack(batch, delete_rows=delete_rows, limit=0),
self.assertRaisesRegex(ServiceWarning, '尚未配置在线测试批次清理上限'),
):
await TestBatchPurgeService.purge(self.db, 7, purge_command(raw_token), make_admin())
delete_rows.assert_not_awaited()
self.db.rollback.assert_awaited_once_with()
async def test_purge_blocks_after_locked_recount_when_rows_exceed_limit(self) -> None:
raw_token = 'l' * 43
batch = make_token_batch(raw_token, self.selection)
delete_rows = AsyncMock()
with (
self._purge_stack(batch, delete_rows=delete_rows, limit=3),
self.assertRaisesRegex(ServiceWarning, '批次数据量超过在线清理上限'),
):
await TestBatchPurgeService.purge(self.db, 7, purge_command(raw_token), make_admin())
delete_rows.assert_not_awaited()
self.db.rollback.assert_awaited_once_with()
async def test_delete_failure_rolls_back_and_keeps_closed(self) -> None:
raw_token = 'f' * 43
selection = make_selection(
biz_pallet_inventory={1, 2, 3},
biz_usage_record={4},
)
batch = make_token_batch(raw_token, selection)
delete_rows = AsyncMock(side_effect=[3, RuntimeError('fk failure')])
with (
self._purge_stack(batch, selection=selection, delete_rows=delete_rows),
self.assertRaisesRegex(RuntimeError, 'fk failure'),
):
await TestBatchPurgeService.purge(self.db, 7, purge_command(raw_token), make_admin())
self.db.rollback.assert_awaited_once_with()
self.db.commit.assert_not_awaited()
self.assertEqual('CLOSED', batch.status)
async def test_delete_count_change_rolls_back_before_marking_purged(self) -> None:
raw_token = 'c' * 43
selection = make_selection(biz_pallet_inventory={1, 2, 3})
batch = make_token_batch(raw_token, selection)
delete_rows = AsyncMock(return_value=2)
with (
self._purge_stack(batch, selection=selection, delete_rows=delete_rows),
self.assertRaises(ServiceException) as caught,
):
await TestBatchPurgeService.purge(self.db, 7, purge_command(raw_token), make_admin())
self.assertEqual('biz_pallet_inventory 清理数量发生变化', caught.exception.message)
self.assertEqual('CLOSED', batch.status)
self.db.rollback.assert_awaited_once_with()
async def test_purge_deletes_links_then_foreign_key_order_and_saves_summary(self) -> None:
raw_token = 'p' * 43
batch = make_token_batch(raw_token, self.selection)
actions: list[str] = []
async def delete_warehouse_links(_db: AsyncSession, _ids: frozenset[int]) -> None:
actions.append('biz_user_warehouse')
async def delete_site_links(_db: AsyncSession, _ids: frozenset[int]) -> None:
actions.append('biz_user_site')
async def delete_rows(_db: AsyncSession, table_name: str, _batch_id: int) -> int:
actions.append(table_name)
return len(self.selection[table_name])
async def delete_batch_users(_db: AsyncSession, _batch_id: int) -> None:
actions.append('biz_test_batch_user')
self.db.commit.side_effect = lambda: actions.append('commit')
with (
patch.object(TestBatchDao, 'get_batch_for_update', AsyncMock(return_value=batch)),
patch.object(TestBatchPurgeDao, 'load_selection', AsyncMock(return_value=self.selection)),
patch.object(TestBatchPurgeDao, 'delete_user_warehouse_links', delete_warehouse_links),
patch.object(TestBatchPurgeDao, 'delete_user_site_links', delete_site_links),
patch.object(TestBatchPurgeDao, 'delete_table_rows', delete_rows),
patch.object(TestBatchDao, 'delete_batch_users', delete_batch_users),
patch(
'module_business.service.test_batch_purge_service.TestBatchConfig',
purge_config(10_000),
),
):
result = await TestBatchPurgeService.purge(
self.db,
7,
purge_command(raw_token),
make_admin(),
)
self.assertEqual(CrudResponseModel(is_success=True, message='测试批次清理成功'), result)
self.assertEqual(
['biz_user_warehouse', 'biz_user_site', *EXPECTED_DELETE_ORDER, 'biz_test_batch_user', 'commit'],
actions,
)
self.assertEqual('PURGED', batch.status)
self.assertEqual(1, batch.purged_by)
self.assertIsInstance(batch.purged_at, datetime)
self.assertEqual(
{table_name: len(self.selection[table_name]) for table_name in BATCH_MODELS},
json.loads(batch.purge_summary),
)
self.assertIsNone(batch.preview_action)
self.assertIsNone(batch.preview_token_hash)
self.assertIsNone(batch.preview_scope_hash)
self.assertIsNone(batch.preview_expires_at)
self.db.rollback.assert_not_awaited()
async def test_repeated_purge_of_purged_batch_is_stable(self) -> None:
batch = make_batch(status='PURGED')
delete_rows = AsyncMock()
with (
patch.object(TestBatchDao, 'get_batch_for_update', AsyncMock(return_value=batch)),
patch.object(TestBatchPurgeDao, 'delete_table_rows', delete_rows),
):
result = await TestBatchPurgeService.purge(
self.db,
7,
purge_command('r' * 43),
make_admin(),
)
self.assertEqual('该测试批次已清理', result.message)
delete_rows.assert_not_awaited()
self.db.commit.assert_awaited_once_with()
self.db.rollback.assert_not_awaited()
async def test_two_concurrent_purge_requests_delete_only_once(self) -> None:
raw_token = 'q' * 43
selection = make_selection()
batch = make_token_batch(raw_token, selection)
transaction_lock = asyncio.Lock()
first_db = AsyncMock(spec=AsyncSession)
second_db = AsyncMock(spec=AsyncSession)
async def lock_batch(_db: AsyncSession, _batch_id: int) -> SimpleNamespace:
await transaction_lock.acquire()
return batch
def release_transaction() -> None:
transaction_lock.release()
first_db.commit.side_effect = release_transaction
second_db.commit.side_effect = release_transaction
delete_rows = AsyncMock(return_value=0)
with (
patch.object(TestBatchDao, 'get_batch_for_update', side_effect=lock_batch),
patch.object(TestBatchPurgeDao, 'load_selection', AsyncMock(return_value=selection)),
patch.object(TestBatchPurgeDao, 'delete_user_warehouse_links', AsyncMock()),
patch.object(TestBatchPurgeDao, 'delete_user_site_links', AsyncMock()),
patch.object(TestBatchPurgeDao, 'delete_table_rows', delete_rows),
patch.object(TestBatchDao, 'delete_batch_users', AsyncMock()),
patch(
'module_business.service.test_batch_purge_service.TestBatchConfig',
purge_config(10_000),
),
):
results = await asyncio.wait_for(
asyncio.gather(
TestBatchPurgeService.purge(first_db, 7, purge_command(raw_token), make_admin()),
TestBatchPurgeService.purge(second_db, 7, purge_command(raw_token), make_admin()),
),
timeout=2,
)
self.assertEqual(
{'测试批次清理成功', '该测试批次已清理'},
{result.message for result in results},
)
self.assertEqual(len(BATCH_MODELS), delete_rows.await_count)
if __name__ == '__main__':
unittest.main()