"""Тесты кооперативного SIGTERM-drain в длинных sweep/full-load циклах (#1182 Phase 3a). Проверяют, что при выставленном shutdown_requested() sweep'ы scrape_pipeline.py останавливаются ЧИСТО на своих between-unit checkpoint'ах: - финализируют run как mark_done(partial) — НЕ mark_cancelled (это не user-cancel); - НЕ обрабатывают последующие anchor'ы / бакеты / дорогие фазы (IMV / detail); - user-cancel ветка (is_cancelled) не задействуется. Тесты network-free: внутренние scrape-вызовы и lifecycle скрейперов monkeypatch'атся, shutdown_requested() патчится контролируемым флагом, переключающимся в True ПОСЛЕ первого unit'а (а не с самого старта) — чтобы первый unit прошёл, а на втором сработал drain. """ from __future__ import annotations import os os.environ.setdefault("DATABASE_URL", "postgresql+psycopg://test:test@localhost:5432/test") from typing import Any from unittest.mock import AsyncMock, MagicMock import pytest from app.services import scrape_pipeline from app.services.scrapers.avito import AvitoScraper from app.services.scrapers.base import ScrapedLot from app.services.scrapers.cian import CianScraper # Два anchor'а: первый успешен, второй должен быть отрезан drain'ом. TWO_ANCHORS: list[tuple[float, float, str]] = [ (56.8400, 60.6050, "First"), (56.7950, 60.5300, "Second"), ] # ── Fake scrape_runs (записывает финализаторы) ────────────────────────────── class _FakeRuns: def __init__(self) -> None: self.heartbeats: list[dict[str, int]] = [] self.done: dict[str, int] | None = None self.failed: tuple[str, dict[str, int]] | None = None self.banned: tuple[str, dict[str, int]] | None = None self.cancelled_calls: int = 0 self.is_cancelled_calls: int = 0 def is_cancelled(self, _db: Any, _run_id: int) -> bool: self.is_cancelled_calls += 1 return False def update_heartbeat(self, _db: Any, _run_id: int, counters: dict[str, int]) -> None: self.heartbeats.append(dict(counters)) def mark_done(self, _db: Any, _run_id: int, counters: dict[str, int]) -> None: self.done = dict(counters) def mark_failed(self, _db: Any, _run_id: int, error: str, counters: dict[str, int]) -> None: self.failed = (error, dict(counters)) def mark_banned(self, _db: Any, _run_id: int, error: str, counters: dict[str, int]) -> None: self.banned = (error, dict(counters)) def mark_cancelled(self, _db: Any, _run_id: int) -> bool: self.cancelled_calls += 1 return True def _install_runs(monkeypatch: pytest.MonkeyPatch, fake: _FakeRuns) -> None: monkeypatch.setattr(scrape_pipeline.scrape_runs, "is_cancelled", fake.is_cancelled) monkeypatch.setattr(scrape_pipeline.scrape_runs, "update_heartbeat", fake.update_heartbeat) monkeypatch.setattr(scrape_pipeline.scrape_runs, "mark_done", fake.mark_done) monkeypatch.setattr(scrape_pipeline.scrape_runs, "mark_failed", fake.mark_failed) monkeypatch.setattr(scrape_pipeline.scrape_runs, "mark_banned", fake.mark_banned) monkeypatch.setattr(scrape_pipeline.scrape_runs, "mark_cancelled", fake.mark_cancelled) def _fake_lot(source: str, offer_id: str) -> ScrapedLot: return ScrapedLot( source=source, source_url=f"https://example.ru/{offer_id}/", source_id=offer_id, address="Екатеринбург, улица Тестовая, 1", price_rub=4_000_000, area_m2=42.0, rooms=1, ) # ── Avito city sweep: drain на границе anchor'а ───────────────────────────── async def test_avito_city_sweep_drains_on_shutdown(monkeypatch: pytest.MonkeyPatch) -> None: """SIGTERM после первого anchor'а → sweep останавливается на втором anchor'е: - второй anchor НЕ фетчится (drain до начала его фаз); - run финализируется mark_done(partial) с anchors_done==1; - mark_cancelled НЕ вызывается (это drain, не user-cancel); - mark_failed / mark_banned не вызываются. """ import curl_cffi.requests as _cffi_mod state = {"flag": False} monkeypatch.setattr(scrape_pipeline, "shutdown_requested", lambda: state["flag"]) monkeypatch.setattr(scrape_pipeline.asyncio, "sleep", AsyncMock()) fetch_calls: list[str] = [] async def fake_fetch_around( self: AvitoScraper, lat: float, lon: float, radius_m: int, pages: int = 1, delay_override_sec: float | None = None, ) -> list[ScrapedLot]: fetch_calls.append(f"{lat:.4f}") lots = [_fake_lot("avito", f"{lat:.4f}-{i}") for i in range(3)] # SIGTERM приходит ПОСЛЕ первого anchor'а — следующий top-of-loop checkpoint # увидит флаг и сделает drain. state["flag"] = True return lots monkeypatch.setattr(AvitoScraper, "fetch_around", fake_fetch_around) monkeypatch.setattr(scrape_pipeline, "save_listings", lambda _db, lots, **_kw: (len(lots), 0)) # Stub shared curl_cffi AsyncSession (non-browser fetch mode). fake_session = AsyncMock() fake_session.__aenter__ = AsyncMock(return_value=fake_session) fake_session.__aexit__ = AsyncMock(return_value=None) monkeypatch.setattr(_cffi_mod, "AsyncSession", lambda **_kw: fake_session) # IMV-фаза не должна вызываться при drain (enrich_imv=False всё равно). monkeypatch.setattr( "app.services.house_imv_backfill.process_houses_imv_batch", AsyncMock(side_effect=AssertionError("IMV phase must not run after drain")), ) fake = _FakeRuns() _install_runs(monkeypatch, fake) counters = await scrape_pipeline.run_avito_city_sweep( db=MagicMock(), run_id=1, anchors=TWO_ANCHORS, enrich_houses=False, detail_top_n=0, enrich_imv=False, request_delay_sec=0.0, ) # Второй anchor не дошёл до SERP — fetch вызван ровно один раз. assert fetch_calls == ["56.8400"], f"только первый anchor должен фетчиться; got {fetch_calls}" # Partial-финализация: обработан только первый anchor. assert counters.anchors_done == 1 assert counters.lots_fetched == 3 # mark_done(partial) вызван (drain-ветка), и счётчики partial. assert fake.done is not None, "drain должен вызвать mark_done(partial)" assert fake.done["anchors_done"] == 1 # Это НЕ user-cancel. assert fake.cancelled_calls == 0, "mark_cancelled не должен вызываться при drain" assert fake.failed is None assert fake.banned is None # ── Cian city sweep: drain на границе anchor'а ────────────────────────────── @pytest.fixture def _stub_cian_lifecycle(monkeypatch: pytest.MonkeyPatch) -> None: def _init(self: CianScraper) -> None: self.request_delay_sec = 0.0 self._cffi = None self._browser = None async def _aenter(self: CianScraper) -> CianScraper: return self async def _aexit(self: CianScraper, *_args: Any) -> None: return None monkeypatch.setattr(CianScraper, "__init__", _init) monkeypatch.setattr(CianScraper, "__aenter__", _aenter) monkeypatch.setattr(CianScraper, "__aexit__", _aexit) async def test_cian_city_sweep_drains_on_shutdown( monkeypatch: pytest.MonkeyPatch, _stub_cian_lifecycle: None ) -> None: """SIGTERM после первого anchor'а → cian sweep drain на втором anchor'е → mark_done(partial).""" state = {"flag": False} monkeypatch.setattr(scrape_pipeline, "shutdown_requested", lambda: state["flag"]) monkeypatch.setattr(scrape_pipeline.asyncio, "sleep", AsyncMock()) serp_calls: list[str] = [] async def fake_fetch_multi( self: CianScraper, lat: float, lon: float, r: int, pages: int = 3 ) -> list[ScrapedLot]: serp_calls.append(f"{lat:.4f}") state["flag"] = True return [_fake_lot("cian", f"{lat:.4f}-1")] monkeypatch.setattr(CianScraper, "fetch_around_multi_room", fake_fetch_multi) monkeypatch.setattr( scrape_pipeline, "save_listings", lambda _db, lots, *, run_id=None: (len(lots), 0) ) monkeypatch.setattr( "app.services.scrapers.cian_detail.fetch_detail", AsyncMock(return_value=None) ) fake = _FakeRuns() _install_runs(monkeypatch, fake) counters = await scrape_pipeline.run_cian_city_sweep( db=MagicMock(), run_id=2, anchors=TWO_ANCHORS, detail_top_n=0, enrich_houses=False, request_delay_sec=0.0, ) assert serp_calls == ["56.8400"], f"только первый anchor должен фетчиться; got {serp_calls}" assert counters.anchors_done == 1 assert fake.done is not None, "drain должен вызвать mark_done(partial)" assert fake.done["anchors_done"] == 1 assert fake.cancelled_calls == 0 assert fake.failed is None assert fake.banned is None # ── Avito full load: drain через sentinel в on_bucket ─────────────────────── @pytest.fixture def _stub_avito_lifecycle(monkeypatch: pytest.MonkeyPatch) -> None: def _init(self: AvitoScraper) -> None: self.request_delay_sec = 0.0 self._cffi = None self._browser = None async def _aenter(self: AvitoScraper) -> AvitoScraper: return self async def _aexit(self: AvitoScraper, *_args: Any) -> None: return None monkeypatch.setattr(AvitoScraper, "__init__", _init) monkeypatch.setattr(AvitoScraper, "__aenter__", _aenter) monkeypatch.setattr(AvitoScraper, "__aexit__", _aexit) async def test_avito_full_load_drains_on_shutdown( monkeypatch: pytest.MonkeyPatch, _stub_avito_lifecycle: None ) -> None: """SIGTERM после первого бакета → on_bucket поднимает sentinel → mark_done(partial): - бакет №2 поднимает RuntimeError("shutdown") (НЕ "cancelled"), бакет №3 не достигается; - done_buckets == ["b1"] (бакет, упавший на drain-checkpoint, в done не попадает); - run финализируется mark_done(partial) с unique_fetched==1; - mark_cancelled / mark_failed / mark_banned не вызываются. """ state = {"flag": False} monkeypatch.setattr(scrape_pipeline, "shutdown_requested", lambda: state["flag"]) saved_buckets: list[str] = [] def fake_save(_db: Any, lots: Any, **_kw: Any) -> tuple[int, int]: saved_buckets.append("x") return (len(lots), 0) monkeypatch.setattr(scrape_pipeline, "save_listings", fake_save) async def fake_fetch_all(self: AvitoScraper, **kwargs: Any) -> None: on_bucket = kwargs["on_bucket"] # Бакет №1 — сохраняется штатно. on_bucket("b1", [_fake_lot("avito", "b1-1")]) # SIGTERM приходит после первого бакета. state["flag"] = True # Бакет №2 — on_bucket должен поднять RuntimeError("shutdown"). on_bucket("b2", [_fake_lot("avito", "b2-1")]) # Бакет №3 — НЕ должен быть достигнут. on_bucket("b3", [_fake_lot("avito", "b3-1")]) monkeypatch.setattr(AvitoScraper, "fetch_all_secondary", fake_fetch_all) fake = _FakeRuns() _install_runs(monkeypatch, fake) counters = await scrape_pipeline.run_avito_full_load( db=MagicMock(), run_id=3, request_delay_sec=0.0, ) # Только первый бакет сохранён (b2 упал на checkpoint до save, b3 не достигнут). assert saved_buckets == ["x"], f"только b1 должен сохраниться; got {saved_buckets}" assert counters.unique_fetched == 1 assert counters.saved_inserted == 1 # mark_done(partial) с done_buckets только b1. assert fake.done is not None, "drain должен вызвать mark_done(partial)" assert fake.done.get("done_buckets") == ["b1"] assert fake.done["unique_fetched"] == 1 # НЕ user-cancel, не failed/banned. assert fake.cancelled_calls == 0 assert fake.failed is None assert fake.banned is None