Добавляет scripts/geocode_deals_nominatim.py для ~14% сделок всё ещё
lat IS NULL после house-centroid прохода (Step 2) — улицы без геокодированных
домов, которые centroid-join не покрывает. Переиспользует geocoder.geocode
(cache → Cadastral FDW → Yandex → Nominatim), дедупит по адресу (distinct
улицы определяют число вызовов, не 7k строк), и стампит geocode_tried_at на
КАЖДОЙ попытке (успех И чистый промах) чтобы un-geocodable строки не крутились
по 1 req/s. Resume: lat IS NULL + 30-day staleness, per-address SAVEPOINT,
--limit/--dry-run/--batch/--stale-days.
Pre-push review (code-reviewer): ✅ APPROVE, no critical — проверены SQL-safety,
no-infinite-reloop (mark-tried на miss, no-stamp на raised exc), window-cap,
geom-not-manual, dedup, async. Tests: 17 passed. Ruff clean.
456 lines
17 KiB
Python
456 lines
17 KiB
Python
"""Unit tests for issue #569 Step 3 — geocode_deals_nominatim.py.
|
||
|
||
Coverage (per the issue's test plan):
|
||
- a NULL-coords deal gets geocoded → UPDATE deals issued with lat/lon +
|
||
geocode_tried_at, counted as geocoded, deals_updated reflects rowcount.
|
||
- geocode() returns None → geocode_tried_at still stamped (a "mark tried"
|
||
UPDATE with no lat/lon), no coords UPDATE, counted as geocode_failed.
|
||
- --dry-run issues no UPDATE / no commit, counters still move.
|
||
- the resume filter excludes already-tried-recently rows (SQL carries the
|
||
lat IS NULL + staleness predicate, and binds --stale-days through).
|
||
- per-address SAVEPOINT isolation: one failing UPDATE doesn't abort the
|
||
batch.
|
||
- dedup: distinct addresses drive geocode call count (one call per address).
|
||
- main() wiring: SessionLocal, --limit bind, returns geocoded count.
|
||
|
||
No real Postgres. `geocode()` is async → patched with AsyncMock (mirrors
|
||
tests/tasks/test_geocode_missing.py). The Session is a MagicMock that routes
|
||
SELECT side-effects by SQL substring and records UPDATE binds (mirrors
|
||
tests/scripts/test_geocode_deals_from_houses.py).
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
from unittest.mock import AsyncMock, MagicMock, patch
|
||
|
||
# Settings needs a DSN at import time — set a dummy before any app.* import.
|
||
os.environ.setdefault("DATABASE_URL", "postgresql+psycopg://test:test@localhost/test_db")
|
||
|
||
from app.services.geocoder import GeocodeResult
|
||
from scripts.geocode_deals_nominatim import (
|
||
AddressGroup,
|
||
Stats,
|
||
_mark_deals_tried,
|
||
_run_backfill,
|
||
_select_pending_addresses,
|
||
_update_deals_geocoded,
|
||
main,
|
||
)
|
||
|
||
_GEOCODE_PATH = "scripts.geocode_deals_nominatim.geocode"
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Helpers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _result(provider: str = "nominatim", lat: float = 56.838, lon: float = 60.605) -> GeocodeResult:
|
||
return GeocodeResult(
|
||
lat=lat,
|
||
lon=lon,
|
||
full_address="Екатеринбург, ул. Тестовая",
|
||
provider=provider, # type: ignore[arg-type]
|
||
confidence="approximate",
|
||
)
|
||
|
||
|
||
def _make_db_mock(
|
||
*,
|
||
address_rows: list[dict] | None = None,
|
||
total_streets: int = 0,
|
||
total_rows: int = 0,
|
||
update_rowcount: int = 1,
|
||
) -> tuple[MagicMock, list[dict], list[dict]]:
|
||
"""MagicMock Session that:
|
||
|
||
- returns `address_rows` for the candidate SELECT (`FROM deals` GROUP BY)
|
||
- returns (total_streets, total_rows) for the dry-run COUNT
|
||
- records coords-UPDATE binds into `coord_updates`
|
||
- records mark-tried-UPDATE binds into `tried_updates`
|
||
- every UPDATE returns a result with `.rowcount == update_rowcount`
|
||
- supports `db.begin_nested()` as a context manager
|
||
|
||
Returns (db, coord_updates, tried_updates).
|
||
"""
|
||
address_rows = address_rows or []
|
||
coord_updates: list[dict] = []
|
||
tried_updates: list[dict] = []
|
||
|
||
db = MagicMock()
|
||
db.begin_nested.return_value.__enter__ = lambda self: self
|
||
db.begin_nested.return_value.__exit__ = lambda self, *a: False
|
||
|
||
def execute_side_effect(sql, params=None):
|
||
sql_str = str(sql)
|
||
result = MagicMock()
|
||
if "UPDATE deals" in sql_str:
|
||
# Distinguish the coords UPDATE (sets lat) from the mark-tried UPDATE.
|
||
if "SET lat" in sql_str:
|
||
coord_updates.append(dict(params) if params else {})
|
||
else:
|
||
tried_updates.append(dict(params) if params else {})
|
||
result.rowcount = update_rowcount
|
||
return result
|
||
if "COUNT(DISTINCT address)" in sql_str:
|
||
result.first.return_value = (total_streets, total_rows)
|
||
return result
|
||
if "FROM deals" in sql_str:
|
||
result.mappings.return_value.all.return_value = address_rows
|
||
return result
|
||
return result
|
||
|
||
db.execute.side_effect = execute_side_effect
|
||
db.commit = MagicMock()
|
||
db.rollback = MagicMock()
|
||
db.close = MagicMock()
|
||
return db, coord_updates, tried_updates
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# _run_backfill — happy path: geocoded → coords UPDATE + tried_at stamped
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
async def test_run_backfill_geocoded_issues_coords_update():
|
||
groups = [AddressGroup(address="Екатеринбург, Тестовая", deals_count=5)]
|
||
db, coord_updates, tried_updates = _make_db_mock(update_rowcount=5)
|
||
|
||
with patch(_GEOCODE_PATH, new_callable=AsyncMock, return_value=_result("nominatim")):
|
||
stats = await _run_backfill(db, groups, batch="b1", dry_run=False)
|
||
|
||
assert stats.geocoded == 1
|
||
assert stats.geocode_failed == 0
|
||
assert stats.processed == 1
|
||
assert stats.deals_updated == 5
|
||
assert stats.cache_misses == 1
|
||
assert stats.cache_hits == 0
|
||
# Exactly one coords UPDATE, bound with lat/lon/addr; no mark-tried UPDATE.
|
||
assert len(coord_updates) == 1
|
||
assert coord_updates[0]["addr"] == "Екатеринбург, Тестовая"
|
||
assert coord_updates[0]["lat"] == 56.838
|
||
assert coord_updates[0]["lon"] == 60.605
|
||
assert tried_updates == []
|
||
assert db.commit.call_count == 1
|
||
|
||
|
||
async def test_run_backfill_cache_hit_counted_separately():
|
||
groups = [AddressGroup(address="Екатеринбург, Кэшевая", deals_count=2)]
|
||
db, coord_updates, _ = _make_db_mock(update_rowcount=2)
|
||
|
||
with patch(_GEOCODE_PATH, new_callable=AsyncMock, return_value=_result("cache")):
|
||
stats = await _run_backfill(db, groups, batch="b", dry_run=False)
|
||
|
||
assert stats.geocoded == 1
|
||
assert stats.cache_hits == 1
|
||
assert stats.cache_misses == 0
|
||
assert len(coord_updates) == 1
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# geocode None → geocode_tried_at stamped, NO coords UPDATE, counted failed
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
async def test_run_backfill_geocode_none_stamps_tried_no_coords():
|
||
groups = [AddressGroup(address="Екатеринбург, Несуществующая", deals_count=3)]
|
||
db, coord_updates, tried_updates = _make_db_mock(update_rowcount=3)
|
||
|
||
with patch(_GEOCODE_PATH, new_callable=AsyncMock, return_value=None):
|
||
stats = await _run_backfill(db, groups, batch="b", dry_run=False)
|
||
|
||
assert stats.geocode_failed == 1
|
||
assert stats.geocoded == 0
|
||
assert stats.processed == 1
|
||
assert stats.deals_updated == 0
|
||
# No coords UPDATE…
|
||
assert coord_updates == []
|
||
# …but geocode_tried_at IS stamped via the mark-tried UPDATE.
|
||
assert len(tried_updates) == 1
|
||
assert tried_updates[0]["addr"] == "Екатеринбург, Несуществующая"
|
||
assert db.commit.call_count == 1
|
||
|
||
|
||
async def test_mark_deals_tried_sql_stamps_tried_at_only():
|
||
"""The mark-tried writer sets geocode_tried_at and nothing else (no lat/lon)."""
|
||
db = MagicMock()
|
||
db.execute.return_value = MagicMock(rowcount=4)
|
||
n = _mark_deals_tried(db, address="Екатеринбург, X")
|
||
args, _kw = db.execute.call_args
|
||
sql_str = str(args[0])
|
||
assert "UPDATE deals" in sql_str
|
||
assert "geocode_tried_at = NOW()" in sql_str
|
||
assert "SET lat" not in sql_str
|
||
assert "lat IS NULL" in sql_str # idempotency / no-clobber guard
|
||
assert n == 4
|
||
|
||
|
||
async def test_run_backfill_geocode_raises_not_stamped():
|
||
"""A raised geocode error → counted failed but NOT stamped (retry sooner)."""
|
||
groups = [AddressGroup(address="Екатеринбург, Сетевая", deals_count=1)]
|
||
db, coord_updates, tried_updates = _make_db_mock()
|
||
|
||
with patch(_GEOCODE_PATH, new_callable=AsyncMock, side_effect=RuntimeError("network")):
|
||
stats = await _run_backfill(db, groups, batch="b", dry_run=False)
|
||
|
||
assert stats.geocode_failed == 1
|
||
assert stats.geocoded == 0
|
||
# Neither a coords UPDATE nor a mark-tried UPDATE — leave it for next run.
|
||
assert coord_updates == []
|
||
assert tried_updates == []
|
||
assert db.commit.call_count == 0
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# --dry-run — no DB writes
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
async def test_run_backfill_dry_run_issues_no_update():
|
||
groups = [
|
||
AddressGroup(address="Екатеринбург, Тестовая", deals_count=4),
|
||
AddressGroup(address="Екатеринбург, Пустая", deals_count=1),
|
||
]
|
||
db, coord_updates, tried_updates = _make_db_mock()
|
||
|
||
# First address resolves, second is a miss — neither must write in dry-run.
|
||
with patch(
|
||
_GEOCODE_PATH,
|
||
new_callable=AsyncMock,
|
||
side_effect=[_result("nominatim"), None],
|
||
):
|
||
stats = await _run_backfill(db, groups, batch="dry", dry_run=True)
|
||
|
||
assert stats.geocoded == 1
|
||
assert stats.geocode_failed == 1
|
||
assert stats.processed == 2
|
||
assert coord_updates == []
|
||
assert tried_updates == []
|
||
assert db.commit.call_count == 0
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# dedup — one geocode call per distinct address
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
async def test_run_backfill_one_call_per_distinct_address():
|
||
groups = [
|
||
AddressGroup(address="Екатеринбург, А", deals_count=10),
|
||
AddressGroup(address="Екатеринбург, Б", deals_count=2),
|
||
]
|
||
db, coord_updates, _ = _make_db_mock(update_rowcount=1)
|
||
|
||
with patch(
|
||
_GEOCODE_PATH, new_callable=AsyncMock, return_value=_result("nominatim")
|
||
) as mock_geo:
|
||
stats = await _run_backfill(db, groups, batch="b", dry_run=False)
|
||
|
||
# 2 distinct addresses → exactly 2 geocode calls (not 12 = sum of deals).
|
||
assert mock_geo.call_count == 2
|
||
assert stats.geocoded == 2
|
||
assert len(coord_updates) == 2
|
||
|
||
|
||
async def test_run_backfill_skips_too_short_address_without_geocoding():
|
||
groups = [AddressGroup(address="ек", deals_count=1)] # < 3 chars
|
||
db, coord_updates, _ = _make_db_mock()
|
||
|
||
with patch(_GEOCODE_PATH, new_callable=AsyncMock) as mock_geo:
|
||
stats = await _run_backfill(db, groups, batch="b", dry_run=False)
|
||
|
||
assert stats.skipped == 1
|
||
assert stats.processed == 0
|
||
mock_geo.assert_not_called()
|
||
assert coord_updates == []
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# per-address SAVEPOINT — one bad UPDATE doesn't abort the batch
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
async def test_run_backfill_db_write_failure_isolated_to_address():
|
||
groups = [
|
||
AddressGroup(address="Екатеринбург, Первая", deals_count=1),
|
||
AddressGroup(address="Екатеринбург, Вторая", deals_count=1),
|
||
]
|
||
db, coord_updates, _ = _make_db_mock(update_rowcount=1)
|
||
|
||
# Make the FIRST coords UPDATE raise, the rest succeed.
|
||
call = {"n": 0}
|
||
real_side_effect = db.execute.side_effect
|
||
|
||
def failing_execute(sql, params=None):
|
||
sql_str = str(sql)
|
||
if "UPDATE deals" in sql_str and "SET lat" in sql_str:
|
||
call["n"] += 1
|
||
if call["n"] == 1:
|
||
raise RuntimeError("constraint blew up")
|
||
return real_side_effect(sql, params)
|
||
|
||
db.execute.side_effect = failing_execute
|
||
|
||
with patch(_GEOCODE_PATH, new_callable=AsyncMock, return_value=_result("nominatim")):
|
||
stats = await _run_backfill(db, groups, batch="b", dry_run=False)
|
||
|
||
# Both addresses geocoded; one write failed but was isolated via SAVEPOINT.
|
||
assert stats.geocoded == 2
|
||
assert db.rollback.call_count == 1
|
||
# Only the second address's coords UPDATE was recorded.
|
||
assert [u["addr"] for u in coord_updates] == ["Екатеринбург, Вторая"]
|
||
# deals_updated only counts the successful write.
|
||
assert stats.deals_updated == 1
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# _update_deals_geocoded — bind shape + geom NOT set manually
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def test_update_deals_geocoded_sets_lat_lon_tried_at_not_geom():
|
||
db = MagicMock()
|
||
db.execute.return_value = MagicMock(rowcount=7)
|
||
n = _update_deals_geocoded(db, address="Екатеринбург, Y", lat=56.1, lon=60.2)
|
||
args, _kw = db.execute.call_args
|
||
sql_str = str(args[0])
|
||
binds = args[1]
|
||
assert "UPDATE deals" in sql_str
|
||
assert "geocode_tried_at = NOW()" in sql_str
|
||
# geom must NOT be set manually — the deals_set_geom_trg trigger fills it.
|
||
assert "geom" not in sql_str
|
||
assert "lat IS NULL" in sql_str # no-clobber guard
|
||
assert binds == {"addr": "Екатеринбург, Y", "lat": 56.1, "lon": 60.2}
|
||
assert n == 7
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# resume filter — candidate SQL carries lat IS NULL + staleness predicate
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def test_select_pending_addresses_filters_null_and_stale():
|
||
"""The candidate query must combine lat IS NULL with the staleness window
|
||
and bind --stale-days, so already-tried-recently rows are excluded."""
|
||
db, _, _ = _make_db_mock(
|
||
address_rows=[{"address": "Екатеринбург, Z", "deals_count": 3}]
|
||
)
|
||
|
||
groups = _select_pending_addresses(db, limit=2000, stale_days=30)
|
||
|
||
assert groups == [AddressGroup(address="Екатеринбург, Z", deals_count=3)]
|
||
# Inspect the SQL + binds of the SELECT.
|
||
select_call = next(
|
||
c for c in db.execute.call_args_list if "FROM deals" in str(c[0][0])
|
||
)
|
||
sql_str = str(select_call[0][0])
|
||
binds = select_call[0][1]
|
||
assert "lat IS NULL" in sql_str
|
||
assert "geocode_tried_at IS NULL" in sql_str
|
||
assert "make_interval(days => CAST(:stale_days AS int))" in sql_str
|
||
assert "GROUP BY address" in sql_str
|
||
assert binds["stale_days"] == 30
|
||
assert binds["limit"] == 2000
|
||
|
||
|
||
def test_select_pending_addresses_passes_custom_stale_days():
|
||
db, _, _ = _make_db_mock(address_rows=[])
|
||
_select_pending_addresses(db, limit=500, stale_days=7)
|
||
select_call = next(
|
||
c for c in db.execute.call_args_list if "FROM deals" in str(c[0][0])
|
||
)
|
||
binds = select_call[0][1]
|
||
assert binds["stale_days"] == 7
|
||
assert binds["limit"] == 500
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# main() — wiring, --limit, --dry-run, return value
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
async def test_main_geocodes_and_returns_count():
|
||
address_rows = [
|
||
{"address": "Екатеринбург, Малышева", "deals_count": 4},
|
||
{"address": "Екатеринбург, Ленина", "deals_count": 2},
|
||
]
|
||
db, coord_updates, _ = _make_db_mock(address_rows=address_rows, update_rowcount=4)
|
||
|
||
with (
|
||
patch("scripts.geocode_deals_nominatim.SessionLocal", return_value=db),
|
||
patch(_GEOCODE_PATH, new_callable=AsyncMock, return_value=_result("nominatim")),
|
||
):
|
||
n = await main(["--batch", "test_main"])
|
||
|
||
assert n == 2 # both distinct addresses geocoded
|
||
assert len(coord_updates) == 2
|
||
assert {u["addr"] for u in coord_updates} == {
|
||
"Екатеринбург, Малышева",
|
||
"Екатеринбург, Ленина",
|
||
}
|
||
|
||
|
||
async def test_main_dry_run_writes_nothing():
|
||
address_rows = [{"address": "Екатеринбург, Малышева", "deals_count": 3}]
|
||
db, coord_updates, tried_updates = _make_db_mock(
|
||
address_rows=address_rows, total_streets=1, total_rows=3
|
||
)
|
||
|
||
with (
|
||
patch("scripts.geocode_deals_nominatim.SessionLocal", return_value=db),
|
||
patch(_GEOCODE_PATH, new_callable=AsyncMock, return_value=_result("nominatim")),
|
||
):
|
||
n = await main(["--dry-run"])
|
||
|
||
# dry-run returns the would-geocode count, writes nothing.
|
||
assert n == 1
|
||
assert coord_updates == []
|
||
assert tried_updates == []
|
||
assert db.commit.call_count == 0
|
||
|
||
|
||
async def test_main_respects_limit_bind():
|
||
db, _, _ = _make_db_mock(address_rows=[])
|
||
|
||
with (
|
||
patch("scripts.geocode_deals_nominatim.SessionLocal", return_value=db),
|
||
patch(_GEOCODE_PATH, new_callable=AsyncMock),
|
||
):
|
||
await main(["--limit", "50"])
|
||
|
||
select_call = next(
|
||
c for c in db.execute.call_args_list if "FROM deals" in str(c[0][0])
|
||
)
|
||
assert select_call[0][1]["limit"] == 50
|
||
|
||
|
||
async def test_main_no_pending_returns_zero():
|
||
db, coord_updates, _ = _make_db_mock(address_rows=[])
|
||
|
||
with (
|
||
patch("scripts.geocode_deals_nominatim.SessionLocal", return_value=db),
|
||
patch(_GEOCODE_PATH, new_callable=AsyncMock) as mock_geo,
|
||
):
|
||
n = await main([])
|
||
|
||
assert n == 0
|
||
assert coord_updates == []
|
||
mock_geo.assert_not_called()
|
||
db.close.assert_called_once()
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Stats dataclass
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def test_stats_defaults_zero():
|
||
s = Stats()
|
||
assert s.processed == 0
|
||
assert s.geocoded == 0
|
||
assert s.geocode_failed == 0
|
||
assert s.skipped == 0
|
||
assert s.deals_updated == 0
|
||
assert s.cache_hits == 0
|
||
assert s.cache_misses == 0
|