gendesign/tradein-mvp/backend/tests/test_3170_yandex_combo_checkpoint.py

"""#3170: чекпоинт combo в yandex_city_sweep обязан ставиться ПОСЛЕ save, не до.

На origin/main `_on_combo` (run_yandex_city_sweep) делал `done_combos.add(combo_label)`
ДО вызова save_listings. Отказ save_listings перехватывался (`except Exception`),
логировался, шёл db.rollback() — и прогон продолжался, но combo уже был отмечен
пройденным и уходил в heartbeat done_buckets. Следующий resume брал skip_combos
из done_buckets (членство в множестве, само по себе корректно) и пропускал этот
combo НАВСЕГДА — молча, прогон завершался штатно, просто сегмент выдачи (сегмент
× комнатность × ценовой диапазон) не собирался никогда.

Тот же инвариант уже был закреплён для страницы в run_avito_newbuilding_sweep
(_saved_ok, см. test_3074_avito_newbuilding_checkpoint.py), якоря в
run_avito_city_sweep (_anchor_ok) и бакета в run_cian_full_load (_mark_bucket) —
здесь применяем его к combo в yandex-свипе.

ИНВАРИАНТЫ, РАДИ КОТОРЫХ ТЕСТ:
1. Combo, чей save_listings упал, НЕ попадает в done_buckets (чекпоинт).
2. Combo, пройденный успешно (включая пустую выдачу — #3074), попадает.
3. Heartbeat пишется в любом случае (и при отказе save тоже) — иначе
   reap_zombies посчитает живой прогон мёртвым.
"""

from __future__ import annotations

import os

# Settings собирается автофикстурой conftest'а и требует database_url. Выставляем
# до остальных импортов — так же, как в test_3074_yandex_sweep_checkpoint.py.
os.environ.setdefault("DATABASE_URL", "postgresql+psycopg://test:test@localhost:5432/test")

import json
import types
from typing import Any
from unittest.mock import MagicMock, patch

import pytest


class _SweepFakeDb:
    """Резюм-SELECT отдаёт counters предшественника; heartbeat'ы записываются."""

    def __init__(self, prev_counters: dict[str, Any] | None = None) -> None:
        self.prev_counters = prev_counters or {}
        self.heartbeats: list[dict[str, Any]] = []

    def execute(self, _stmt: Any, params: dict[str, Any] | None = None) -> Any:
        if params and "counters" in params:
            self.heartbeats.append(json.loads(params["counters"]))
            return MagicMock()
        if params and "rid" in params:
            return MagicMock(fetchone=lambda: types.SimpleNamespace(counters=self.prev_counters))
        return MagicMock()  # is_cancelled/сторожа — best-effort, глотаем

    def commit(self) -> None: ...
    def rollback(self) -> None: ...


def _lot(tag: str) -> MagicMock:
    m = MagicMock()
    m.source_id = tag
    return m


class _FakeSweepScraper:
    """Двойник YandexRealtyScraper: прогоняет заданный список combo через on_combo."""

    combos: list[tuple[str, list[Any]]] = []  # noqa: RUF012 — тестовый сборник
    gate_fetch_attempts = 1
    gate_fetch_failures = 0

    def __init__(self, *_a: Any, **_kw: Any) -> None:
        pass

    async def __aenter__(self) -> _FakeSweepScraper:
        return self

    async def __aexit__(self, *_exc: Any) -> None:
        return None

    async def fetch_around_multi_room(self, *_a: Any, **kw: Any) -> list:
        on_combo = kw.get("on_combo")
        if on_combo is not None:
            for label, lots in _FakeSweepScraper.combos:
                on_combo(label, lots)
        return []


def _make_save(fail_labels: set[str]):
    """save_listings, падающий для лотов указанных combo-меток (тег лота == label)."""

    def _save(_db, new_lots, **_kw):
        tags = {lot.source_id for lot in new_lots}
        if tags & fail_labels:
            raise RuntimeError(f"save_listings упал: {tags & fail_labels}")
        return (len(new_lots), 0)

    return _save


async def _run(db: _SweepFakeDb, *, fail_labels: set[str]) -> _SweepFakeDb:
    from scraper_kit.orchestration import pipeline as pl

    # record_yandex_price_history складывается в counters.price_history_rows (int) —
    # дефолтный MagicMock() возврата не годится для json.dumps в heartbeat.
    enrichment = MagicMock()
    enrichment.record_yandex_price_history.return_value = 0

    with (
        patch.object(pl, "YandexRealtyScraper", _FakeSweepScraper),
        patch.object(pl.runs, "is_cancelled", lambda *_a: False),
        patch.object(pl, "save_listings", _make_save(fail_labels)),
    ):
        await pl.run_yandex_city_sweep(
            db,  # type: ignore[arg-type]
            run_id=9001,
            config=types.SimpleNamespace(scraper_proxy_url=None),
            matcher=MagicMock(),
            enrichment=enrichment,
            enrich_address=False,
            resume_run_id=None,
        )
    return db


def _last_checkpoint(db: _SweepFakeDb) -> list[str]:
    with_ckpt = [hb for hb in db.heartbeats if "done_buckets" in hb]
    assert with_ckpt, "ни один heartbeat не унёс done_buckets — чекпоинт не персистится"
    return with_ckpt[-1]["done_buckets"]


@pytest.mark.asyncio
async def test_combo_with_failed_save_does_not_enter_checkpoint() -> None:
    """Combo, чей save_listings упал, НЕ должен попасть в done_buckets.

    Иначе следующий resume возьмёт его из skip_combos и пропустит навсегда —
    молча, потому что прогон завершится штатно.
    """
    label_ok = "vtorichka/room_1:None-3000000"
    label_fail = "vtorichka/room_2:None-3000000"
    _FakeSweepScraper.combos = [
        (label_ok, [_lot(label_ok)]),
        (label_fail, [_lot(label_fail)]),
    ]

    db = await _run(_SweepFakeDb(), fail_labels={label_fail})

    ckpt = _last_checkpoint(db)
    assert label_fail not in ckpt, "combo с упавшим save попал в чекпоинт — покрытие потеряно молча"
    assert label_ok in ckpt


@pytest.mark.asyncio
async def test_combo_saved_successfully_enters_checkpoint() -> None:
    """Combo, пройденный до конца (save успешен, включая пустую выдачу),
    обязан попасть в done_buckets — иначе следующий resume перечитывает
    уже собранное вечно."""
    label_nonempty = "vtorichka/room_1:None-3000000"
    label_empty = "vtorichka/room_2:None-3000000"
    _FakeSweepScraper.combos = [
        (label_nonempty, [_lot(label_nonempty)]),
        (label_empty, []),  # #3074: пустой, но полностью пройденный combo
    ]

    db = await _run(_SweepFakeDb(), fail_labels=set())

    ckpt = _last_checkpoint(db)
    assert label_nonempty in ckpt
    assert label_empty in ckpt, "пустой, но пройденный combo не попал в чекпоинт"