Profeto/tests/test_injuries_inserted_count.py

"""回归测试: 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 ===")