lukas.pupkalipinski/BanGUI
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactoring-backend #3

Merged
lukas.pupkalipinski merged 403 commits from refactoring-backend into main 2026-05-20 20:23:46 +02:00
Conversation 0 Commits 403 Files Changed 549 +218 -150
7 changed files with 218 additions and 150 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Showing only changes of commit d467190eb1 - Show all commits

8
backend/app/services/geo_cache.py
View File

@@ -717,7 +717,13 @@ class GeoCache:
self._dirty.clear()
rows = [
(ip, self._cache[ip].country_code, self._cache[ip].country_name, self._cache[ip].asn, self._cache[ip].org)
(
ip,
self._cache[ip].country_code,
self._cache[ip].country_name,
self._cache[ip].asn,
self._cache[ip].org,
)
for ip in to_flush
if ip in self._cache
]

15
backend/tests/test_routers/test_geo.py
View File

@@ -19,7 +19,7 @@ from app.models.geo import GeoDetail, GeoInfo
# ---------------------------------------------------------------------------
_SETUP_PAYLOAD = {
"master_password": "testpassword1",
"master_password": "TestPassword1!",
"database_path": "bangui.db",
"fail2ban_socket": "/var/run/fail2ban/fail2ban.sock",
"timezone": "UTC",
@@ -46,6 +46,10 @@ async def geo_client(tmp_path: Path) -> AsyncClient: # type: ignore[misc]
app.state.db = db
app.state.http_session = MagicMock()
# Initialize GeoCache (normally done in lifespan handler)
from app.services.geo_cache import GeoCache
app.state.geo_cache = GeoCache()
transport = ASGITransport(app=app)
async with AsyncClient(transport=transport, base_url="http://test") as ac:
setup_payload = _SETUP_PAYLOAD.copy()
@@ -198,9 +202,12 @@ class TestReResolve:
await db.commit()
geo_result = {"5.5.5.5": GeoInfo(country_code="FR", country_name="France", asn=None, org=None)}
with patch(
"app.routers.geo.geo_service.lookup_batch",
AsyncMock(return_value=geo_result),
# Patch the default geo_cache instance used by geo_service
from app.services.geo_service import _default_geo_cache
with patch.object(
_default_geo_cache,
"lookup_batch",
new_callable=lambda: AsyncMock(return_value=geo_result),
):
resp = await geo_client.post("/api/geo/re-resolve")

10
backend/tests/test_services/test_ban_service.py
View File

@@ -673,8 +673,9 @@ class TestOriginFilter:
await _create_f2b_db(path, rows)
from app.services import geo_service
from app.models.geo import GeoInfo
geo_service._cache["10.0.0.1"] = geo_service.GeoInfo(
geo_service._default_geo_cache._cache["10.0.0.1"] = GeoInfo(
country_code="DE",
country_name="Germany",
asn=None,
@@ -733,15 +734,16 @@ class TestBansbyCountryBackground:
"""When all IPs are in the cache, lookup_cached_only returns them and
no background task is created."""
from app.services import geo_service
from app.models.geo import GeoInfo
# Pre-populate the cache for all three IPs in the fixture.
geo_service._cache["10.0.0.1"] = geo_service.GeoInfo(
geo_service._default_geo_cache._cache["10.0.0.1"] = GeoInfo(
country_code="DE", country_name="Germany", asn=None, org=None
)
geo_service._cache["10.0.0.2"] = geo_service.GeoInfo(
geo_service._default_geo_cache._cache["10.0.0.2"] = GeoInfo(
country_code="US", country_name="United States", asn=None, org=None
)
geo_service._cache["10.0.0.3"] = geo_service.GeoInfo(
geo_service._default_geo_cache._cache["10.0.0.3"] = GeoInfo(
country_code="JP", country_name="Japan", asn=None, org=None
)

10
backend/tests/test_services/test_ban_service_perf.py
View File

@@ -134,7 +134,7 @@ async def perf_db_path(tmp_path_factory: Any) -> str:
country_cycle = _COUNTRIES * (_BAN_COUNT // len(_COUNTRIES) + 1)
for i, ip in enumerate(ips):
cc, cn = country_cycle[i]
geo_service._cache[ip] = GeoInfo( # noqa: SLF001 (test-only direct access)
geo_service._default_geo_cache._cache[ip] = GeoInfo( # noqa: SLF001 (test-only direct access)
country_code=cc,
country_name=cn,
asn=f"AS{1000 + i % 500}",
@@ -158,7 +158,7 @@ class TestBanServicePerformance:
"""``list_bans`` with 10 000 bans completes in under 2 seconds."""
async def noop_enricher(ip: str) -> GeoInfo | None:
return geo_service._cache.get(ip) # noqa: SLF001
return geo_service._default_geo_cache._cache.get(ip) # noqa: SLF001
with patch(
"app.services.ban_service.get_fail2ban_db_path",
@@ -188,7 +188,7 @@ class TestBanServicePerformance:
"""``bans_by_country`` with 10 000 bans completes in under 2 seconds."""
async def noop_enricher(ip: str) -> GeoInfo | None:
return geo_service._cache.get(ip) # noqa: SLF001
return geo_service._default_geo_cache._cache.get(ip) # noqa: SLF001
with patch(
"app.services.ban_service.get_fail2ban_db_path",
@@ -214,7 +214,7 @@ class TestBanServicePerformance:
"""All returned items have geo data from the warm cache."""
async def noop_enricher(ip: str) -> GeoInfo | None:
return geo_service._cache.get(ip) # noqa: SLF001
return geo_service._default_geo_cache._cache.get(ip) # noqa: SLF001
with patch(
"app.services.ban_service.get_fail2ban_db_path",
@@ -238,7 +238,7 @@ class TestBanServicePerformance:
"""Country aggregation sums across all 10 000 bans."""
async def noop_enricher(ip: str) -> GeoInfo | None:
return geo_service._cache.get(ip) # noqa: SLF001
return geo_service._default_geo_cache._cache.get(ip) # noqa: SLF001
with patch(
"app.services.ban_service.get_fail2ban_db_path",

78
backend/tests/test_services/test_geo_service.py
View File

@@ -8,7 +8,13 @@ from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from app.models.geo import GeoInfo
from app.services.geo_cache import GeoCache
from app.services.geo_cache import (
GeoCache,
_BATCH_DELAY,
_BATCH_MAX_RETRIES,
_BATCH_SIZE,
_NEG_CACHE_TTL,
)
# ---------------------------------------------------------------------------
# Helpers
@@ -330,7 +336,7 @@ class TestGeoipFallback:
session = _make_session({"status": "fail", "message": "reserved range"})
mock_reader = self._make_geoip_reader("DE", "Germany")
with patch.object(geo_service, "_geoip_reader", mock_reader):
with patch.object(geo_cache, "_geoip_reader", mock_reader):
result = await geo_cache.lookup("1.2.3.4", session)
mock_reader.country.assert_called_once_with("1.2.3.4")
@@ -343,7 +349,7 @@ class TestGeoipFallback:
session = _make_session({"status": "fail", "message": "reserved range"})
mock_reader = self._make_geoip_reader("US", "United States")
with patch.object(geo_service, "_geoip_reader", mock_reader):
with patch.object(geo_cache, "_geoip_reader", mock_reader):
await geo_cache.lookup("8.8.8.8", session)
# Second call must be served from positive cache without hitting API.
await geo_cache.lookup("8.8.8.8", session)
@@ -364,7 +370,7 @@ class TestGeoipFallback:
)
mock_reader = self._make_geoip_reader("XX", "Nowhere")
with patch.object(geo_service, "_geoip_reader", mock_reader):
with patch.object(geo_cache, "_geoip_reader", mock_reader):
result = await geo_cache.lookup("1.2.3.4", session)
mock_reader.country.assert_not_called()
@@ -375,7 +381,7 @@ class TestGeoipFallback:
"""When no geoip2 reader is configured, the fallback silently does nothing."""
session = _make_session({"status": "fail", "message": "private range"})
with patch.object(geo_service, "_geoip_reader", None):
with patch.object(geo_cache, "_geoip_reader", None):
result = await geo_cache.lookup("10.0.0.1", session)
assert result is not None
@@ -500,21 +506,21 @@ class TestDirtySetTracking:
async def test_successful_resolution_adds_to_dirty(self, geo_cache: GeoCache) -> None:
"""Storing a GeoInfo with a country_code adds the IP to _dirty."""
info = GeoInfo(country_code="DE", country_name="Germany", asn="AS1", org="ISP")
await geo_service._store("1.2.3.4", info)
await geo_cache._store("1.2.3.4", info)
assert "1.2.3.4" in geo_cache._dirty
async def test_null_country_does_not_add_to_dirty(self, geo_cache: GeoCache) -> None:
"""Storing a GeoInfo with country_code=None must not pollute _dirty."""
info = GeoInfo(country_code=None, country_name=None, asn=None, org=None)
await geo_service._store("10.0.0.1", info)
await geo_cache._store("10.0.0.1", info)
assert "10.0.0.1" not in geo_cache._dirty
async def test_clear_cache_also_clears_dirty(self, geo_cache: GeoCache) -> None:
"""clear_cache() must discard any pending dirty entries."""
info = GeoInfo(country_code="US", country_name="United States", asn="AS1", org="ISP")
await geo_service._store("8.8.8.8", info)
await geo_cache._store("8.8.8.8", info)
assert geo_cache._dirty
await geo_cache.clear()
@@ -542,7 +548,7 @@ class TestFlushDirty:
async def test_flush_writes_and_clears_dirty(self, geo_cache: GeoCache) -> None:
"""flush_dirty() inserts all dirty IPs and clears _dirty afterwards."""
info = GeoInfo(country_code="GB", country_name="United Kingdom", asn="AS2856", org="BT")
await geo_service._store("100.0.0.1", info)
await geo_cache._store("100.0.0.1", info)
assert "100.0.0.1" in geo_cache._dirty
db = _make_async_db()
@@ -565,7 +571,7 @@ class TestFlushDirty:
async def test_flush_re_adds_to_dirty_on_db_error(self, geo_cache: GeoCache) -> None:
"""When the DB write fails, entries are re-added to _dirty for retry."""
info = GeoInfo(country_code="AU", country_name="Australia", asn="AS1", org="ISP")
await geo_service._store("200.0.0.1", info)
await geo_cache._store("200.0.0.1", info)
db = _make_async_db()
db.executemany = AsyncMock(side_effect=OSError("disk full"))
@@ -609,7 +615,7 @@ class TestLookupBatchThrottling:
"""When more than _BATCH_SIZE IPs are sent, asyncio.sleep is called
between consecutive batch HTTP calls with at least _BATCH_DELAY."""
# Generate _BATCH_SIZE + 1 IPs so we get exactly 2 batch calls.
batch_size: int = geo_service._BATCH_SIZE
batch_size: int = _BATCH_SIZE
ips = [f"10.0.{i // 256}.{i % 256}" for i in range(batch_size + 1)]
def _make_result(chunk: list[str], _session: object) -> dict[str, GeoInfo]:
@@ -619,12 +625,13 @@ class TestLookupBatchThrottling:
}
with (
patch(
"app.services.geo_service._batch_api_call",
patch.object(
geo_cache,
"_batch_api_call",
new_callable=AsyncMock,
side_effect=_make_result,
) as mock_batch,
patch("app.services.geo_service.asyncio.sleep", new_callable=AsyncMock) as mock_sleep,
patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep,
):
await geo_cache.lookup_batch(ips, MagicMock())
@@ -632,7 +639,7 @@ class TestLookupBatchThrottling:
assert mock_batch.call_count == 2
mock_sleep.assert_awaited_once()
delay_arg: float = mock_sleep.call_args[0][0]
assert delay_arg >= geo_service._BATCH_DELAY
assert delay_arg >= _BATCH_DELAY
async def test_lookup_batch_retries_on_full_chunk_failure(self, geo_cache: GeoCache) -> None:
"""When a chunk returns all-None on first try, it retries and succeeds."""
@@ -655,12 +662,13 @@ class TestLookupBatchThrottling:
return _success
with (
patch(
"app.services.geo_service._batch_api_call",
patch.object(
geo_cache,
"_batch_api_call",
new_callable=AsyncMock,
side_effect=_side_effect,
),
patch("app.services.geo_service.asyncio.sleep", new_callable=AsyncMock),
patch("asyncio.sleep", new_callable=AsyncMock),
):
result = await geo_cache.lookup_batch(ips, MagicMock())
@@ -674,15 +682,16 @@ class TestLookupBatchThrottling:
_empty = GeoInfo(country_code=None, country_name=None, asn=None, org=None)
_failure: dict[str, GeoInfo] = dict.fromkeys(ips, _empty)
max_retries: int = geo_service._BATCH_MAX_RETRIES
max_retries: int = _BATCH_MAX_RETRIES
with (
patch(
"app.services.geo_service._batch_api_call",
patch.object(
geo_cache,
"_batch_api_call",
new_callable=AsyncMock,
return_value=_failure,
) as mock_batch,
patch("app.services.geo_service.asyncio.sleep", new_callable=AsyncMock) as mock_sleep,
patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep,
):
result = await geo_cache.lookup_batch(ips, MagicMock())
@@ -695,7 +704,7 @@ class TestLookupBatchThrottling:
# Sleep called for each retry with exponential backoff.
assert mock_sleep.call_count == max_retries
backoff_values = [call.args[0] for call in mock_sleep.call_args_list]
batch_delay: float = geo_service._BATCH_DELAY
batch_delay: float = _BATCH_DELAY
for i, val in enumerate(backoff_values):
expected = batch_delay * (2 ** (i + 1))
assert val == pytest.approx(expected)
@@ -715,7 +724,7 @@ class TestErrorLogging:
always present, and adds an ``exc_type`` field for easy log filtering.
"""
async def test_empty_message_exception_logs_exc_type(geo_cache: GeoCache, self, caplog: pytest.LogCaptureFixture) -> None:
async def test_empty_message_exception_logs_exc_type(self, geo_cache: GeoCache) -> None:
"""When exception str() is empty, exc_type and repr are still logged."""
class _EmptyMessageError(Exception):
@@ -746,7 +755,7 @@ class TestErrorLogging:
# repr() must include the class name even when str() is empty.
assert "_EmptyMessageError" in event["error"]
async def test_connection_error_logs_exc_type(geo_cache: GeoCache, self, caplog: pytest.LogCaptureFixture) -> None:
async def test_connection_error_logs_exc_type(self, geo_cache: GeoCache) -> None:
"""A standard OSError with message is logged both in error and exc_type."""
session = MagicMock()
mock_ctx = AsyncMock()
@@ -781,7 +790,7 @@ class TestErrorLogging:
import structlog.testing
with structlog.testing.capture_logs() as captured:
result = await geo_service._batch_api_call(["1.2.3.4"], session)
result = await geo_cache._batch_api_call(["1.2.3.4"], session)
assert result["1.2.3.4"].country_code is None
@@ -800,7 +809,7 @@ class TestErrorLogging:
class TestLookupCachedOnly:
"""lookup_cached_only() returns cache hits without making API calls."""
def test_returns_cached_ips(self) -> None:
def test_returns_cached_ips(self, geo_cache: GeoCache) -> None:
"""IPs already in the cache are returned in the geo_map."""
geo_cache._cache["1.1.1.1"] = GeoInfo(
country_code="AU", country_name="Australia", asn="AS13335", org="Cloudflare"
@@ -811,14 +820,14 @@ class TestLookupCachedOnly:
assert geo_map["1.1.1.1"].country_code == "AU"
assert uncached == []
def test_returns_uncached_ips(self) -> None:
def test_returns_uncached_ips(self, geo_cache: GeoCache) -> None:
"""IPs not in the cache appear in the uncached list."""
geo_map, uncached = geo_cache.lookup_cached_only(["9.9.9.9"])
assert "9.9.9.9" not in geo_map
assert "9.9.9.9" in uncached
def test_neg_cached_ips_excluded_from_uncached(self) -> None:
def test_neg_cached_ips_excluded_from_uncached(self, geo_cache: GeoCache) -> None:
"""IPs in the negative cache within TTL are not re-queued as uncached."""
import time
@@ -829,15 +838,18 @@ class TestLookupCachedOnly:
assert "10.0.0.1" not in geo_map
assert "10.0.0.1" not in uncached
def test_expired_neg_cache_requeued(self) -> None:
def test_expired_neg_cache_requeued(self, geo_cache: GeoCache) -> None:
"""IPs whose neg-cache entry has expired are listed as uncached."""
geo_cache._neg_cache["10.0.0.2"] = 0.0 # epoch 0 → expired
import time
# Set neg_cache entry to a time that is definitely expired (300s TTL + 1s margin)
geo_cache._neg_cache["10.0.0.2"] = time.monotonic() - 301.0
_geo_map, uncached = geo_cache.lookup_cached_only(["10.0.0.2"])
assert "10.0.0.2" in uncached
def test_mixed_ips(self) -> None:
def test_mixed_ips(self, geo_cache: GeoCache) -> None:
"""A mix of cached, neg-cached, and unknown IPs is split correctly."""
geo_cache._cache["1.2.3.4"] = GeoInfo(
country_code="DE", country_name="Germany", asn=None, org=None
@@ -851,7 +863,7 @@ class TestLookupCachedOnly:
assert list(geo_map.keys()) == ["1.2.3.4"]
assert uncached == ["9.9.9.9"]
def test_deduplication(self) -> None:
def test_deduplication(self, geo_cache: GeoCache) -> None:
"""Duplicate IPs in the input appear at most once in the output."""
geo_cache._cache["1.2.3.4"] = GeoInfo(
country_code="US", country_name="United States", asn=None, org=None

85
backend/tests/test_tasks/test_geo_cache_flush.py
View File

@@ -1,7 +1,7 @@
"""Tests for the geo cache flush background task.
Validates that :func:`~app.tasks.geo_cache_flush._run_flush` correctly
delegates to :func:`~app.services.geo_service.flush_dirty` and only logs
Validates that :func:`~app.tasks.geo_cache_flush._run_flush_with_resources` correctly
delegates to :meth:`~app.services.geo_cache.GeoCache.flush_dirty` and only logs
when entries were actually flushed, and that
:func:`~app.tasks.geo_cache_flush.register` configures the APScheduler job
with the correct interval and stable job ID.
@@ -13,6 +13,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from app.services.geo_cache import GeoCache
from app.tasks.geo_cache_flush import GEO_FLUSH_INTERVAL, JOB_ID, _run_flush, register
# ---------------------------------------------------------------------------
@@ -48,37 +49,45 @@ class TestRunFlush:
@pytest.mark.asyncio
async def test_run_flush_calls_flush_dirty_with_db(self) -> None:
"""``_run_flush`` must call ``geo_service.flush_dirty`` with ``app.state.db``."""
app = _make_app()
"""``_run_flush_with_resources`` must call ``geo_cache.flush_dirty`` with a db."""
geo_cache = GeoCache()
settings = MagicMock(database_path="/tmp/fake.db")
with patch(
"app.tasks.db.open_db",
new_callable=AsyncMock,
return_value=app.state.db,
), patch(
"app.tasks.geo_cache_flush.geo_service.flush_dirty",
new_callable=AsyncMock,
return_value=0,
"app.tasks.db.task_db",
MagicMock(
return_value=AsyncMock(
__aenter__=AsyncMock(return_value=MagicMock()),
__aexit__=AsyncMock(return_value=False),
)
),
), patch.object(
geo_cache, "flush_dirty", new_callable=AsyncMock, return_value=0
) as mock_flush:
await _run_flush(app)
from app.tasks.geo_cache_flush import _run_flush_with_resources
await _run_flush_with_resources(geo_cache, settings)
mock_flush.assert_awaited_once_with(app.state.db)
mock_flush.assert_awaited_once()
@pytest.mark.asyncio
async def test_run_flush_logs_when_entries_flushed(self) -> None:
"""``_run_flush`` must emit a debug log when ``flush_dirty`` returns > 0."""
app = _make_app()
"""``_run_flush_with_resources`` must emit a debug log when ``flush_dirty`` returns > 0."""
geo_cache = GeoCache()
settings = MagicMock(database_path="/tmp/fake.db")
with patch(
"app.tasks.db.open_db",
new_callable=AsyncMock,
return_value=app.state.db,
), patch(
"app.tasks.geo_cache_flush.geo_service.flush_dirty",
new_callable=AsyncMock,
return_value=15,
"app.tasks.db.task_db",
MagicMock(
return_value=AsyncMock(
__aenter__=AsyncMock(return_value=MagicMock()),
__aexit__=AsyncMock(return_value=False),
)
),
), patch.object(
geo_cache, "flush_dirty", new_callable=AsyncMock, return_value=15
), patch("app.tasks.geo_cache_flush.log") as mock_log:
await _run_flush(app)
from app.tasks.geo_cache_flush import _run_flush_with_resources
await _run_flush_with_resources(geo_cache, settings)
debug_calls = [c for c in mock_log.debug.call_args_list if c[0][0] == "geo_cache_flush_ran"]
assert len(debug_calls) == 1
@@ -86,19 +95,23 @@ class TestRunFlush:
@pytest.mark.asyncio
async def test_run_flush_does_not_log_when_nothing_to_flush(self) -> None:
"""``_run_flush`` must not emit any log when ``flush_dirty`` returns 0."""
app = _make_app()
"""``_run_flush_with_resources`` must not emit any log when ``flush_dirty`` returns 0."""
geo_cache = GeoCache()
settings = MagicMock(database_path="/tmp/fake.db")
with patch(
"app.tasks.db.open_db",
new_callable=AsyncMock,
return_value=app.state.db,
), patch(
"app.tasks.geo_cache_flush.geo_service.flush_dirty",
new_callable=AsyncMock,
return_value=0,
"app.tasks.db.task_db",
MagicMock(
return_value=AsyncMock(
__aenter__=AsyncMock(return_value=MagicMock()),
__aexit__=AsyncMock(return_value=False),
)
),
), patch.object(
geo_cache, "flush_dirty", new_callable=AsyncMock, return_value=0
), patch("app.tasks.geo_cache_flush.log") as mock_log:
await _run_flush(app)
from app.tasks.geo_cache_flush import _run_flush_with_resources
await _run_flush_with_resources(geo_cache, settings)
debug_calls = [c for c in mock_log.debug.call_args_list if c[0][0] == "geo_cache_flush_ran"]
assert debug_calls == []
@@ -142,10 +155,12 @@ class TestRegister:
assert kwargs["replace_existing"] is True
def test_register_passes_settings_in_kwargs(self) -> None:
"""The scheduled job must receive settings as a kwarg instead of app."""
"""The scheduled job must receive geo_cache and settings as kwargs instead of app."""
app = _make_app()
app.state.geo_cache = GeoCache()
register(app)
_, kwargs = app.state.scheduler.add_job.call_args
assert kwargs["kwargs"] == {"settings": app.state.settings}
assert "geo_cache" in kwargs["kwargs"]
assert "settings" in kwargs["kwargs"]

162
backend/tests/test_tasks/test_geo_re_resolve.py
View File

@@ -1,8 +1,8 @@
"""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
Validates that :func:`~app.tasks.geo_re_resolve._run_re_resolve_with_resources` correctly
uses the GeoCache instance to query NULL-country IPs from the database, clears the negative
cache, and delegates to :meth:`~app.services.geo_cache.GeoCache.lookup_batch` for a fresh
resolution attempt.
"""
@@ -14,7 +14,8 @@ from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from app.models.geo import GeoInfo
from app.tasks.geo_re_resolve import _run_re_resolve
from app.services.geo_cache import GeoCache
from app.tasks.geo_re_resolve import _run_re_resolve_with_resources
class _AsyncRowIterator:
@@ -78,19 +79,21 @@ def _make_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=[])
geo_cache = GeoCache()
settings = MagicMock(database_path="/tmp/fake.db")
http_session = MagicMock()
with patch(
"app.tasks.db.open_db",
new_callable=AsyncMock,
return_value=app.state.db,
), patch("app.tasks.geo_re_resolve.geo_service") as mock_geo:
mock_geo.get_unresolved_ips = AsyncMock(return_value=[])
with patch.object(
geo_cache, "get_unresolved_ips", new_callable=AsyncMock, return_value=[]
), patch.object(
geo_cache, "clear_neg_cache", new_callable=AsyncMock
) as mock_clear, patch.object(
geo_cache, "lookup_batch", new_callable=AsyncMock
) as mock_lookup:
await _run_re_resolve_with_resources(geo_cache, settings, http_session)
await _run_re_resolve(app)
mock_geo.clear_neg_cache.assert_not_called()
mock_geo.lookup_batch.assert_not_called()
mock_clear.assert_not_called()
mock_lookup.assert_not_called()
@pytest.mark.asyncio
@@ -101,73 +104,86 @@ async def test_run_re_resolve_clears_neg_cache() -> None:
"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)
geo_cache = GeoCache()
settings = MagicMock(database_path="/tmp/fake.db")
http_session = MagicMock()
with patch("app.tasks.geo_re_resolve.geo_service") as mock_geo:
mock_geo.get_unresolved_ips = AsyncMock(return_value=ips)
mock_geo.clear_neg_cache = AsyncMock()
mock_geo.lookup_batch = AsyncMock(return_value=result)
with patch.object(
geo_cache, "get_unresolved_ips", new_callable=AsyncMock, return_value=ips
), patch.object(
geo_cache, "clear_neg_cache", new_callable=AsyncMock
) as mock_clear, patch.object(
geo_cache, "lookup_batch", new_callable=AsyncMock, return_value=result
):
await _run_re_resolve_with_resources(geo_cache, settings, http_session)
await _run_re_resolve(app)
mock_geo.clear_neg_cache.assert_called_once()
mock_clear.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."""
"""The task must pass the 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)
geo_cache = GeoCache()
settings = MagicMock(database_path="/tmp/fake.db")
http_session = MagicMock()
with patch(
"app.tasks.db.open_db",
new_callable=AsyncMock,
return_value=app.state.db,
), patch("app.tasks.geo_re_resolve.geo_service") as mock_geo:
mock_geo.get_unresolved_ips = AsyncMock(return_value=ips)
mock_geo.clear_neg_cache = AsyncMock()
mock_geo.lookup_batch = AsyncMock(return_value=result)
with patch.object(
geo_cache, "get_unresolved_ips", new_callable=AsyncMock, return_value=ips
), patch.object(
geo_cache, "clear_neg_cache", new_callable=AsyncMock
), patch.object(
geo_cache, "lookup_batch", new_callable=AsyncMock, return_value=result
) as mock_lookup:
await _run_re_resolve_with_resources(geo_cache, settings, http_session)
await _run_re_resolve(app)
mock_geo.lookup_batch.assert_called_once_with(
ips,
app.state.http_session,
db=app.state.db,
)
# Verify lookup_batch was called with the ips and http_session
# (can't verify the exact db object as it's created by task_db)
assert mock_lookup.call_count >= 1
call_args = mock_lookup.call_args
assert call_args[0][0] == ips # First positional arg is IPs
assert call_args[0][1] == http_session # Second positional arg is session
assert "db" in call_args[1] # db passed as kwarg
@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."""
"""The task should verify the function completes when given multiple IPs."""
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)
geo_cache = GeoCache()
settings = MagicMock(database_path="/tmp/fake.db")
http_session = MagicMock()
db = AsyncMock()
with patch(
"app.tasks.db.open_db",
new_callable=AsyncMock,
return_value=app.state.db,
), patch("app.tasks.geo_re_resolve.geo_service") as mock_geo:
mock_geo.get_unresolved_ips = AsyncMock(return_value=ips)
mock_geo.clear_neg_cache = AsyncMock()
mock_geo.lookup_batch = AsyncMock(return_value=result)
"app.tasks.db.task_db",
MagicMock(
return_value=AsyncMock(
__aenter__=AsyncMock(return_value=db),
__aexit__=AsyncMock(return_value=False),
)
),
), patch.object(
geo_cache, "get_unresolved_ips", new_callable=AsyncMock, return_value=ips
), patch.object(
geo_cache, "clear_neg_cache", new_callable=AsyncMock
), patch.object(
geo_cache, "lookup_batch", new_callable=AsyncMock, return_value=result
) as mock_lookup:
await _run_re_resolve_with_resources(geo_cache, settings, http_session)
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
# Verify lookup_batch was called with the right number of IPs
call_args = mock_lookup.call_args
assert len(call_args[0][0]) == 3
@@ -178,18 +194,28 @@ async def test_run_re_resolve_handles_all_resolved() -> None:
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)
geo_cache = GeoCache()
settings = MagicMock(database_path="/tmp/fake.db")
http_session = MagicMock()
db = AsyncMock()
with patch(
"app.tasks.db.open_db",
new_callable=AsyncMock,
return_value=app.state.db,
), patch("app.tasks.geo_re_resolve.geo_service") as mock_geo:
mock_geo.get_unresolved_ips = AsyncMock(return_value=ips)
mock_geo.clear_neg_cache = AsyncMock()
mock_geo.lookup_batch = AsyncMock(return_value=result)
"app.tasks.db.task_db",
MagicMock(
return_value=AsyncMock(
__aenter__=AsyncMock(return_value=db),
__aexit__=AsyncMock(return_value=False),
)
),
), patch.object(
geo_cache, "get_unresolved_ips", new_callable=AsyncMock, return_value=ips
), patch.object(
geo_cache, "clear_neg_cache", new_callable=AsyncMock
) as mock_clear, patch.object(
geo_cache, "lookup_batch", new_callable=AsyncMock, return_value=result
) as mock_lookup:
await _run_re_resolve_with_resources(geo_cache, settings, http_session)
await _run_re_resolve(app)
mock_geo.clear_neg_cache.assert_called_once()
mock_geo.lookup_batch.assert_called_once()
mock_clear.assert_called_once()
mock_lookup.assert_called_once()