diff --git a/backend/app/services/geo_cache.py b/backend/app/services/geo_cache.py index 73a9ead..04419b4 100644 --- a/backend/app/services/geo_cache.py +++ b/backend/app/services/geo_cache.py @@ -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 ] diff --git a/backend/tests/test_routers/test_geo.py b/backend/tests/test_routers/test_geo.py index aad1419..c718a4e 100644 --- a/backend/tests/test_routers/test_geo.py +++ b/backend/tests/test_routers/test_geo.py @@ -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") diff --git a/backend/tests/test_services/test_ban_service.py b/backend/tests/test_services/test_ban_service.py index e28a2f3..b98f803 100644 --- a/backend/tests/test_services/test_ban_service.py +++ b/backend/tests/test_services/test_ban_service.py @@ -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 ) diff --git a/backend/tests/test_services/test_ban_service_perf.py b/backend/tests/test_services/test_ban_service_perf.py index 058c868..9a4df51 100644 --- a/backend/tests/test_services/test_ban_service_perf.py +++ b/backend/tests/test_services/test_ban_service_perf.py @@ -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", diff --git a/backend/tests/test_services/test_geo_service.py b/backend/tests/test_services/test_geo_service.py index 2684c5b..89994b3 100644 --- a/backend/tests/test_services/test_geo_service.py +++ b/backend/tests/test_services/test_geo_service.py @@ -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 diff --git a/backend/tests/test_tasks/test_geo_cache_flush.py b/backend/tests/test_tasks/test_geo_cache_flush.py index d0569ac..9fff3ce 100644 --- a/backend/tests/test_tasks/test_geo_cache_flush.py +++ b/backend/tests/test_tasks/test_geo_cache_flush.py @@ -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"] diff --git a/backend/tests/test_tasks/test_geo_re_resolve.py b/backend/tests/test_tasks/test_geo_re_resolve.py index 5f188df..0b2f942 100644 --- a/backend/tests/test_tasks/test_geo_re_resolve.py +++ b/backend/tests/test_tasks/test_geo_re_resolve.py @@ -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()