bot-backend
3320805723
Extend the cooperative shutdown-drain (Phase 2) to the long sweep and
full-load loops in scrape_pipeline.py so they stop CLEANLY on SIGTERM
instead of being hard-cancelled at the 80s _drain_inflight timeout.
At every between-unit checkpoint that already polls scrape_runs.is_cancelled
(user-cancel) we add an ALONGSIDE shutdown_requested() branch:
- anchor-loop sweeps (avito/yandex/cian city) + avito newbuilding +
domclick pre-SERP: update_heartbeat + mark_done(partial) + return;
- full-load on_bucket callbacks (cian/yandex/avito): raise a distinct
RuntimeError("shutdown") sentinel, handled next to the existing
"cancelled" sentinel -> mark_done(partial), skipping the detail phase;
- cian full-load detail loop: break -> existing post-loop mark_done.
Drain finalizes as mark_done(partial), NOT mark_cancelled (that stays
user-cancel only). is_cancelled -> user-cancel semantics are unchanged, and
behaviour is identical when shutdown_requested() is False. No resume_cursor
or paused status yet (Phase 3b, deferred): this is drain-coverage only.
Tests: tests/test_sweep_drain.py covers both structural variants -
run_avito_city_sweep / run_cian_city_sweep (anchor-loop) and
run_avito_full_load (on_bucket sentinel): shutdown flips True after the
first unit, asserting the sweep stops at the checkpoint, calls mark_done
with partial counters, never calls mark_cancelled, and does not process
later anchors/buckets/phases.
310 lines
13 KiB
Python
310 lines
13 KiB
Python
"""Тесты кооперативного 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
|