BanGUI/backend/tests/test_tasks/test_geo_re_resolve.py

"""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()