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

"""#3404: атрибуция прогона к узлу (scrape_runs.proxy_id) + гашение бана вместо DELETE.

Offline-тесты (без live БД) в стиле test_proxy_pool.py / test_3390_single_runs_module.py:
stateful fake-сессия эмулирует таблицы scrape_proxies / scrape_proxy_source_bans /
scrape_runs, интерпретируя SQL по ключевым фрагментам, так что acquire/mark_banned/
clear_source_bans/attribute_run_proxy проверяются по фактическому изменению состояния,
а не по замоканному возврату.

Покрытие:
  1. clear_source_bans гасит строку (banned_until<=now, ban_count=0, cleared_at
     заполнен), а НЕ удаляет — строка остаётся в таблице.
  2. Эскалация 1:1: бан → clear → бан той же пары снова стартует с ровно
     SOURCE_BAN_BASE_HOURS, а не удвоенного срока (главный тест issue).
  3. Погашенная строка не мешает acquire(source) — узел выдаётся как обычно.
  4. attribute_run_proxy доводит proxy_id прогона до scrape_runs через acquire();
     путь без run_id (health-checker, run_id=None) оставляет scrape_runs нетронутым.
  5. Смена узла за прогон (повторный acquire тем же run_id на другой узел)
     отражается в counters['proxy_ids'] без дублей.
"""

from __future__ import annotations

import os

os.environ.setdefault("DATABASE_URL", "postgresql+psycopg://test:test@localhost:5432/test")

from datetime import UTC, datetime, timedelta
from typing import Any

from app.services.proxy_pool import (
    SOURCE_BAN_BASE_HOURS,
    acquire,
    attribute_run_proxy,
    clear_source_bans,
    mark_banned,
)

# ── stateful fake session (scrape_proxies + scrape_proxy_source_bans + scrape_runs) ──


class _FakeResult:
    def __init__(self, rows: list[Any]) -> None:
        self._rows = rows

    def mappings(self) -> _FakeResult:
        return self

    def fetchone(self) -> Any:
        return self._rows[0] if self._rows else None

    def first(self) -> Any:
        return self._rows[0] if self._rows else None

    def fetchall(self) -> list[Any]:
        return [type("Row", (), r)() if isinstance(r, dict) else r for r in self._rows]

    def all(self) -> list[Any]:
        return list(self._rows)


def _proxy(
    pid: int,
    *,
    affinity: str = "avito",
    enabled: bool = True,
    leased_by: int | None = None,
) -> dict[str, Any]:
    return {
        "id": pid,
        "url": f"http://proxy{pid}",
        "kind": "residential",
        "rotate_url": None,
        "browser_unfit_since": None,
        "enabled": enabled,
        "consecutive_fails": 0,
        "provider_affinity": affinity,
        "leased_by": leased_by,
        "leased_at": None,
        "expires_at": None,
        "last_ok_at": None,
    }


class RunAttributionDb:
    """Мини-Postgres: scrape_proxies + scrape_proxy_source_bans + scrape_runs.

    Гейты по SQL-подстрокам скопированы из проверяемого кода (proxy_pool.py) — как в
    test_proxy_pool.py/test_3390: значение проверяется по факту исполнения реального
    statement'а, а не зашито ожиданием теста.
    """

    def __init__(self, proxies: list[dict[str, Any]]) -> None:
        self.proxies = proxies
        self.bans: list[dict[str, Any]] = []
        self.runs: dict[int, dict[str, Any]] = {}
        self.set_local_statements: list[str] = []

    def add_run(self, run_id: int) -> None:
        self.runs[run_id] = {"proxy_id": None, "counters": {}}

    def _by_id(self, pid: int) -> dict[str, Any] | None:
        return next((r for r in self.proxies if r["id"] == pid), None)

    def _ban(self, pid: int, source: str) -> dict[str, Any] | None:
        return next((b for b in self.bans if b["proxy_id"] == pid and b["source"] == source), None)

    def _has_active_ban(self, pid: int, source: str) -> bool:
        b = self._ban(pid, source)
        return b is not None and b["banned_until"] > datetime.now(UTC)

    def execute(self, stmt: Any, params: dict[str, Any] | None = None) -> _FakeResult:
        sql = " ".join(str(stmt).split())
        p = params or {}

        if "pg_advisory_xact_lock" in sql:
            return _FakeResult([])

        if "expires_at <= now()" in sql and "FOR UPDATE SKIP LOCKED" not in sql:
            return _FakeResult([])  # без просроченных узлов в этих тестах

        if "FOR UPDATE SKIP LOCKED" in sql:  # acquire SELECT (primary/fallback)
            provider = p["provider"]
            max_fails = p["max_fails"]
            primary = "provider_affinity IN" in sql
            cands = [
                r
                for r in self.proxies
                if r["enabled"]
                and r["consecutive_fails"] < max_fails
                and r["leased_by"] is None
                and not self._has_active_ban(r["id"], provider)
                and (r["provider_affinity"] in (provider, "any") if primary else True)
            ]
            cands.sort(
                key=lambda r: (
                    r["browser_unfit_since"] is not None,
                    r["last_ok_at"] is None,
                    r["last_ok_at"] or datetime.min.replace(tzinfo=UTC),
                    r["id"],
                )
            )
            return _FakeResult(cands[:1])

        if "SET leased_by = CAST(:run_id" in sql:  # acquire lease UPDATE
            row = self._by_id(p["id"])
            if row is not None:
                row["leased_by"] = p["run_id"]
                row["leased_at"] = datetime.now(UTC)
            return _FakeResult([])

        if sql.startswith("SET LOCAL"):
            # attribute_run_proxy ставит lock_timeout перед UPDATE scrape_runs:
            # ждать на блокировке строки прогона в пути выдачи прокси нельзя.
            # Для фейка это no-op, но проглатывать молча нечестно — гейт ниже
            # ловит любой ДРУГОЙ незнакомый SQL.
            self.set_local_statements.append(sql)
            return _FakeResult([])

        if "SET leased_by = NULL" in sql:  # release
            row = self._by_id(p["id"])
            if row is not None:
                row["leased_by"] = None
                row["leased_at"] = None
            return _FakeResult([])

        if "INSERT INTO scrape_proxy_source_bans" in sql:  # mark_banned upsert
            proxy_id, source = p["proxy_id"], p["source"]
            if self._by_id(proxy_id) is None:
                return _FakeResult([])
            max_fails = p["max_fails"]
            still_available = any(
                r["id"] != proxy_id
                and r["enabled"]
                and r["consecutive_fails"] < max_fails
                and not self._has_active_ban(r["id"], source)
                for r in self.proxies
            )
            if not still_available:
                return _FakeResult([])  # protected — последний узел

            now = datetime.now(UTC)
            ban = self._ban(proxy_id, source)
            reason = p["reason"]
            if ban is None:
                ban = {
                    "proxy_id": proxy_id,
                    "source": source,
                    "ban_count": 1,
                    "banned_until": now + timedelta(hours=p["base_hours"]),
                    "reason": reason,
                    "cleared_at": None,
                    "cleared_reason": None,
                }
                self.bans.append(ban)
            else:
                is_active_foreign = (
                    ban["banned_until"] > now
                    and ban.get("reason") != reason
                    and reason != p.get("live_reason")
                )
                if is_active_foreign:
                    return _FakeResult([])  # deferred — чужой активный владелец
                ban["ban_count"] += 1
                hours = min(p["base_hours"] * 2 ** (ban["ban_count"] - 1), p["max_hours"])
                ban["banned_until"] = now + timedelta(hours=hours)
                ban["reason"] = reason
                ban["cleared_at"] = None
                ban["cleared_reason"] = None
            return _FakeResult(
                [{"ban_count": ban["ban_count"], "banned_until": ban["banned_until"]}]
            )

        if "UPDATE scrape_proxy_source_bans" in sql and "cleared_at" in sql:  # clear_source_bans
            proxy_id = p["proxy_id"]
            source = p.get("source")
            only_reason = p.get("only_reason")
            reason = p["reason"]
            now = datetime.now(UTC)
            cleared = []
            for b in self.bans:
                if b["proxy_id"] != proxy_id:
                    continue
                if source is not None and b["source"] != source:
                    continue
                if only_reason is not None and b.get("reason") != only_reason:
                    continue
                if (
                    b.get("cleared_at") is not None
                    and b["ban_count"] == 0
                    and b["banned_until"] <= now
                ):
                    continue  # гейт покоя — уже погашена, no-op
                b["banned_until"] = now
                b["ban_count"] = 0
                b["cleared_at"] = now
                b["cleared_reason"] = reason
                cleared.append({"source": b["source"]})
            return _FakeResult(cleared)

        if "UPDATE scrape_runs" in sql:  # attribute_run_proxy (#3404)
            run = self.runs.get(p["run_id"])
            if run is None:
                return _FakeResult([])
            run["proxy_id"] = p["proxy_id"]
            proxy_ids: list[int] = run["counters"].get("proxy_ids", [])
            if p["proxy_id"] not in proxy_ids:
                proxy_ids = [*proxy_ids, p["proxy_id"]]
            run["counters"]["proxy_ids"] = proxy_ids
            return _FakeResult([{"id": p["run_id"]}])

        raise AssertionError(f"unhandled SQL in fake session: {sql[:120]}")

    def commit(self) -> None:
        pass

    def rollback(self) -> None:
        pass


# ── 1. clear_source_bans гасит, а не удаляет ────────────────────────────────


def test_clear_source_bans_gates_not_deletes() -> None:
    db = RunAttributionDb([_proxy(1), _proxy(2)])
    assert mark_banned(db, 1, source="avito", reason="banned:avito") == "banned"  # type: ignore[arg-type]
    assert len(db.bans) == 1

    cleared = clear_source_bans(db, 1, reason="manual enable")  # type: ignore[arg-type]

    assert cleared == 1
    assert len(db.bans) == 1, "строка должна остаться на месте, не удалиться"
    ban = db.bans[0]
    assert ban["banned_until"] <= datetime.now(UTC)
    assert ban["ban_count"] == 0
    assert ban["cleared_at"] is not None
    assert ban["cleared_reason"] == "manual enable"


# ── 2. Эскалация сохранилась 1:1 (главный тест issue) ───────────────────────


def test_escalation_resets_to_base_after_clear() -> None:
    db = RunAttributionDb([_proxy(1), _proxy(2)])

    assert mark_banned(db, 1, source="avito", reason="banned:avito") == "banned"  # type: ignore[arg-type]
    first_ban = db.bans[0]
    assert first_ban["ban_count"] == 1
    first_span = first_ban["banned_until"] - datetime.now(UTC)
    assert timedelta(hours=SOURCE_BAN_BASE_HOURS - 1) < first_span <= timedelta(
        hours=SOURCE_BAN_BASE_HOURS
    )

    clear_source_bans(db, 1, reason="manual enable")  # type: ignore[arg-type]
    assert db.bans[0]["ban_count"] == 0

    assert mark_banned(db, 1, source="avito", reason="banned:avito") == "banned"  # type: ignore[arg-type]
    second_ban = db.bans[0]

    assert second_ban["ban_count"] == 1, "после clear эскалация обязана стартовать заново"
    second_span = second_ban["banned_until"] - datetime.now(UTC)
    assert timedelta(hours=SOURCE_BAN_BASE_HOURS - 1) < second_span <= timedelta(
        hours=SOURCE_BAN_BASE_HOURS
    ), f"срок {second_span} обязан быть БАЗОВЫМ ({SOURCE_BAN_BASE_HOURS}ч), а не удвоенным"


# ── 3. Погашенная строка не мешает выдаче узла ──────────────────────────────


def test_acquire_ignores_cleared_ban() -> None:
    db = RunAttributionDb([_proxy(1), _proxy(2)])
    mark_banned(db, 1, source="avito", reason="banned:avito")  # type: ignore[arg-type]
    clear_source_bans(db, 1, reason="manual enable")  # type: ignore[arg-type]
    # узел 2 занят чужим прогоном — иначе acquire мог бы честно выдать его вместо
    # узла 1, и тест перестал бы проверять именно "погашенный бан не блокирует".
    db.proxies[1]["leased_by"] = 999

    lease = acquire(db, "avito")  # type: ignore[arg-type]

    assert lease is not None, "погашенный бан не должен блокировать acquire"
    assert lease.id == 1


def test_acquire_still_blocks_active_ban_after_clear_gate_check() -> None:
    """Контроль ложноположительного теста выше: АКТИВНЫЙ (не погашенный) бан acquire
    по-прежнему блокирует — это не сломано этим же изменением."""
    db = RunAttributionDb([_proxy(1), _proxy(2)])
    mark_banned(db, 1, source="avito", reason="banned:avito")  # type: ignore[arg-type]

    lease = acquire(db, "avito")  # type: ignore[arg-type]

    assert lease is not None
    assert lease.id == 2, "узел 1 активно забанен по avito — выдан должен быть узел 2"


# ── 4. proxy_id прогона доезжает до scrape_runs ─────────────────────────────


def test_acquire_attributes_run_proxy() -> None:
    db = RunAttributionDb([_proxy(1)])
    db.add_run(42)

    lease = acquire(db, "avito", run_id=42)  # type: ignore[arg-type]

    assert lease is not None
    assert db.runs[42]["proxy_id"] == lease.id == 1
    assert db.runs[42]["counters"]["proxy_ids"] == [1]
    # Атрибуция обязана ограничить ожидание блокировки: строку прогона параллельно
    # пишет heartbeat из другой сессии, а путь выдачи прокси ждать не может.
    assert any("lock_timeout" in stmt for stmt in db.set_local_statements)


def test_acquire_without_run_id_leaves_scrape_runs_untouched() -> None:
    """health-checker и прочие не-run вызовы (run_id=None) не трогают scrape_runs."""
    db = RunAttributionDb([_proxy(1)])
    db.add_run(42)

    lease = acquire(db, "avito")  # type: ignore[arg-type]

    assert lease is not None
    assert db.runs[42]["proxy_id"] is None
    assert db.runs[42]["counters"] == {}


# ── 5. Смена узла за прогон отражается в counters ───────────────────────────


def test_mid_run_proxy_rotation_recorded_in_counters() -> None:
    """Ре-acquire тем же run_id на другой узел (старый lease ещё держится, как при
    ротации в browser_fetcher — `_LEASE_ROTATE_AFTER_FAILS`) — оба узла в proxy_ids,
    proxy_id несёт ПОСЛЕДНИЙ (сверено с attribute_run_proxy docstring, а не с ТЗ)."""
    db = RunAttributionDb([_proxy(1), _proxy(2)])
    db.add_run(7)

    first = acquire(db, "avito", run_id=7)  # type: ignore[arg-type]
    assert first is not None
    # старый lease НЕ освобождён (leased_by=7) — второй acquire тем же run_id
    # неизбежно возьмёт другой свободный узел, детерминированно.
    assert db._by_id(first.id)["leased_by"] == 7

    second = acquire(db, "avito", run_id=7)  # type: ignore[arg-type]
    assert second is not None
    assert second.id != first.id, "тестовая обвязка обязана взять ДРУГОЙ узел"

    assert db.runs[7]["proxy_id"] == second.id
    assert set(db.runs[7]["counters"]["proxy_ids"]) == {first.id, second.id}


def test_attribute_run_proxy_idempotent_same_proxy() -> None:
    """Повторная атрибуция ТЕМ ЖЕ узлом не дублирует id в proxy_ids."""
    db = RunAttributionDb([_proxy(1)])
    db.add_run(9)

    attribute_run_proxy(db, 9, 1)  # type: ignore[arg-type]
    attribute_run_proxy(db, 9, 1)  # type: ignore[arg-type]

    assert db.runs[9]["counters"]["proxy_ids"] == [1]


def test_attribute_run_proxy_best_effort_on_missing_run() -> None:
    """run_id не найден (гонка с финализацией) — best-effort no-op, исключение не летит."""
    db = RunAttributionDb([_proxy(1)])
    # run 999 никогда не создавался в этой fake-БД
    attribute_run_proxy(db, 999, 1)  # type: ignore[arg-type]  # не должно бросить