import asyncio
import logging
import time
from typing import List

import pytest

import ray
from ray import serve
from ray._common.utils import get_or_create_event_loop
from ray.serve._private.common import DeploymentID, ReplicaID
from ray.serve._private.config import DeploymentConfig
from ray.serve._private.constants import SERVE_LOGGER_NAME
from ray.serve.batching import _BatchQueue
from ray.serve.exceptions import RayServeException

# Setup the global replica context for the test.
default_deployment_config = DeploymentConfig()
ray.serve.context._set_internal_replica_context(
    replica_id=ReplicaID(unique_id="test", deployment_id=DeploymentID(name="test")),
    servable_object=None,
    _deployment_config=default_deployment_config,
    rank=0,
    world_size=1,
)


class FakeStream:
    def __init__(self):
        self.messages = []

    def write(self, buf):
        self.messages.append(buf)

    def reset_message(self):
        self.messages = []


# We use a single event loop for the entire test session. Without this
# fixture, the event loop is sometimes prematurely terminated by pytest.
@pytest.fixture(scope="session")
def event_loop():
    loop = get_or_create_event_loop()
    yield loop
    loop.close()


@pytest.mark.asyncio
async def test_decorator_validation():
    @serve.batch
    async def function():
        pass

    @serve.batch(max_batch_size=10, batch_wait_timeout_s=1.5)
    async def function2():
        pass

    class Class:
        @serve.batch
        async def method(self):
            pass

    class Class2:
        @serve.batch(max_batch_size=10, batch_wait_timeout_s=1.5)
        async def method(self):
            pass

    with pytest.raises(TypeError, match="async def"):

        @serve.batch
        def non_async_function():
            pass

    with pytest.raises(TypeError, match="async def"):

        class NotAsync:
            @serve.batch
            def method(self, requests):
                pass

    with pytest.raises(ValueError):

        class ZeroBatch:
            @serve.batch(max_batch_size=0)
            async def method(self, requests):
                pass

    with pytest.raises(TypeError):

        class FloatNonIntBatch:
            @serve.batch(max_batch_size=1.1)
            async def method(self, requests):
                pass

    class FloatIntegerBatch:
        @serve.batch(max_batch_size=1.0)
        async def method(self, requests):
            pass

    with pytest.raises(ValueError):

        class NegativeTimeout:
            @serve.batch(batch_wait_timeout_s=-0.1)
            async def method(self, requests):
                pass

    class FloatZeroTimeout:
        @serve.batch(batch_wait_timeout_s=0.0)
        async def method(self, requests):
            pass

    class IntZeroTimeout:
        @serve.batch(batch_wait_timeout_s=0)
        async def method(self, requests):
            pass

    with pytest.raises(TypeError):

        class NonTimeout:
            @serve.batch(batch_wait_timeout_s="a")
            async def method(self, requests):
                pass


@pytest.mark.asyncio
@pytest.mark.parametrize("use_class", [True, False])
async def test_batch_size_one_long_timeout(use_class):
    @serve.batch(max_batch_size=1, batch_wait_timeout_s=1000)
    async def long_timeout(requests):
        if "raise" in requests:
            _ = 1 / 0
        return requests

    class LongTimeout:
        @serve.batch(max_batch_size=1, batch_wait_timeout_s=1000)
        async def long_timeout(self, requests):
            if "raise" in requests:
                _ = 1 / 0
            return requests

    cls = LongTimeout()

    async def call(arg):
        if use_class:
            return await cls.long_timeout(arg)
        else:
            return await long_timeout(arg)

    assert await call("hi") == "hi"
    with pytest.raises(ZeroDivisionError):
        await call("raise")


@pytest.mark.asyncio
@pytest.mark.parametrize("use_class", [True, False])
async def test_batch_size_multiple_zero_timeout(use_class):
    block_execution_event = asyncio.Event()

    @serve.batch(max_batch_size=2, batch_wait_timeout_s=0)
    async def zero_timeout(requests):
        await block_execution_event.wait()
        if "raise" in requests:
            _ = 1 / 0
        return requests

    class ZeroTimeout:
        @serve.batch(max_batch_size=2, batch_wait_timeout_s=0)
        async def zero_timeout(self, requests):
            await block_execution_event.wait()
            if "raise" in requests:
                _ = 1 / 0
            return requests

    cls = ZeroTimeout()

    async def call(arg):
        if use_class:
            return await cls.zero_timeout(arg)
        else:
            return await zero_timeout(arg)

    block_execution_event.set()
    assert await call("hi") == "hi"
    with pytest.raises(ZeroDivisionError):
        await call("raise")
    block_execution_event.clear()

    # Check that 2 requests will be executed together if available.
    # The first should cause a size-one batch to be executed, then
    # the next two should be executed together (signaled by both
    # having the exception).
    t1 = get_or_create_event_loop().create_task(call("hi1"))
    with pytest.raises(asyncio.TimeoutError):
        await asyncio.wait_for(asyncio.shield(t1), timeout=0.001)

    t2 = get_or_create_event_loop().create_task(call("hi2"))
    t3 = get_or_create_event_loop().create_task(call("raise"))

    block_execution_event.set()
    assert await t1 == "hi1"

    with pytest.raises(ZeroDivisionError):
        await t2
    with pytest.raises(ZeroDivisionError):
        await t3


@pytest.mark.asyncio
async def test_batch_timeout_empty_queue():
    """Check that Serve waits when creating batches.

    Serve should wait a full batch_wait_timeout_s after receiving the first
    request in the next batch before processing the batch.
    """

    @serve.batch(max_batch_size=10, batch_wait_timeout_s=0.25)
    async def no_op(requests):
        return ["No-op"] * len(requests)

    num_iterations = 2
    for iteration in range(num_iterations):
        tasks = [get_or_create_event_loop().create_task(no_op(None)) for _ in range(9)]
        done, _ = await asyncio.wait(tasks, timeout=0.05)

        # Due to the long timeout, none of the tasks should finish until a tenth
        # request is submitted
        assert len(done) == 0

        tasks.append(get_or_create_event_loop().create_task(no_op(None)))
        done, _ = await asyncio.wait(tasks, timeout=0.05)

        # All the timeout tasks should be finished
        assert set(tasks) == set(done)
        assert all(t.result() == "No-op" for t in tasks)

        if iteration < num_iterations - 1:
            # Leave queue empty for batch_wait_timeout_s between batches
            time.sleep(0.25)


@pytest.mark.asyncio
async def test_batch_wait_queue_exceeds_batch_size_race_condition():
    """Check that the wait queue can exceed the batch size.

    This test was added to guard against a race condition documented in
    https://github.com/ray-project/ray/pull/42705#discussion_r1466653910.
    """

    @serve.batch(max_batch_size=2, batch_wait_timeout_s=10000)
    async def no_op(requests):
        return ["No-op"] * len(requests)

    tasks = [get_or_create_event_loop().create_task(no_op(None)) for _ in range(10)]

    # All the tasks should finish.
    done, pending = await asyncio.wait(tasks, timeout=0.5)
    assert len(done) == len(tasks)
    assert len(pending) == 0


@pytest.mark.asyncio
@pytest.mark.parametrize("use_class", [True, False])
async def test_batch_size_multiple_long_timeout(use_class):
    @serve.batch(max_batch_size=3, batch_wait_timeout_s=1000)
    async def long_timeout(requests):
        if "raise" in requests:
            _ = 1 / 0
        return requests

    class LongTimeout:
        @serve.batch(max_batch_size=3, batch_wait_timeout_s=1000)
        async def long_timeout(self, requests):
            if "raise" in requests:
                _ = 1 / 0
            return requests

    cls = LongTimeout()

    async def call(arg):
        if use_class:
            return await cls.long_timeout(arg)
        else:
            return await long_timeout(arg)

    t1 = get_or_create_event_loop().create_task(call("hi1"))
    t2 = get_or_create_event_loop().create_task(call("hi2"))
    done, pending = await asyncio.wait([t1, t2], timeout=0.1)
    assert len(done) == 0
    t3 = get_or_create_event_loop().create_task(call("hi3"))
    done, pending = await asyncio.wait([t1, t2, t3], timeout=100)
    assert set(done) == {t1, t2, t3}
    assert [t1.result(), t2.result(), t3.result()] == ["hi1", "hi2", "hi3"]

    t1 = get_or_create_event_loop().create_task(call("hi1"))
    t2 = get_or_create_event_loop().create_task(call("raise"))
    done, pending = await asyncio.wait([t1, t2], timeout=0.1)
    assert len(done) == 0
    t3 = get_or_create_event_loop().create_task(call("hi3"))
    done, pending = await asyncio.wait([t1, t2, t3], timeout=100)
    assert set(done) == {t1, t2, t3}
    assert all(isinstance(t.exception(), ZeroDivisionError) for t in done)
    with pytest.raises(ZeroDivisionError):
        t1.result()
    with pytest.raises(ZeroDivisionError):
        t2.result()
    with pytest.raises(ZeroDivisionError):
        t3.result()


@pytest.mark.asyncio
@pytest.mark.parametrize("mode", ["args", "kwargs", "mixed", "out-of-order"])
@pytest.mark.parametrize("use_class", [True, False])
async def test_batch_args_kwargs(mode, use_class):
    if use_class:

        class MultipleArgs:
            @serve.batch(max_batch_size=2, batch_wait_timeout_s=1000)
            async def method(self, key1, key2):
                return [(key1[i], key2[i]) for i in range(len(key1))]

        instance = MultipleArgs()
        func = instance.method

    else:

        @serve.batch(max_batch_size=2, batch_wait_timeout_s=1000)
        async def func(key1, key2):
            return [(key1[i], key2[i]) for i in range(len(key1))]

    if mode == "args":
        coros = [func("hi1", "hi2"), func("hi3", "hi4")]
    elif mode == "kwargs":
        coros = [func(key1="hi1", key2="hi2"), func(key1="hi3", key2="hi4")]
    elif mode == "mixed":
        coros = [func("hi1", key2="hi2"), func("hi3", key2="hi4")]
    elif mode == "out-of-order":
        coros = [func(key2="hi2", key1="hi1"), func(key2="hi4", key1="hi3")]

    result = await asyncio.gather(*coros)
    assert result == [("hi1", "hi2"), ("hi3", "hi4")]


@pytest.mark.asyncio
@pytest.mark.parametrize("use_class", [True, False])
@pytest.mark.parametrize("use_gen", [True, False])
async def test_batch_cancellation(use_class, use_gen):
    block_requests = asyncio.Event()
    request_was_cancelled = True

    async def unary_implementation(key1, key2):
        nonlocal block_requests, request_was_cancelled
        await block_requests.wait()
        request_was_cancelled = False
        return [(key1[i], key2[i]) for i in range(len(key1))]

    async def streaming_implementation(key1, key2):
        nonlocal block_requests, request_was_cancelled
        await block_requests.wait()
        request_was_cancelled = False
        yield [(key1[i], key2[i]) for i in range(len(key1))]

    if use_class:

        class MultipleArgs:
            @serve.batch(max_batch_size=2, batch_wait_timeout_s=1000)
            async def unary_method(self, key1, key2):
                return await unary_implementation(key1, key2)

            @serve.batch(max_batch_size=2, batch_wait_timeout_s=1000)
            async def streaming_method(self, key1, key2):
                async for value in streaming_implementation(key1, key2):
                    yield value

        instance = MultipleArgs()

        if use_gen:
            func = instance.streaming_method
        else:
            func = instance.unary_method

    else:

        @serve.batch(max_batch_size=2, batch_wait_timeout_s=1000)
        async def unary_func(key1, key2):
            return await unary_implementation(key1, key2)

        @serve.batch(max_batch_size=2, batch_wait_timeout_s=1000)
        async def streaming_func(key1, key2):
            async for value in streaming_implementation(key1, key2):
                yield value

        if use_gen:
            func = streaming_func
        else:
            func = unary_func

    if use_gen:
        gens = [func("hi1", "hi2"), func("hi3", "hi4")]
        tasks = [asyncio.create_task(gen.__anext__()) for gen in gens]
    else:
        tasks = [
            asyncio.create_task(func("hi1", "hi2")),
            asyncio.create_task(func("hi3", "hi4")),
        ]

    print("Submitted requests.")

    # The requests should be blocked on the long request_timeout
    done, pending = await asyncio.wait(tasks, timeout=0.01)
    assert len(done) == 0
    assert len(pending) == 2

    print("Requests are blocked, as expected.")

    # Cancel the first request. The second request should still be blocked on
    # the long request_timeout
    tasks[0].cancel()
    pending, done = await asyncio.wait(tasks, timeout=0.01)
    assert len(done) == 1
    assert len(pending) == 1

    print("Cancelled first request.")

    # Cancel the second request. Both requests should be done.
    tasks[1].cancel()
    done, pending = await asyncio.wait(tasks, timeout=0.01)
    assert len(done) == 2
    assert len(pending) == 0

    print("Cancelled second request. Sending new requests with no timeout.")

    # Sanity check that the request was actually cancelled.
    assert request_was_cancelled

    # Unblock requests. The requests should succeed.
    block_requests.set()

    if use_gen:
        gens = [func("hi1", "hi2"), func("hi3", "hi4")]
        tasks = [asyncio.create_task(gen.__anext__()) for gen in gens]
    else:
        tasks = [
            asyncio.create_task(func("hi1", "hi2")),
            asyncio.create_task(func("hi3", "hi4")),
        ]

    result = await asyncio.gather(*tasks)
    assert result == [("hi1", "hi2"), ("hi3", "hi4")]


@pytest.mark.asyncio
@pytest.mark.parametrize("use_class", [True, False])
@pytest.mark.parametrize("use_gen", [True, False])
async def test_cancellation_after_error(use_class, use_gen):
    """Cancelling a request after it errors should be supported."""

    raise_error = asyncio.Event()

    async def unary_implementation(key1, key2):
        if not raise_error.is_set():
            raise ValueError()
        return [(key1[i], key2[i]) for i in range(len(key1))]

    async def streaming_implementation(key1, key2):
        if not raise_error.is_set():
            raise ValueError()
        yield [(key1[i], key2[i]) for i in range(len(key1))]

    if use_class:

        class MultipleArgs:
            @serve.batch(max_batch_size=2, batch_wait_timeout_s=1000)
            async def unary_method(self, key1, key2):
                return await unary_implementation(key1, key2)

            @serve.batch(max_batch_size=2, batch_wait_timeout_s=1000)
            async def streaming_method(self, key1, key2):
                async for value in streaming_implementation(key1, key2):
                    yield value

        instance = MultipleArgs()

        if use_gen:
            func = instance.streaming_method
        else:
            func = instance.unary_method

    else:

        @serve.batch(max_batch_size=2, batch_wait_timeout_s=1000)
        async def unary_func(key1, key2):
            return await unary_implementation(key1, key2)

        @serve.batch(max_batch_size=2, batch_wait_timeout_s=1000)
        async def streaming_func(key1, key2):
            async for value in streaming_implementation(key1, key2):
                yield value

        if use_gen:
            func = streaming_func
        else:
            func = unary_func

    # Submit requests and then cancel them.
    if use_gen:
        gens = [func("hi1", "hi2"), func("hi3", "hi4")]
        tasks = [asyncio.create_task(gen.__anext__()) for gen in gens]
    else:
        tasks = [
            asyncio.create_task(func("hi1", "hi2")),
            asyncio.create_task(func("hi3", "hi4")),
        ]

    print("Submitted initial batch of requests.")

    for task in tasks:
        task.cancel()

    print("Closed initial batch of requests.")

    raise_error.set()

    # Submit requests and check that they still work.
    if use_gen:
        gens = [func("hi1", "hi2"), func("hi3", "hi4")]
        tasks = [asyncio.create_task(gen.__anext__()) for gen in gens]
    else:
        tasks = [
            asyncio.create_task(func("hi1", "hi2")),
            asyncio.create_task(func("hi3", "hi4")),
        ]

    print("Submitted new batch of requests.")

    result = await asyncio.gather(*tasks)
    assert result == [("hi1", "hi2"), ("hi3", "hi4")]


@pytest.mark.asyncio
@pytest.mark.parametrize("use_class", [True, False])
async def test_batch_setters(use_class):
    if use_class:

        class C:
            @serve.batch(max_batch_size=2, batch_wait_timeout_s=1000)
            async def method(self, key1, key2):
                return [(key1[i], key2[i]) for i in range(len(key1))]

        instance = C()
        func = instance.method

    else:

        @serve.batch(max_batch_size=2, batch_wait_timeout_s=1000)
        async def func(key1, key2):
            return [(key1[i], key2[i]) for i in range(len(key1))]

    assert func._get_max_batch_size() == 2
    assert func._get_batch_wait_timeout_s() == 1000

    # @serve.batch should create batches of size 2
    tasks = [
        get_or_create_event_loop().create_task(func("hi1", "hi2")),
        get_or_create_event_loop().create_task(func("hi3", "hi4")),
    ]
    done, pending = await asyncio.wait(tasks, timeout=0.1)
    assert len(pending) == 0
    assert {task.result() for task in done} == {("hi1", "hi2"), ("hi3", "hi4")}

    # Set new values
    func.set_max_batch_size(3)
    func.set_batch_wait_timeout_s(15000)

    assert func._get_max_batch_size() == 3
    assert func._get_batch_wait_timeout_s() == 15000

    # @serve.batch should create batches of size 3
    tasks = [
        get_or_create_event_loop().create_task(func("hi1", "hi2")),
        get_or_create_event_loop().create_task(func("hi3", "hi4")),
        get_or_create_event_loop().create_task(func("hi5", "hi6")),
    ]
    done, pending = await asyncio.wait(tasks, timeout=0.1)
    assert len(pending) == 0
    assert {task.result() for task in done} == {
        ("hi1", "hi2"),
        ("hi3", "hi4"),
        ("hi5", "hi6"),
    }


@pytest.mark.asyncio
async def test_batch_use_earliest_setters():
    """@serve.batch should use the right settings when constructing a batch.

    When the @serve.batch setters get called before a batch has started
    accumulating, the next batch should use the setters' values. When they
    get called while a batch is accumulating, the previous values should be
    used.
    """

    @serve.batch(max_batch_size=2, batch_wait_timeout_s=1000)
    async def func(key1, key2):
        return [(key1[i], key2[i]) for i in range(len(key1))]

    assert func._get_max_batch_size() == 2
    assert func._get_batch_wait_timeout_s() == 1000

    # Set new values
    func.set_max_batch_size(3)
    func.set_batch_wait_timeout_s(15000)

    assert func._get_max_batch_size() == 3
    assert func._get_batch_wait_timeout_s() == 15000

    # Should create batches of size 3, even if setters are called while
    # batch is accumulated
    tasks = [
        get_or_create_event_loop().create_task(func("hi1", "hi2")),
        get_or_create_event_loop().create_task(func("hi3", "hi4")),
    ]

    # Batch should be waiting for last request
    done, pending = await asyncio.wait(tasks, timeout=0.1)
    assert len(done) == 0 and len(pending) == 2

    func.set_max_batch_size(1)
    func.set_batch_wait_timeout_s(0)

    assert func._get_max_batch_size() == 1
    assert func._get_batch_wait_timeout_s() == 0

    # Batch should still be waiting for last request
    done, pending = await asyncio.wait(pending, timeout=0.1)
    assert len(done) == 0 and len(pending) == 2

    # Batch should execute after last request
    pending.add(get_or_create_event_loop().create_task(func("hi5", "hi6")))
    done, pending = await asyncio.wait(pending, timeout=0.1)
    assert len(pending) == 0
    assert {task.result() for task in done} == {
        ("hi1", "hi2"),
        ("hi3", "hi4"),
        ("hi5", "hi6"),
    }

    # Next batch should use updated values
    tasks = [get_or_create_event_loop().create_task(func("hi1", "hi2"))]
    done, pending = await asyncio.wait(tasks, timeout=0.1)
    assert len(done) == 1 and len(pending) == 0
    assert done.pop().result() == ("hi1", "hi2")


@pytest.mark.asyncio
@pytest.mark.parametrize("mode", ["args", "kwargs", "mixed", "out-of-order"])
@pytest.mark.parametrize("use_class", [True, False])
@pytest.mark.parametrize("generator_length", [0, 2, 5])
async def test_batch_generator_basic(mode, use_class, generator_length):
    if use_class:

        class MultipleArgs:
            @serve.batch(max_batch_size=2, batch_wait_timeout_s=1000)
            async def method(self, key1, key2):
                for gen_idx in range(generator_length):
                    yield [(gen_idx, key1[i], key2[i]) for i in range(len(key1))]

        instance = MultipleArgs()
        func = instance.method

    else:

        @serve.batch(max_batch_size=2, batch_wait_timeout_s=1000)
        async def func(key1, key2):
            for gen_idx in range(generator_length):
                yield [(gen_idx, key1[i], key2[i]) for i in range(len(key1))]

    if mode == "args":
        generators = [func("hi1", "hi2"), func("hi3", "hi4")]
    elif mode == "kwargs":
        generators = [func(key1="hi1", key2="hi2"), func(key1="hi3", key2="hi4")]
    elif mode == "mixed":
        generators = [func("hi1", key2="hi2"), func("hi3", key2="hi4")]
    elif mode == "out-of-order":
        generators = [func(key2="hi2", key1="hi1"), func(key2="hi4", key1="hi3")]

    results = [
        [result async for result in generators[0]],
        [result async for result in generators[1]],
    ]

    assert results == [
        [(gen_idx, "hi1", "hi2") for gen_idx in range(generator_length)],
        [(gen_idx, "hi3", "hi4") for gen_idx in range(generator_length)],
    ]


@pytest.mark.asyncio
@pytest.mark.parametrize("error_type", ["runtime_error", "mismatched_lengths"])
async def test_batch_generator_exceptions(error_type):
    GENERATOR_LENGTH = 5
    ERROR_IDX = 2
    ERROR_MSG = "Testing error"

    @serve.batch(max_batch_size=2, batch_wait_timeout_s=1000)
    async def func(key1, key2):
        for gen_idx in range(GENERATOR_LENGTH):
            results = [(gen_idx, key1[i], key2[i]) for i in range(len(key1))]
            if gen_idx == ERROR_IDX:
                if error_type == "runtime_error":
                    raise RuntimeError(ERROR_MSG)
                elif error_type == "mismatched_lengths":
                    yield results * 2
            yield results

    generators = [func("hi1", "hi2"), func("hi3", "hi4")]

    for generator in generators:
        for _ in range(ERROR_IDX):
            await generator.__anext__()

        if error_type == "runtime_error":
            with pytest.raises(RuntimeError, match=ERROR_MSG):
                await generator.__anext__()
        elif error_type == "mismatched_lengths":
            with pytest.raises(RayServeException):
                await generator.__anext__()

        with pytest.raises(StopAsyncIteration):
            await generator.__anext__()


@pytest.mark.asyncio
@pytest.mark.parametrize("stop_token", [StopAsyncIteration, StopIteration])
async def test_batch_generator_early_termination(stop_token):
    NUM_CALLERS = 4

    event = asyncio.Event()

    @serve.batch(max_batch_size=NUM_CALLERS, batch_wait_timeout_s=1000)
    async def sequential_terminator(ids: List[int]):
        """Terminates callers one-after-another in order of call."""

        for num_finished_callers in range(1, NUM_CALLERS + 1):
            event.clear()
            responses = [stop_token for _ in range(num_finished_callers)]
            responses += [ids[idx] for idx in range(num_finished_callers, NUM_CALLERS)]
            yield responses
            await event.wait()

    ids = list(range(NUM_CALLERS))
    generators = [sequential_terminator(id) for id in ids]
    for id, generator in zip(ids, generators):
        async for result in generator:
            assert result == id

        # Each terminated caller frees the sequential_terminator to process
        # another iteration.
        event.set()


@pytest.mark.asyncio
async def test_batch_generator_setters():
    """@serve.batch setters should succeed while the current batch streams."""

    @serve.batch(max_batch_size=2, batch_wait_timeout_s=1000)
    async def yield_three_times(key1, key2):
        for _ in range(3):
            yield [(key1[i], key2[i]) for i in range(len(key1))]

    assert yield_three_times._get_max_batch_size() == 2
    assert yield_three_times._get_batch_wait_timeout_s() == 1000

    args_list = [("hi1", "hi2"), ("hi3", "hi4")]
    coros = [yield_three_times(*args) for args in args_list]

    # Partially consume generators
    for coro, expected_result in zip(coros, args_list):
        for _ in range(2):
            await coro.__anext__() == expected_result

    # Set new values
    yield_three_times.set_max_batch_size(3)
    yield_three_times.set_batch_wait_timeout_s(15000)

    assert yield_three_times._get_max_batch_size() == 3
    assert yield_three_times._get_batch_wait_timeout_s() == 15000

    # Execute three more requests
    args_list_2 = [("hi1", "hi2"), ("hi3", "hi4"), ("hi5", "hi6")]
    coros_2 = [yield_three_times(*args) for args in args_list_2]

    # Finish consuming original requests
    for coro, expected_result in zip(coros, args_list):
        await coro.__anext__() == expected_result
        with pytest.raises(StopAsyncIteration):
            await coro.__anext__()

    # Consume new requests
    for coro, expected_result in zip(coros_2, args_list_2):
        for _ in range(3):
            await coro.__anext__() == expected_result
        with pytest.raises(StopAsyncIteration):
            await coro.__anext__()


@pytest.mark.asyncio
async def test_batch_size_fn_deferred_item_early_break():
    batches_processed = []

    @serve.batch(
        max_batch_size=10,
        batch_wait_timeout_s=0.05,
        batch_size_fn=lambda items: sum(item["size"] for item in items),
    )
    async def batch_handler(requests):
        batches_processed.append([req["value"] for req in requests])
        return [req["value"] for req in requests]

    # Request 1: size=6 (fits in batch)
    # Request 2: size=6 (would make total 12 > 10, should be deferred)
    # Each should be processed in its own batch
    t1 = get_or_create_event_loop().create_task(
        batch_handler({"size": 6, "value": "first"})
    )
    t2 = get_or_create_event_loop().create_task(
        batch_handler({"size": 6, "value": "second"})
    )

    done, pending = await asyncio.wait([t1, t2], timeout=1.0)

    assert len(done) == 2, "Both tasks should complete"
    assert len(pending) == 0

    results = {t1.result(), t2.result()}
    assert results == {"first", "second"}

    # Verify they were processed in separate batches due to size constraint
    assert (
        len(batches_processed) == 2
    ), f"Expected 2 separate batches, got {batches_processed}"


@pytest.mark.asyncio
async def test_batch_size_fn_fail_to_fit():
    batches_processed = []

    @serve.batch(
        max_batch_size=10,
        batch_wait_timeout_s=0.05,
        batch_size_fn=lambda items: sum(item["size"] for item in items),
    )
    async def batch_handler(requests):
        batches_processed.append([req["value"] for req in requests])
        return [req["value"] for req in requests]

    # Request 1: size=6 (fits in batch)
    # Request 2: size=6 (would make total 12 > 10, should be deferred)
    # Each should be processed in its own batch
    t1 = get_or_create_event_loop().create_task(
        batch_handler({"size": 6, "value": "first"})
    )
    t2 = get_or_create_event_loop().create_task(
        batch_handler({"size": 12, "value": "second"})
    )

    t1_result = await t1
    assert t1_result == "first"
    with pytest.raises(RuntimeError):
        await t2


@pytest.mark.asyncio
async def test_batch_size_fn_multiple_items_fit():
    """Test that multiple items are batched together when they fit within max_batch_size."""
    batches_seen = []

    @serve.batch(
        max_batch_size=20,
        batch_wait_timeout_s=0.1,
        batch_size_fn=lambda items: sum(item["size"] for item in items),
    )
    async def batch_handler(requests):
        batch_values = [req["value"] for req in requests]
        batches_seen.append(batch_values)
        return batch_values

    # All three requests fit: 5 + 6 + 7 = 18 <= 20
    t1 = get_or_create_event_loop().create_task(
        batch_handler({"size": 5, "value": "a"})
    )
    t2 = get_or_create_event_loop().create_task(
        batch_handler({"size": 6, "value": "b"})
    )
    t3 = get_or_create_event_loop().create_task(
        batch_handler({"size": 7, "value": "c"})
    )

    done, pending = await asyncio.wait([t1, t2, t3], timeout=1.0)
    assert len(done) == 3
    assert len(pending) == 0

    # All three should be in the same batch since they fit
    assert len(batches_seen) == 1, f"Expected 1 batch, got {len(batches_seen)}"
    assert set(batches_seen[0]) == {"a", "b", "c"}


def test_warn_if_max_batch_size_exceeds_max_ongoing_requests():
    """Test warn_if_max_batch_size_exceeds_max_ongoing_requests() logged the warning
     message correctly.

    When the queue starts with or updated `max_batch_size` to be larger than
    max_ongoing_requests, log the warning to suggest configuring `max_ongoing_requests`.
    When the queue starts with or updated `max_batch_size` to be smaller or equal than
    max_ongoing_requests, no warning should be logged.
    """
    logger = logging.getLogger(SERVE_LOGGER_NAME)
    stream = FakeStream()
    stream_handler = logging.StreamHandler(stream)
    logger.addHandler(stream_handler)
    bound = default_deployment_config.max_ongoing_requests
    over_bound = bound + 1
    under_bound = bound - 1
    over_bound_warning_message = (
        f"`max_batch_size` ({over_bound}) * `max_concurrent_batches` "
        f"({1}) is larger than `max_ongoing_requests` "
        f"({bound}). This means the replica will never achieve "
        "the configured `max_batch_size` concurrently. Please update "
        "`max_ongoing_requests` to be >= `max_batch_size` * `max_concurrent_batches`.\n"
    )

    # Start queue above the bound will log warning. Start at under or at the bound will
    # not log warning
    for max_batch_size in [over_bound, under_bound, bound]:
        queue = _BatchQueue(
            max_batch_size=max_batch_size,
            batch_wait_timeout_s=1000,
            max_concurrent_batches=1,
        )
        if max_batch_size > bound:
            assert over_bound_warning_message in stream.messages
        else:
            assert over_bound_warning_message not in stream.messages
        stream.reset_message()

    # Update queue above the bound will log warning. Update at under or at the bound
    # will not log warning
    for max_batch_size in [over_bound, under_bound, bound]:
        queue.set_max_batch_size(max_batch_size)
        if max_batch_size > bound:
            assert over_bound_warning_message in stream.messages
        else:
            assert over_bound_warning_message not in stream.messages
        stream.reset_message()


if __name__ == "__main__":
    import sys

    sys.exit(pytest.main(["-v", "-s", __file__]))