diff --git a/src/crawlee/request_loaders/_sitemap_request_loader.py b/src/crawlee/request_loaders/_sitemap_request_loader.py index f0a6af0647..bc39211058 100644 --- a/src/crawlee/request_loaders/_sitemap_request_loader.py +++ b/src/crawlee/request_loaders/_sitemap_request_loader.py @@ -302,6 +302,12 @@ async def abort_loading(self) -> None: with suppress(asyncio.CancelledError): await self._loading_task + async def persist_state(self) -> None: + """Persist the loader's state to the key-value store. Does nothing without a `persist_state_key`.""" + # Loads the persisted record first, so that a call before the first use writes the restored state. + await self._get_state() + await self._state.persist_state() + async def close(self) -> None: """Close the request loader.""" self._closed = True diff --git a/tests/unit/request_loaders/test_sitemap_request_loader.py b/tests/unit/request_loaders/test_sitemap_request_loader.py index 48596f8a12..39bf5d6ad6 100644 --- a/tests/unit/request_loaders/test_sitemap_request_loader.py +++ b/tests/unit/request_loaders/test_sitemap_request_loader.py @@ -234,3 +234,19 @@ async def test_close_before_first_use_stops_listening(monkeypatch: pytest.Monkey await event_manager.wait_for_all_listeners_to_complete() set_value.assert_not_called() + + +async def test_persist_state_writes_record(key_value_store: KeyValueStore) -> None: + """`persist_state()` writes the state with the in-progress request to the key-value store.""" + sitemap_url = f'{DEFAULT_URL}sitemap.xml' + client, _ = make_status_stream_client({sitemap_url: [(200, get_basic_sitemap().encode())]}) + loader = SitemapRequestLoader([sitemap_url], http_client=client, persist_state_key='sitemap-state') + request = await loader.fetch_next_request() + assert request is not None + + await loader.persist_state() + record = await key_value_store.get_value('sitemap-state') + await loader.close() + + assert record['inProgress'] == [request.url] + assert list(record['requestData']) == [request.url]