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

"""Прогон на egress-пути знает свой узел (#3404, хвост PR #3405).

PR #3405 писал `scrape_runs.proxy_id` только из `proxy_pool.acquire()`. Прогоны, которые
берут прокси через `proxy_egress.resolve_proxy_url` (без аренды: yandex_detail_backfill,
yandex_address_backfill, curl-ветка avito_detail_backfill), атрибуцию не получали —
прод 17.09: у yandex_detail_backfill proxy_id NULL в 56 прогонах из 56.

Живой Postgres, настоящие сессии (атрибуция идёт своей сессией, поэтому строки
коммитятся и удаляются в finally). Без БД — skip.
"""

from __future__ import annotations

import os

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

import json
from collections.abc import Iterator
from typing import Any

import pytest
from scraper_kit.orchestration.run_context import current_run_id
from sqlalchemy import create_engine, text
from sqlalchemy.orm import Session, sessionmaker

from app.services import proxy_egress


def _live_engine() -> Any | None:
    dsn = os.environ.get("TEST_DATABASE_URL") or os.environ.get("DATABASE_URL", "")
    if not dsn or "localhost:5432/test" in dsn:
        return None
    try:
        engine = create_engine(dsn, future=True)
        with engine.connect() as conn:
            conn.execute(text("SELECT proxy_id FROM scrape_runs LIMIT 1"))
        return engine
    except Exception:
        return None


_ENGINE = _live_engine()
pytestmark = pytest.mark.skipif(_ENGINE is None, reason="no reachable Postgres test DB")


@pytest.fixture
def env(monkeypatch: pytest.MonkeyPatch) -> Iterator[dict[str, Any]]:
    """Узел пула + два прогона: `run` — текущий, `other` — строка, которую вызывающий
    держит изменённой без коммита."""
    assert _ENGINE is not None
    factory = sessionmaker(bind=_ENGINE, future=True)
    monkeypatch.setattr(proxy_egress, "_SessionLocal", factory)
    setup = factory()
    proxy_id = setup.execute(
        text(
            "INSERT INTO scrape_proxies (url, provider_affinity) "
            "VALUES ('http://t3404-' || gen_random_uuid(), 'any') RETURNING id"
        )
    ).scalar_one()
    run_ids = [
        setup.execute(
            text(
                "INSERT INTO scrape_runs (source, status) "
                "VALUES ('test_3404', 'running') RETURNING id"
            )
        ).scalar_one()
        for _ in range(2)
    ]
    setup.commit()
    token = current_run_id.set(None)
    try:
        yield {"proxy_id": int(proxy_id), "run": int(run_ids[0]), "other": int(run_ids[1])}
    finally:
        current_run_id.reset(token)
        setup.rollback()
        setup.execute(text("DELETE FROM scrape_runs WHERE id = ANY(:ids)"), {"ids": run_ids})
        setup.execute(text("DELETE FROM scrape_proxies WHERE id = :id"), {"id": proxy_id})
        setup.commit()
        setup.close()


def _run_row(run_id: int) -> Any:
    assert _ENGINE is not None
    with _ENGINE.connect() as conn:
        return conn.execute(
            text("SELECT proxy_id, counters FROM scrape_runs WHERE id = :id"), {"id": run_id}
        ).one()


def _node_id_of(url: str) -> int:
    assert _ENGINE is not None
    with _ENGINE.connect() as conn:
        return int(
            conn.execute(
                text("SELECT id FROM scrape_proxies WHERE url = :u"), {"u": url}
            ).scalar_one()
        )


def test_resolve_within_run_writes_node_to_scrape_runs(env: dict[str, Any]) -> None:
    """Главный случай: прогон идёт, egress выбран — прогон знает узел.

    На старом коде proxy_id остаётся NULL."""
    current_run_id.set(env["run"])
    caller = Session(bind=_ENGINE, future=True)
    try:
        url = proxy_egress.resolve_proxy_url(caller, "yandex")
    finally:
        caller.close()

    assert url is not None
    node = _node_id_of(url)
    row = _run_row(env["run"])
    assert row.proxy_id == node
    counters = row.counters if isinstance(row.counters, dict) else json.loads(row.counters)
    assert counters.get("proxy_ids") == [node]


def test_resolve_outside_run_writes_nothing(env: dict[str, Any]) -> None:
    """Без прогона (админка, проверка кук) — scrape_runs не трогаем."""
    caller = Session(bind=_ENGINE, future=True)
    try:
        assert proxy_egress.resolve_proxy_url(caller, "yandex") is not None
    finally:
        caller.close()

    assert _run_row(env["run"]).proxy_id is None


def test_attribution_does_not_commit_callers_transaction(env: dict[str, Any]) -> None:
    """db вызывающего — долгоживущая сессия прогона посреди работы. Атрибуция не имеет
    права закоммитить его незавершённые изменения (или откатить их на сбое)."""
    current_run_id.set(env["run"])
    caller = Session(bind=_ENGINE, future=True)
    try:
        caller.execute(
            text("UPDATE scrape_runs SET counters = '{\"t3404\": 1}'::jsonb WHERE id = :id"),
            {"id": env["other"]},
        )
        proxy_egress.resolve_proxy_url(caller, "yandex")
        caller.rollback()
    finally:
        caller.close()

    other = _run_row(env["other"]).counters or {}
    assert "t3404" not in other, "атрибуция закоммитила чужую транзакцию"
    assert _run_row(env["run"]).proxy_id is not None