test: add missing Crawlee and platform E2E tests

tests/e2e/actor_source_base/server.py

@@ -17,6 +17,9 @@
     /redirect-target - Page the redirect lands on
     /slow - Answers only after 10 minutes, to keep a request in flight while a run is interrupted
     /delayed/<n> - Answers after 5 seconds with a link to /delayed/<n+1>, so a request is always in flight
+    /files/<size>.<ext> - A binary file of <size> bytes, with the content type of <ext> (bin, pdf or png)
+    /start-urls.txt - A plain-text list of URLs, for the `requestsFromUrl` request list source. The URLs point to
+        example.com, since the URL extraction only matches hosts with a top-level domain.
 
 The homepage includes both direct product links (for Scrapy spiders that look for /products/ links
 on the start page) and category links (for testing crawl depth with Crawlee crawlers).
@@ -42,16 +45,31 @@
     '3': {'name': 'Widget C', 'price': '$39.99', 'description': 'A premium widget for professionals'},
 }
 
+_FILE_CONTENT_TYPES = {
+    'bin': 'application/octet-stream',
+    'pdf': 'application/pdf',
+    'png': 'image/png',
+}
 
-async def _send_html(send: Send, html: str, status: int = 200) -> None:
+
+def file_body(size: int) -> bytes:
+    """Return the body of the `/files/<size>.<ext>` route. Tests import it to compare the downloaded bytes."""
+    return bytes(i % 256 for i in range(size))
+
+
+async def _send_body(send: Send, body: bytes, content_type: str, status: int = 200) -> None:
     await send(
         {
             'type': 'http.response.start',
             'status': status,
-            'headers': [[b'content-type', b'text/html; charset=utf-8']],
+            'headers': [[b'content-type', content_type.encode()]],
         }
     )
-    await send({'type': 'http.response.body', 'body': html.encode()})
+    await send({'type': 'http.response.body', 'body': body})
+
+
+async def _send_html(send: Send, html: str, status: int = 200) -> None:
+    await _send_body(send, html.encode(), 'text/html; charset=utf-8', status)
 
 
 async def _send_redirect(send: Send, location: str) -> None:
@@ -149,6 +167,19 @@ async def app(scope: dict[str, Any], _receive: Receive, send: Send) -> None:
             f'<a href="/delayed/{n + 1}">Next</a>'
             f'</body></html>',
         )
+    elif path.startswith('/files/'):
+        size, _, extension = path.split('/')[-1].partition('.')
+        content_type = _FILE_CONTENT_TYPES.get(extension)
+        if not size.isdigit() or content_type is None:
+            await _send_html(send, '<html><body>Not Found</body></html>', 404)
+            return
+        await _send_body(send, file_body(int(size)), content_type)
+    elif path == '/start-urls.txt':
+        await _send_body(
+            send,
+            b'Products to scrape:\nhttps://example.com/products/2\nhttps://example.com/products/3\n',
+            'text/plain; charset=utf-8',
+        )
     elif path.startswith('/deep/'):
         try:
             n = int(path.split('/')[-1])

tests/e2e/test_actor_events.py

@@ -1,11 +1,15 @@
 from __future__ import annotations
 
 import asyncio
+from datetime import timedelta
 from typing import TYPE_CHECKING
 
+from .._utils import poll_until_condition
 from apify import Actor
 
 if TYPE_CHECKING:
+    from apify_client import ApifyClientAsync
+
     from .conftest import MakeActorFunction, RunActorFunction
 
 
@@ -167,3 +171,85 @@ async def handler(context: ParselCrawlingContext) -> None:
     assert run_log is not None
     assert 'error occurred during closing of asynchronous generator' not in run_log
     assert 'asynchronous generator is already running' not in run_log
+
+
+async def test_aborting_event_on_graceful_abort(
+    make_actor: MakeActorFunction,
+    apify_client_async: ApifyClientAsync,
+) -> None:
+    """A graceful abort from the API reaches the `ABORTING` listener before the platform stops the run."""
+
+    async def main() -> None:
+        from crawlee.events._types import Event
+
+        async def on_aborting() -> None:
+            await Actor.set_value('ABORTING_RECEIVED', {'received': True})
+
+        async with Actor:
+            Actor.on(Event.ABORTING, on_aborting)
+            Actor.log.info('Waiting for the abort')
+            await asyncio.sleep(300)
+
+    actor = await make_actor(label='actor-aborting-event', main_func=main)
+    run = await actor.start()
+    run_client = apify_client_async.run(run.id)
+
+    # Abort the run only once the listener is registered. The container startup time is highly variable, so poll the
+    # log with a growing interval.
+    log_client = run_client.log()
+    log = await poll_until_condition(
+        log_client.get,
+        lambda log: bool(log and 'Waiting for the abort' in log),
+        timeout=300,
+        poll_interval=2,
+        backoff_factor=1.2,
+    )
+    assert log is not None
+    assert 'Waiting for the abort' in log, f'The run did not start waiting in time:\n{log}'
+
+    await run_client.abort(gracefully=True)
+    run_result = await run_client.wait_for_finish(wait_duration=timedelta(seconds=600))
+    assert run_result is not None
+    assert run_result.status == 'ABORTED'
+
+    record = await run_client.key_value_store().get_record('ABORTING_RECEIVED')
+    assert record is not None, 'The ABORTING listener was not called.'
+    assert record['value'] == {'received': True}
+
+
+async def test_migrating_listeners_run_before_reboot(
+    make_actor: MakeActorFunction,
+    run_actor: RunActorFunction,
+) -> None:
+    """`Actor.reboot()` calls the `MIGRATING` listeners and the `PERSIST_STATE` listeners with `is_migrating=True`."""
+
+    async def main() -> None:
+        from crawlee.events._types import Event, EventPersistStateData
+
+        async def on_migrating() -> None:
+            await Actor.set_value('MIGRATING_RECEIVED', {'received': True})
+
+        async def on_persist_state(data: EventPersistStateData) -> None:
+            if data.is_migrating:
+                await Actor.set_value('PERSIST_STATE_MIGRATING', {'received': True})
+
+        async with Actor:
+            run = await Actor.apify_client.run(Actor.configuration.actor_run_id or '').get()
+            assert run is not None, 'The Actor could not read its own run.'
+
+            if run.stats.reboot_count == 0:
+                Actor.on(Event.MIGRATING, on_migrating)
+                Actor.on(Event.PERSIST_STATE, on_persist_state)
+                await Actor.reboot()
+
+    actor = await make_actor(label='actor-migrating-event', main_func=main)
+    run_result = await run_actor(actor)
+
+    assert run_result.status == 'SUCCEEDED'
+    assert run_result.stats.reboot_count == 1
+
+    kvs = actor.last_run().key_value_store()
+    for key in ('MIGRATING_RECEIVED', 'PERSIST_STATE_MIGRATING'):
+        record = await kvs.get_record(key)
+        assert record is not None, f'The listener did not write {key}.'
+        assert record['value'] == {'received': True}

tests/e2e/test_actor_key_value_store.py

@@ -140,3 +140,36 @@ async def main() -> None:
     run_result = await run_actor(actor)
 
     assert run_result.status == 'SUCCEEDED'
+
+
+async def test_use_state_survives_reboot(
+    make_actor: MakeActorFunction,
+    run_actor: RunActorFunction,
+) -> None:
+    """The state from `Actor.use_state()` is persisted before a reboot and loaded again in the next container."""
+
+    async def main() -> None:
+        async with Actor:
+            run = await Actor.apify_client.run(Actor.configuration.actor_run_id or '').get()
+            assert run is not None, 'The Actor could not read its own run.'
+
+            state = await Actor.use_state({'containers': 0})
+            containers = state['containers']
+            assert isinstance(containers, int), f'state={dict(state)}'
+            state['containers'] = containers + 1
+
+            if run.stats.reboot_count == 0:
+                state['before_reboot'] = ['a', 'b']
+                await Actor.reboot()
+            else:
+                assert state.get('before_reboot') == ['a', 'b'], f'state={dict(state)}'
+
+    actor = await make_actor(label='use-state-reboot', main_func=main)
+    run_result = await run_actor(actor)
+
+    assert run_result.status == 'SUCCEEDED'
+    assert run_result.stats.reboot_count == 1
+
+    state_record = await actor.last_run().key_value_store().get_record('APIFY_GLOBAL_STATE')
+    assert state_record is not None
+    assert state_record['value'] == {'containers': 2, 'before_reboot': ['a', 'b']}

tests/e2e/test_actor_request_list.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from .._utils import poll_until_condition
+from apify import Actor
+
+if TYPE_CHECKING:
+    from apify_client import ApifyClientAsync
+
+    from .conftest import MakeActorFunction, RunActorFunction
+
+_EXPECTED_REQUESTS = [
+    ('https://example.com/products/1', 'direct'),
+    ('https://example.com/products/2', 'remote'),
+    ('https://example.com/products/3', 'remote'),
+]
+
+
+async def test_request_list_from_actor_input(
+    make_actor: MakeActorFunction,
+    run_actor: RunActorFunction,
+    apify_client_async: ApifyClientAsync,
+) -> None:
+    """Direct URLs and a `requestsFromUrl` list from the Actor input reach the crawler with their user data."""
+
+    async def main() -> None:
+        from crawlee.crawlers import BasicCrawler, BasicCrawlingContext
+
+        from apify.request_loaders import ApifyRequestList
+
+        async with Actor:
+            actor_input = await Actor.get_input() or {}
+            request_list = await ApifyRequestList.open(
+                request_list_sources_input=actor_input.get('requestListSources', []),
+            )
+            crawler = BasicCrawler(request_manager=await request_list.to_tandem())
+
+            # The handler only records the requests. It doesn't fetch them, so the URLs don't have to resolve.
+            @crawler.router.default_handler
+            async def handler(context: BasicCrawlingContext) -> None:
+                await context.push_data({'url': context.request.url, 'source': context.request.user_data.get('source')})
+
+            await crawler.run()
+
+    actor = await make_actor(label='request-list', main_func=main)
+    run_result = await run_actor(
+        actor,
+        run_input={
+            'requestListSources': [
+                {'url': 'https://example.com/products/1', 'userData': {'source': 'direct'}},
+                {'requestsFromUrl': 'http://localhost:8080/start-urls.txt', 'userData': {'source': 'remote'}},
+            ],
+        },
+    )
+    assert run_result.status == 'SUCCEEDED'
+
+    dataset = actor.last_run().dataset()
+    items = await poll_until_condition(
+        dataset.list_items,
+        lambda page: len(page.items) >= len(_EXPECTED_REQUESTS),
+        timeout=30,
+    )
+    assert sorted((item['url'], item['source']) for item in items.items) == _EXPECTED_REQUESTS
+
+    # The tandem moves every request from the list into the default RQ and marks it as handled there.
+    assert run_result.default_request_queue_id is not None
+    rq_client = apify_client_async.request_queue(run_result.default_request_queue_id)
+    requests = await rq_client.list_requests()
+    handled_urls = sorted(request.url for request in requests.items if request.handled_at is not None)
+    assert handled_urls == [url for url, _ in _EXPECTED_REQUESTS]

tests/e2e/test_crawlee/actor_source/main_delayed_chain_crawler.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from crawlee import ConcurrencySettings
+from crawlee.crawlers import ParselCrawler, ParselCrawlingContext
+
+from apify import Actor
+
+if TYPE_CHECKING:
+    from crawlee.crawlers import BasicCrawlingContext
+
+IN_FLIGHT_LOG_MARKER = 'Sending the request to the server'
+"""Logged before every navigation, so a test knows which request is in flight."""
+
+
+async def main() -> None:
+    async with Actor:
+        # One request at a time, so the test knows what is in flight when it interrupts the run.
+        crawler = ParselCrawler(concurrency_settings=ConcurrencySettings(max_concurrency=1, desired_concurrency=1))
+
+        @crawler.pre_navigation_hook
+        async def mark_in_flight(context: BasicCrawlingContext) -> None:
+            Actor.log.info(f'{IN_FLIGHT_LOG_MARKER}: {context.request.url}')
+
+        @crawler.router.default_handler
+        async def handler(context: ParselCrawlingContext) -> None:
+            await context.push_data({'url': context.request.url})
+            await context.enqueue_links()
+
+        await crawler.run(['http://localhost:8080/delayed/1'])

tests/e2e/test_crawlee/actor_source/main_file_download_crawler.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+from crawlee.crawlers import FileDownloadCrawler, FileDownloadCrawlingContext
+from crawlee.storages import RequestQueue
+
+from apify import Actor
+
+_BUFFERED_URLS = ['http://localhost:8080/files/1024.pdf', 'http://localhost:8080/files/2048.png']
+_STREAMED_URLS = ['http://localhost:8080/files/5242880.bin']
+
+
+async def _store_file(context: FileDownloadCrawlingContext, content: bytes, chunk_count: int) -> None:
+    key = context.request.url.rsplit('/', 1)[-1]
+    kvs = await context.get_key_value_store()
+    await kvs.set_value(key, content, content_type=context.http_response.headers.get('content-type'))
+    await context.push_data({'url': context.request.url, 'key': key, 'chunk_count': chunk_count})
+
+
+async def main() -> None:
+    async with Actor:
+        buffered_crawler = FileDownloadCrawler()
+
+        @buffered_crawler.router.default_handler
+        async def buffered_handler(context: FileDownloadCrawlingContext) -> None:
+            content = await context.http_response.read()
+            await _store_file(context, content, chunk_count=1)
+
+        await buffered_crawler.run(_BUFFERED_URLS)
+
+        # The default RQ already holds the buffered crawl, so the streamed crawl gets a queue of its own.
+        streamed_crawler = FileDownloadCrawler(stream=True, request_manager=await RequestQueue.open(alias='streamed'))
+
+        @streamed_crawler.router.default_handler
+        async def streamed_handler(context: FileDownloadCrawlingContext) -> None:
+            chunks = [chunk async for chunk in context.http_response.read_stream()]
+            await _store_file(context, b''.join(chunks), chunk_count=len(chunks))
+
+        await streamed_crawler.run(_STREAMED_URLS)

tests/e2e/test_crawlee/actor_source/main_reboot_crawler.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+from crawlee import ConcurrencySettings
+from crawlee.crawlers import ParselCrawler, ParselCrawlingContext
+
+from apify import Actor
+
+REBOOT_URL = 'http://localhost:8080/about'
+"""The page whose handler reboots the first container, so part of the site is crawled before and part after."""
+
+
+async def main() -> None:
+    async with Actor:
+        run = await Actor.apify_client.run(Actor.configuration.actor_run_id or '').get()
+        assert run is not None, 'The Actor could not read its own run.'
+        is_first_container = run.stats.reboot_count == 0
+
+        # One request at a time, so no other request is in progress when the handler reboots the run.
+        crawler = ParselCrawler(
+            max_crawl_depth=2,
+            concurrency_settings=ConcurrencySettings(max_concurrency=1, desired_concurrency=1),
+        )
+
+        @crawler.router.default_handler
+        async def handler(context: ParselCrawlingContext) -> None:
+            if context.request.url == REBOOT_URL and is_first_container:
+                request_manager = await crawler.get_request_manager()
+                await request_manager.reclaim_request(context.request)
+                await Actor.reboot()
+
+            await context.enqueue_links()
+
+            if '/products/' in context.request.url:
+                await context.push_data(
+                    {
+                        'url': context.request.url,
+                        'name': context.selector.css('h1::text').get('').strip(),
+                    }
+                )
+
+        await crawler.run(['http://localhost:8080/'])
+
+        await Actor.set_value(
+            'CRAWLER_RESULT',
+            {'requests_finished': crawler.statistics.state.requests_finished},
+        )