Repository: apify/crawlee-python
File: tests/unit/_autoscaling/test_autoscaled_pool.py (Python)

Code before:
        assert pool.desired_concurrency > 1

        # After 0.5s, the concurrency should reach max concurrency
        await asyncio.sleep(0.3)
        assert pool.desired_concurrency == 4

        # The concurrency should guarantee completion of more than 10 tasks (a single worker would complete ~5)
        assert done_count > 10

        # After 0.7s, the pretend overload should have kicked in and there should be a drop in desired concurrency
        await asyncio.sleep(0.2)
        assert pool.desired_concurrency < 4

        # After a full second, the pool should scale down all the way to 1
        await asyncio.sleep(0.3)
        assert pool.desired_concurrency == 1
    finally:
        pool_run_task.cancel()
        with suppress(asyncio.CancelledError):
            await pool_run_task


async def test_autoscales_uses_desired_concurrency_ratio(system_status: SystemStatus | Mock) -> None:
    """Test that desired concurrency ratio can limit desired concurrency.

    This test crates situation where only one task is ready and then no other task is ever ready.
    This creates situation where the system could scale up desired concurrency, but it will not do so because
    desired_concurrency_ratio=1 means that first the system would have to increase current concurrency to same number as
    desired concurrency and due to no other task ever being ready, it will never happen. Thus desired concurrency will
    stay 2 as was the initial setup, even though other conditions would allow the increase. (max_concurrency=4,
    system being idle)."""

    async def run() -> None:
        await asyncio.sleep(0.1)

    is_task_ready_iterator = chain([future(True)], repeat(future(False)))

    def is_task_ready_function() -> Awaitable[bool]:
        return next(is_task_ready_iterator)

    def get_historical_system_info() -> SystemInfo:
        return SystemInfo(
            cpu_info=LoadRatioInfo(limit_ratio=0.9, actual_ratio=0.3),
            memory_info=LoadRatioInfo(limit_ratio=0.9, actual_ratio=0.3),
            event_loop_info=LoadRatioInfo(limit_ratio=0.9, actual_ratio=0.3),
            client_info=LoadRatioInfo(limit_ratio=0.9, actual_ratio=0.3),
        )

    cast(Mock, system_status.get_historical_system_info).side_effect = get_historical_system_info

    pool = AutoscaledPool(

Review comment (suggestion):
```suggestion
    This test creates situation where only one task is ready and then no other task is ever ready.
```

Code after:
        assert pool.desired_concurrency > 1

        # After 0.5s, the concurrency should reach max concurrency
        await asyncio.sleep(0.3)
        assert pool.desired_concurrency == 4

        # The concurrency should guarantee completion of more than 10 tasks (a single worker would complete ~5)
        assert done_count > 10

        # After 0.7s, the pretend overload should have kicked in and there should be a drop in desired concurrency
        await asyncio.sleep(0.2)
        assert pool.desired_concurrency < 4

        # After a full second, the pool should scale down all the way to 1
        await asyncio.sleep(0.3)
        assert pool.desired_concurrency == 1
    finally:
        pool_run_task.cancel()
        with suppress(asyncio.CancelledError):
            await pool_run_task


async def test_autoscales_uses_desired_concurrency_ratio(system_status: SystemStatus | Mock) -> None:
    """Test that desired concurrency ratio can limit desired concurrency.

    This test creates situation where only one task is ready and then no other task is ever ready.
    This creates situation where the system could scale up desired concurrency, but it will not do so because
    desired_concurrency_ratio=1 means that first the system would have to increase current concurrency to same number as
    desired concurrency and due to no other task ever being ready, it will never happen. Thus desired concurrency will
    stay 2 as was the initial setup, even though other conditions would allow the increase. (max_concurrency=4,
    system being idle)."""

    async def run() -> None:
        await asyncio.sleep(0.1)

    is_task_ready_iterator = chain([future(True)], repeat(future(False)))

    def is_task_ready_function() -> Awaitable[bool]:
        return next(is_task_ready_iterator)

    def get_historical_system_info() -> SystemInfo:
        return SystemInfo(
            cpu_info=LoadRatioInfo(limit_ratio=0.9, actual_ratio=0.3),
            memory_info=LoadRatioInfo(limit_ratio=0.9, actual_ratio=0.3),
            event_loop_info=LoadRatioInfo(limit_ratio=0.9, actual_ratio=0.3),
            client_info=LoadRatioInfo(limit_ratio=0.9, actual_ratio=0.3),
        )

    cast(Mock, system_status.get_historical_system_info).side_effect = get_historical_system_info

    pool = AutoscaledPool(