"""Tests for the geo re-resolve background task. Validates that :func:`~app.tasks.geo_re_resolve._run_re_resolve` correctly queries NULL-country IPs from the database, clears the negative cache, and delegates to :func:`~app.services.geo_service.lookup_batch` for a fresh resolution attempt. """ from __future__ import annotations from typing import Any from unittest.mock import AsyncMock, MagicMock, patch import pytest from app.services.geo_service import GeoInfo from app.tasks.geo_re_resolve import _run_re_resolve class _AsyncRowIterator: """Minimal async iterator over a list of row tuples.""" def __init__(self, rows: list[tuple[str]]) -> None: self._iter = iter(rows) def __aiter__(self) -> _AsyncRowIterator: return self async def __anext__(self) -> tuple[str]: try: return next(self._iter) except StopIteration: raise StopAsyncIteration # noqa: B904 def _make_app( unresolved_ips: list[str], lookup_result: dict[str, GeoInfo] | None = None, ) -> MagicMock: """Build a minimal mock ``app`` with ``state.db`` and ``state.http_session``. The mock database returns *unresolved_ips* when the re-resolve task queries ``SELECT ip FROM geo_cache WHERE country_code IS NULL``. Args: unresolved_ips: IPs to return from the mocked DB query. lookup_result: Value returned by the mocked ``lookup_batch``. Defaults to an empty dict. Returns: A :class:`unittest.mock.MagicMock` that mimics ``fastapi.FastAPI``. """ if lookup_result is None: lookup_result = {} rows = [(ip,) for ip in unresolved_ips] cursor = _AsyncRowIterator(rows) # db.execute() returns an async context manager yielding the cursor. ctx = AsyncMock() ctx.__aenter__ = AsyncMock(return_value=cursor) ctx.__aexit__ = AsyncMock(return_value=False) db = AsyncMock() db.execute = MagicMock(return_value=ctx) http_session = MagicMock() app = MagicMock() app.state.db = db app.state.http_session = http_session return app @pytest.mark.asyncio async def test_run_re_resolve_no_unresolved_ips_skips() -> None: """The task should return immediately when no NULL-country IPs exist.""" app = _make_app(unresolved_ips=[]) with patch("app.tasks.geo_re_resolve.geo_service") as mock_geo: await _run_re_resolve(app) mock_geo.clear_neg_cache.assert_not_called() mock_geo.lookup_batch.assert_not_called() @pytest.mark.asyncio async def test_run_re_resolve_clears_neg_cache() -> None: """The task must clear the negative cache before calling lookup_batch.""" ips = ["1.2.3.4", "5.6.7.8"] result: dict[str, GeoInfo] = { "1.2.3.4": GeoInfo(country_code="DE", country_name="Germany", asn="AS3320", org="DTAG"), "5.6.7.8": GeoInfo(country_code="US", country_name="United States", asn="AS15169", org="Google"), } app = _make_app(unresolved_ips=ips, lookup_result=result) with patch("app.tasks.geo_re_resolve.geo_service") as mock_geo: mock_geo.lookup_batch = AsyncMock(return_value=result) await _run_re_resolve(app) mock_geo.clear_neg_cache.assert_called_once() @pytest.mark.asyncio async def test_run_re_resolve_calls_lookup_batch_with_db() -> None: """The task must pass the real db to lookup_batch for persistence.""" ips = ["10.0.0.1", "10.0.0.2"] result: dict[str, GeoInfo] = { "10.0.0.1": GeoInfo(country_code="FR", country_name="France", asn=None, org=None), "10.0.0.2": GeoInfo(country_code=None, country_name=None, asn=None, org=None), } app = _make_app(unresolved_ips=ips, lookup_result=result) with patch("app.tasks.geo_re_resolve.geo_service") as mock_geo: mock_geo.lookup_batch = AsyncMock(return_value=result) await _run_re_resolve(app) mock_geo.lookup_batch.assert_called_once_with( ips, app.state.http_session, db=app.state.db, ) @pytest.mark.asyncio async def test_run_re_resolve_logs_correct_counts(caplog: Any) -> None: """The task should log the number retried and number resolved.""" ips = ["1.1.1.1", "2.2.2.2", "3.3.3.3"] result: dict[str, GeoInfo] = { "1.1.1.1": GeoInfo(country_code="AU", country_name="Australia", asn=None, org=None), "2.2.2.2": GeoInfo(country_code="JP", country_name="Japan", asn=None, org=None), "3.3.3.3": GeoInfo(country_code=None, country_name=None, asn=None, org=None), } app = _make_app(unresolved_ips=ips, lookup_result=result) with patch("app.tasks.geo_re_resolve.geo_service") as mock_geo: mock_geo.lookup_batch = AsyncMock(return_value=result) await _run_re_resolve(app) # Verify lookup_batch was called (the logging assertions rely on # structlog which is hard to capture in caplog; instead we verify # the function ran to completion and the counts are correct by # checking that lookup_batch received the right number of IPs). call_args = mock_geo.lookup_batch.call_args assert len(call_args[0][0]) == 3 @pytest.mark.asyncio async def test_run_re_resolve_handles_all_resolved() -> None: """When every IP resolves successfully the task should complete normally.""" ips = ["4.4.4.4"] result: dict[str, GeoInfo] = { "4.4.4.4": GeoInfo(country_code="GB", country_name="United Kingdom", asn=None, org=None), } app = _make_app(unresolved_ips=ips, lookup_result=result) with patch("app.tasks.geo_re_resolve.geo_service") as mock_geo: mock_geo.lookup_batch = AsyncMock(return_value=result) await _run_re_resolve(app) mock_geo.clear_neg_cache.assert_called_once() mock_geo.lookup_batch.assert_called_once()