Repository: apify/crawlee-python
File: tests/unit/basic_crawler/test_basic_crawler.py (Python)

Code before:

    stats = await crawler.run(['http://example.tld/1'])
    assert stats.requests_failed == 1


async def test_context_update_kv_store() -> None:
    crawler = BasicCrawler()

    @crawler.router.default_handler
    async def handler(context: BasicCrawlingContext) -> None:
        store = await context.get_key_value_store()
        await store.set_value('foo', 'bar')

    await crawler.run(['https://hello.world'])

    store = await crawler.get_key_value_store()
    assert (await store.get_value('foo')) == 'bar'


async def test_context_use_state(mock_event_manager: EventManager) -> None:
    crawler = BasicCrawler()

    @crawler.router.default_handler
    async def handler(context: BasicCrawlingContext) -> None:
        await context.use_state('state', {'hello': 'world'})
        await asyncio.sleep(0.1)

    await crawler.run(['https://hello.world'])

    store = await crawler.get_key_value_store()
    assert (await store.get_value('state')) == {'hello': 'world'}


async def test_max_requests_per_crawl(httpbin: str) -> None:
    start_urls = [f'{httpbin}/1', f'{httpbin}/2', f'{httpbin}/3', f'{httpbin}/4', f'{httpbin}/5']
    processed_urls = []

    # Set max_concurrency to 1 to ensure testing max_requests_per_crawl accurately
    crawler = BasicCrawler(
        concurrency_settings=ConcurrencySettings(max_concurrency=1),
        max_requests_per_crawl=3,
    )

    @crawler.router.default_handler
    async def handler(context: BasicCrawlingContext) -> None:
        processed_urls.append(context.request.url)

    stats = await crawler.run(start_urls)

    # Verify that only 3 out of the 5 provided URLs were made
    assert len(processed_urls) == 3

Review comment (suggestion):
Could the sleep go after `crawler.run` instead? I assume it's to ensure key-value store consistency.

Code after:
async def test_crawler_push_data_over_limit() -> None:
    crawler = BasicCrawler()

    @crawler.router.default_handler
    async def handler(context: BasicCrawlingContext) -> None:
        # Push a roughly 15MB payload - this should be enough to break the 9MB limit
        await context.push_data({'hello': 'world' * 3 * 1024 * 1024})

    stats = await crawler.run(['http://example.tld/1'])
    assert stats.requests_failed == 1


async def test_context_update_kv_store() -> None:
    crawler = BasicCrawler()

    @crawler.router.default_handler
    async def handler(context: BasicCrawlingContext) -> None:
        store = await context.get_key_value_store()
        await store.set_value('foo', 'bar')

    await crawler.run(['https://hello.world'])

    store = await crawler.get_key_value_store()
    assert (await store.get_value('foo')) == 'bar'


async def test_context_use_state(key_value_store: KeyValueStore) -> None:
    crawler = BasicCrawler()

    @crawler.router.default_handler
    async def handler(context: BasicCrawlingContext) -> None:
        await context.use_state('state', {'hello': 'world'})

    await crawler.run(['https://hello.world'])

    store = await crawler.get_key_value_store()

    assert (await store.get_value('state')) == {'hello': 'world'}


async def test_context_handlers_use_state(key_value_store: KeyValueStore) -> None:
    state_in_handler_one: dict[str, JsonSerializable] = {}
    state_in_handler_two: dict[str, JsonSerializable] = {}
    state_in_handler_three: dict[str, JsonSerializable] = {}

    crawler = BasicCrawler()

    @crawler.router.handler('one')
    async def handler_one(context: BasicCrawlingContext) -> None:
        state = await context.use_state('state', {'hello': 'world'})
        state_in_handler_one.update(state)