# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/001_InMemoryBroker.ipynb.

# %% auto 0
__all__ = ['logger', 'KafkaRecord', 'KafkaPartition', 'KafkaTopic', 'split_list', 'GroupMetadata', 'InMemoryBroker',
           'InMemoryConsumer', 'InMemoryProducer', 'MockBatch']

# %% ../../nbs/001_InMemoryBroker.ipynb 1
import asyncio
import hashlib
import random
import string
import uuid
from contextlib import contextmanager
from dataclasses import dataclass, field
from typing import *
from types import ModuleType

from aiokafka.structs import ConsumerRecord, RecordMetadata, TopicPartition

import fastkafka._components.aiokafka_consumer_loop
import fastkafka._aiokafka_imports
from .._aiokafka_imports import AIOKafkaConsumer, AIOKafkaProducer
from .._components.logger import get_logger
from fastkafka._components.meta import (
    _get_default_kwargs_from_sig,
    classcontextmanager,
    delegates,
    patch,
)

# %% ../../nbs/001_InMemoryBroker.ipynb 3
logger = get_logger(__name__)

# %% ../../nbs/001_InMemoryBroker.ipynb 6
@dataclass
class KafkaRecord:
    topic: str = ""
    partition: int = 0
    key: Optional[bytes] = None
    value: bytes = b""
    offset: int = 0
    timestamp = 0
    timestamp_type = 0
    checksum = 0
    serialized_key_size = 0
    serialized_value_size = 0
    headers: Sequence[Tuple[str, bytes]] = field(default_factory=list)

# %% ../../nbs/001_InMemoryBroker.ipynb 7
class KafkaPartition:
    def __init__(self, *, partition: int, topic: str):
        """
        Initialize a KafkaPartition object.

        Args:
            partition: The partition number.
            topic: The topic name.
        """
        self.partition = partition
        self.topic = topic
        self.messages: List[KafkaRecord] = list()

    def write(self, value: bytes, key: Optional[bytes] = None) -> RecordMetadata:  # type: ignore
        """
        Write a Kafka record to the partition.

        Args:
            value: The value of the record.
            key: The key of the record.

        Returns:
            The record metadata.
        """
        record = KafkaRecord(
            topic=self.topic,
            partition=self.partition,
            value=value,
            key=key,
            offset=len(self.messages),
        )
        record_meta = RecordMetadata(
            topic=self.topic,
            partition=self.partition,
            topic_partition=TopicPartition(topic=self.topic, partition=self.partition),
            offset=len(self.messages),
            timestamp=1680602752070,
            timestamp_type=0,
            log_start_offset=0,
        )
        self.messages.append(record)
        return record_meta

    def read(self, offset: int) -> Tuple[List[KafkaRecord], int]:
        """
        Read Kafka records from the partition starting from the given offset.

        Args:
            offset: The starting offset.

        Returns:
            A tuple containing the list of records and the current offset.
        """
        return self.messages[offset:], len(self.messages)

    def latest_offset(self) -> int:
        """
        Get the latest offset of the partition.

        Returns:
            The latest offset.
        """
        return len(self.messages)

# %% ../../nbs/001_InMemoryBroker.ipynb 11
class KafkaTopic:
    def __init__(self, topic: str, num_partitions: int = 1):
        """
        Initialize a KafkaTopic object.

        Args:
            topic: The topic name.
            num_partitions: The number of partitions in the topic (default: 1).
        """
        self.topic = topic
        self.num_partitions = num_partitions
        self.partitions: List[KafkaPartition] = [
            KafkaPartition(topic=topic, partition=partition_index)
            for partition_index in range(num_partitions)
        ]

    def read(  # type: ignore
        self, partition: int, offset: int
    ) -> Tuple[TopicPartition, List[KafkaRecord], int]:
        """
        Read records from the specified partition and offset.

        Args:
            partition: The partition index.
            offset: The offset from which to start reading.

        Returns:
            A tuple containing the topic partition, list of Kafka records, and the new offset.
        """
        topic_partition = TopicPartition(topic=self.topic, partition=partition)
        records, offset = self.partitions[partition].read(offset)
        return topic_partition, records, offset

    def write_with_partition(  # type: ignore
        self,
        value: bytes,
        partition: int,
    ) -> RecordMetadata:
        """
        Write a record with a specified partition.

        Args:
            value: The value of the record.
            partition: The partition to write the record to.

        Returns:
            The metadata of the written record.
        """
        return self.partitions[partition].write(value)

    def write_with_key(self, value: bytes, key: bytes) -> RecordMetadata:  # type: ignore
        """
        Write a record with a specified key.

        Args:
            value: The value of the record.
            key: The key of the record.

        Returns:
            The metadata of the written record.
        """
        partition = int(hashlib.sha256(key).hexdigest(), 16) % self.num_partitions
        return self.partitions[partition].write(value, key=key)

    def write(  # type: ignore
        self,
        value: bytes,
        *,
        key: Optional[bytes] = None,
        partition: Optional[int] = None,
    ) -> RecordMetadata:
        """
        Write a record to the topic.

        Args:
            value: The value of the record.
            key: The key of the record (optional).
            partition: The partition to write the record to (optional).

        Returns:
            The metadata of the written record.
        """
        if partition is not None:
            return self.write_with_partition(value, partition)

        if key is not None:
            return self.write_with_key(value, key)

        partition = random.randint(0, self.num_partitions - 1)  # nosec
        return self.write_with_partition(value, partition)

    def latest_offset(self, partition: int) -> int:
        """
        Get the latest offset of a partition.

        Args:
            partition: The partition index.

        Returns:
            The latest offset of the partition.
        """
        return self.partitions[partition].latest_offset()

# %% ../../nbs/001_InMemoryBroker.ipynb 17
def split_list(list_to_split: List[Any], split_size: int) -> List[List[Any]]:
    """
    Split a list into smaller lists of a specified size.

    Args:
        list_to_split: The list to split.
        split_size: The size of each split.

    Returns:
        A list of smaller lists.
    """
    return [
        list_to_split[start_index : start_index + split_size]
        for start_index in range(0, len(list_to_split), split_size)
    ]

# %% ../../nbs/001_InMemoryBroker.ipynb 19
class GroupMetadata:
    def __init__(self, num_partitions: int):
        """
        Initialize a GroupMetadata object.

        Args:
            num_partitions: The number of partitions in the group.
        """
        self.num_partitions = num_partitions
        self.partitions_offsets: Dict[int, int] = {}
        self.consumer_ids: List[uuid.UUID] = list()
        self.partition_assignments: Dict[uuid.UUID, List[int]] = {}

    def subscribe(self, consumer_id: uuid.UUID) -> None:
        """
        Subscribe a consumer to the group.

        Args:
            consumer_id: The ID of the consumer.
        """
        self.consumer_ids.append(consumer_id)
        self.rebalance()

    def unsubscribe(self, consumer_id: uuid.UUID) -> None:
        """
        Unsubscribe a consumer from the group.

        Args:
            consumer_id: The ID of the consumer.
        """
        self.consumer_ids.remove(consumer_id)
        self.rebalance()

    def rebalance(self) -> None:
        """
        Rebalance the group's partition assignments.
        """
        if len(self.consumer_ids) == 0:
            self.partition_assignments = {}
        else:
            partitions_per_actor = self.num_partitions // len(self.consumer_ids)
            if self.num_partitions % len(self.consumer_ids) != 0:
                partitions_per_actor += 1
            self.assign_partitions(partitions_per_actor)

    def assign_partitions(self, partitions_per_actor: int) -> None:
        partitions = [i for i in range(self.num_partitions)]

        partitions_split = split_list(partitions, partitions_per_actor)
        self.partition_assignments = {
            self.consumer_ids[i]: partition_split
            for i, partition_split in enumerate(partitions_split)
        }

    def get_partitions(
        self, consumer_id: uuid.UUID
    ) -> Tuple[List[int], Dict[int, Optional[int]]]:
        """
        Get the partition assignments and offsets for a consumer.

        Args:
            consumer_id: The ID of the consumer.

        Returns:
            A tuple containing the partition assignments and offsets.
        """
        partition_assignments = self.partition_assignments.get(consumer_id, [])
        partition_offsets_assignments = {
            partition: self.partitions_offsets.get(partition, None)
            for partition in partition_assignments
        }
        return partition_assignments, partition_offsets_assignments

    def set_offset(self, partition: int, offset: int) -> None:
        """
        Set the offset for a partition.

        Args:
            partition: The partition index.
            offset: The offset to set.
        """
        self.partitions_offsets[partition] = offset

# %% ../../nbs/001_InMemoryBroker.ipynb 22
@classcontextmanager()
class InMemoryBroker:
    def __init__(
        self,
        num_partitions: int = 1,
    ):
        self.num_partitions = num_partitions
        self.topics: Dict[Tuple[str, str], KafkaTopic] = {}
        self.topic_groups: Dict[Tuple[str, str, str], GroupMetadata] = {}
        self.is_started: bool = False

    def connect(self) -> uuid.UUID:
        return uuid.uuid4()

    def dissconnect(self, consumer_id: uuid.UUID) -> None:
        """
        Disconnect a consumer from the broker.

        Args:
            consumer_id: The ID of the consumer.
        """
        pass

    def subscribe(
        self, bootstrap_server: str, topic: str, group: str, consumer_id: uuid.UUID
    ) -> None:
        raise NotImplementedError()

    def unsubscribe(
        self, bootstrap_server: str, topic: str, group: str, consumer_id: uuid.UUID
    ) -> None:
        raise NotImplementedError()

    def read(  # type: ignore
        self,
        *,
        bootstrap_server: str,
        topic: str,
        group: str,
        consumer_id: uuid.UUID,
        auto_offset_reset: str,
    ) -> Dict[TopicPartition, List[KafkaRecord]]:
        raise NotImplementedError()

    def write(  # type: ignore
        self,
        *,
        bootstrap_server: str,
        topic: str,
        value: bytes,
        key: Optional[bytes] = None,
        partition: Optional[int] = None,
    ) -> RecordMetadata:
        raise NotImplementedError()

    @contextmanager
    def lifecycle(self) -> Iterator["InMemoryBroker"]:
        """
        Context manager for the lifecycle of the in-memory broker.

        Yields:
            An instance of the in-memory broker.
        """
        raise NotImplementedError()

    async def _start(self) -> str:
        """
        Start the in-memory broker.

        Returns:
            The address of the broker.
        """
        logger.info("InMemoryBroker._start() called")
        self.__enter__()  # type: ignore
        return "localbroker:0"

    async def _stop(self) -> None:
        """
        Stop the in-memory broker.
        """
        logger.info("InMemoryBroker._stop() called")
        self.__exit__(None, None, None)  # type: ignore

# %% ../../nbs/001_InMemoryBroker.ipynb 23
@patch
def subscribe(
    self: InMemoryBroker,
    bootstrap_server: str,
    topic: str,
    group: str,
    consumer_id: uuid.UUID,
) -> None:
    """
    Subscribe a consumer to a topic group.

    Args:
        bootstrap_server: The bootstrap server address.
        topic: The topic to subscribe to.
        group: The group to join.
        consumer_id: The ID of the consumer.
    """
    if (bootstrap_server, topic) not in self.topics:
        self.topics[(bootstrap_server, topic)] = KafkaTopic(
            topic=topic, num_partitions=self.num_partitions
        )

    group_meta = self.topic_groups.get(
        (bootstrap_server, topic, group), GroupMetadata(self.num_partitions)
    )
    group_meta.subscribe(consumer_id)
    self.topic_groups[(bootstrap_server, topic, group)] = group_meta


@patch
def unsubscribe(
    self: InMemoryBroker,
    bootstrap_server: str,
    topic: str,
    group: str,
    consumer_id: uuid.UUID,
) -> None:
    """
    Unsubscribe a consumer from a topic group.

    Args:
        bootstrap_server: The bootstrap server address.
        topic: The topic to unsubscribe from.
        group: The group to leave.
        consumer_id: The ID of the consumer.
    """
    self.topic_groups[(bootstrap_server, topic, group)].unsubscribe(consumer_id)

# %% ../../nbs/001_InMemoryBroker.ipynb 25
@patch
def write(  # type: ignore
    self: InMemoryBroker,
    *,
    bootstrap_server: str,
    topic: str,
    value: bytes,
    key: Optional[bytes] = None,
    partition: Optional[int] = None,
) -> RecordMetadata:
    """
    Write a message to a topic.

    Args:
        bootstrap_server: The bootstrap server address.
        topic: The topic to write the message to.
        value: The value of the message.
        key: The key associated with the message.
        partition: The partition ID to write the message to.

    Returns:
        The metadata of the written message.
    """
    if (bootstrap_server, topic) not in self.topics:
        self.topics[(bootstrap_server, topic)] = KafkaTopic(
            topic=topic, num_partitions=self.num_partitions
        )

    return self.topics[(bootstrap_server, topic)].write(
        value, key=key, partition=partition
    )

# %% ../../nbs/001_InMemoryBroker.ipynb 27
@patch
def read(  # type: ignore
    self: InMemoryBroker,
    *,
    bootstrap_server: str,
    topic: str,
    group: str,
    consumer_id: uuid.UUID,
    auto_offset_reset: str,
) -> Dict[TopicPartition, List[KafkaRecord]]:
    """
    Read messages from a topic group.

    Args:
        bootstrap_server: The bootstrap server address.
        topic: The topic to read messages from.
        group: The group to read messages for.
        consumer_id: The ID of the consumer.
        auto_offset_reset: The strategy to use when the consumer does not have a valid offset for the group.

    Returns:
        A dictionary containing the messages retrieved from each topic partition.
    """
    group_meta = self.topic_groups[(bootstrap_server, topic, group)]
    partitions, offsets = group_meta.get_partitions(consumer_id)

    if len(partitions) == 0:
        return {}

    partitions_data = {}

    for partition in partitions:
        offset = offsets[partition]

        if offset is None:
            offset = (
                self.topics[(bootstrap_server, topic)].latest_offset(partition)
                if auto_offset_reset == "latest"
                else 0
            )

        topic_partition, data, offset = self.topics[(bootstrap_server, topic)].read(
            partition, offset
        )

        partitions_data[topic_partition] = data
        group_meta.set_offset(partition, offset)

    return partitions_data

# %% ../../nbs/001_InMemoryBroker.ipynb 34
# InMemoryConsumer
class InMemoryConsumer:
    def __init__(
        self,
        broker: InMemoryBroker,
    ) -> None:
        self.broker = broker
        self._id: Optional[uuid.UUID] = None
        self._auto_offset_reset: str = "latest"
        self._group_id: Optional[str] = None
        self._topics: List[str] = list()
        self._bootstrap_servers = ""

    @delegates(AIOKafkaConsumer)
    def __call__(self, **kwargs: Any) -> "InMemoryConsumer":
        defaults = _get_default_kwargs_from_sig(InMemoryConsumer.__call__, **kwargs)
        consume_copy = InMemoryConsumer(self.broker)
        consume_copy._auto_offset_reset = defaults["auto_offset_reset"]
        consume_copy._bootstrap_servers = (
            "".join(defaults["bootstrap_servers"])
            if isinstance(defaults["bootstrap_servers"], list)
            else defaults["bootstrap_servers"]
        )

        consume_copy._group_id = (
            defaults["group_id"]
            if defaults["group_id"] is not None
            else "".join(random.choices(string.ascii_letters, k=10))  # nosec
        )
        return consume_copy

    @delegates(AIOKafkaConsumer.start)
    async def start(self, **kwargs: Any) -> None:
        pass

    @delegates(AIOKafkaConsumer.stop)
    async def stop(self, **kwargs: Any) -> None:
        pass

    @delegates(AIOKafkaConsumer.subscribe)
    def subscribe(self, topics: List[str], **kwargs: Any) -> None:
        raise NotImplementedError()

    @delegates(AIOKafkaConsumer.getmany)
    async def getmany(  # type: ignore
        self, **kwargs: Any
    ) -> Dict[TopicPartition, List[ConsumerRecord]]:
        raise NotImplementedError()

# %% ../../nbs/001_InMemoryBroker.ipynb 37
@patch
@delegates(AIOKafkaConsumer.start)
async def start(self: InMemoryConsumer, **kwargs: Any) -> None:
    """
    Start consuming messages from the connected broker.

    Raises:
        RuntimeError: If start() has already been called without calling stop() first.
    """
    logger.info("AIOKafkaConsumer patched start() called()")
    if self._id is not None:
        raise RuntimeError(
            "Consumer start() already called! Run consumer stop() before running start() again"
        )
    self._id = self.broker.connect()

# %% ../../nbs/001_InMemoryBroker.ipynb 40
@patch  # type: ignore
@delegates(AIOKafkaConsumer.subscribe)
def subscribe(self: InMemoryConsumer, topics: List[str], **kwargs: Any) -> None:
    """
    Subscribe to a list of topics for consuming messages.

    Args:
        topics: A list of topics to subscribe to.

    Raises:
        RuntimeError: If start() has not been called before calling subscribe().
    """
    logger.info("AIOKafkaConsumer patched subscribe() called")
    if self._id is None:
        raise RuntimeError("Consumer start() not called! Run consumer start() first")
    logger.info(f"AIOKafkaConsumer.subscribe(), subscribing to: {topics}")
    for topic in topics:
        self.broker.subscribe(
            bootstrap_server=self._bootstrap_servers,
            consumer_id=self._id,
            topic=topic,
            group=self._group_id,  # type: ignore
        )
        self._topics.append(topic)

# %% ../../nbs/001_InMemoryBroker.ipynb 43
@patch
@delegates(AIOKafkaConsumer.stop)
async def stop(self: InMemoryConsumer, **kwargs: Any) -> None:
    """
    Stop consuming messages from the connected broker.

    Raises:
        RuntimeError: If start() has not been called before calling stop().
    """
    logger.info("AIOKafkaConsumer patched stop() called")
    if self._id is None:
        raise RuntimeError("Consumer start() not called! Run consumer start() first")
    for topic in self._topics:
        self.broker.unsubscribe(
            bootstrap_server=self._bootstrap_servers,
            topic=topic,
            group=self._group_id,  # type: ignore
            consumer_id=self._id,
        )

# %% ../../nbs/001_InMemoryBroker.ipynb 46
@patch
@delegates(AIOKafkaConsumer.getmany)
async def getmany(  # type: ignore
    self: InMemoryConsumer, **kwargs: Any
) -> Dict[TopicPartition, List[ConsumerRecord]]:
    """
    Retrieve messages from the subscribed topics.

    Returns:
        A dictionary containing the retrieved messages from each topic partition.

    Raises:
        RuntimeError: If start() has not been called before calling getmany().
    """
    await asyncio.sleep(0)
    for topic in self._topics:
        return self.broker.read(
            bootstrap_server=self._bootstrap_servers,
            topic=topic,
            consumer_id=self._id,  # type: ignore
            group=self._group_id,  # type: ignore
            auto_offset_reset=self._auto_offset_reset,
        )

# %% ../../nbs/001_InMemoryBroker.ipynb 49
class InMemoryProducer:
    def __init__(self, broker: InMemoryBroker, **kwargs: Any) -> None:
        self.broker = broker
        self.id: Optional[uuid.UUID] = None
        self._bootstrap_servers = ""

    @delegates(AIOKafkaProducer)
    def __call__(self, **kwargs: Any) -> "InMemoryProducer":
        defaults = _get_default_kwargs_from_sig(InMemoryConsumer.__call__, **kwargs)
        producer_copy = InMemoryProducer(self.broker)
        producer_copy._bootstrap_servers = (
            "".join(defaults["bootstrap_servers"])
            if isinstance(defaults["bootstrap_servers"], list)
            else defaults["bootstrap_servers"]
        )
        return producer_copy

    @delegates(AIOKafkaProducer.start)
    async def start(self, **kwargs: Any) -> None:
        raise NotImplementedError()

    @delegates(AIOKafkaProducer.stop)
    async def stop(self, **kwargs: Any) -> None:
        raise NotImplementedError()

    @delegates(AIOKafkaProducer.send)
    async def send(  # type: ignore
        self,
        topic: str,
        msg: bytes,
        key: Optional[bytes] = None,
        **kwargs: Any,
    ):
        raise NotImplementedError()

    @delegates(AIOKafkaProducer.partitions_for)
    async def partitions_for(self, topic: str) -> List[int]:
        raise NotImplementedError()

    @delegates(AIOKafkaProducer._partition)
    def _partition(
        self, topic: str, arg1: Any, arg2: Any, arg3: Any, key: bytes, arg4: Any
    ) -> int:
        raise NotImplementedError()

    @delegates(AIOKafkaProducer.create_batch)
    def create_batch(self) -> "MockBatch":
        raise NotImplementedError()

    @delegates(AIOKafkaProducer.send_batch)
    async def send_batch(self, batch: "MockBatch", topic: str, partition: Any) -> None:
        raise NotImplementedError()

# %% ../../nbs/001_InMemoryBroker.ipynb 52
@patch  # type: ignore
@delegates(AIOKafkaProducer.start)
async def start(self: InMemoryProducer, **kwargs: Any) -> None:
    """
    Start the in-memory producer.

    Raises:
        RuntimeError: If start() has already been called without calling stop() first.
    """
    logger.info("AIOKafkaProducer patched start() called()")
    if self.id is not None:
        raise RuntimeError(
            "Producer start() already called! Run producer stop() before running start() again"
        )
    self.id = self.broker.connect()

# %% ../../nbs/001_InMemoryBroker.ipynb 55
@patch  # type: ignore
@delegates(AIOKafkaProducer.stop)
async def stop(self: InMemoryProducer, **kwargs: Any) -> None:
    """
    Stop the in-memory producer.

    Raises:
        RuntimeError: If start() has not been called before calling stop().
    """
    logger.info("AIOKafkaProducer patched stop() called")
    if self.id is None:
        raise RuntimeError("Producer start() not called! Run producer start() first")

# %% ../../nbs/001_InMemoryBroker.ipynb 58
@patch
@delegates(AIOKafkaProducer.send)
async def send(  # type: ignore
    self: InMemoryProducer,
    topic: str,
    msg: bytes,
    key: Optional[bytes] = None,
    partition: Optional[int] = None,
    **kwargs: Any,
):  # asyncio.Task[RecordMetadata]
    """
    Send a message to the specified topic.

    Args:
        topic: The topic to send the message to.
        msg: The message to send.
        key: The key associated with the message (optional).
        partition: The partition to send the message to (optional).
        **kwargs: Additional arguments to be passed to AIOKafkaProducer.send().

    Returns:
        A task that resolves to the RecordMetadata of the sent message.

    Raises:
        RuntimeError: If start() has not been called before calling send().
    """
    if self.id is None:
        raise RuntimeError("Producer start() not called! Run producer start() first")

    record = self.broker.write(
        bootstrap_server=self._bootstrap_servers,
        topic=topic,
        value=msg,
        key=key,
        partition=partition,
    )

    async def _f(record: ConsumerRecord = record) -> RecordMetadata:  # type: ignore
        return record

    return asyncio.create_task(_f())

# %% ../../nbs/001_InMemoryBroker.ipynb 60
@patch
@delegates(AIOKafkaProducer.partitions_for)
async def partitions_for(self: InMemoryProducer, topic: str) -> List[int]:
    """
    Retrieve the list of partitions for the specified topic.

    Args:
        topic: The topic to get the partitions for.

    Returns:
        A list of partition IDs.
    """
    return [i for i in range(self.broker.num_partitions)]

# %% ../../nbs/001_InMemoryBroker.ipynb 62
@patch
@delegates(AIOKafkaProducer._partition)
def _partition(
    self: InMemoryProducer,
    topic: str,
    arg1: Any,
    arg2: Any,
    arg3: Any,
    key: bytes,
    arg4: Any,
) -> int:
    """
    Determine the partition to which the message should be sent.

    Args:
        topic: The topic to send the message to.
        arg1, arg2, arg3, arg4: Additional arguments passed to the original AIOKafkaProducer._partition().

    Returns:
        The partition ID.
    """
    return int(hashlib.sha256(key).hexdigest(), 16) % self.broker.num_partitions

# %% ../../nbs/001_InMemoryBroker.ipynb 64
class MockBatch:
    def __init__(self) -> None:
        """
        Initialize an instance of MockBatch.
        """
        self._batch: List[Tuple] = list()

    def append(  # type: ignore
        self, key: Optional[bytes], value: bytes, timestamp: int
    ) -> RecordMetadata:
        """
        Append a message to the batch.

        Args:
            key: The key associated with the message (optional).
            value: The value of the message.
            timestamp: The timestamp of the message.

        Returns:
            The RecordMetadata of the appended message.
        """
        self._batch.append((key, value))
        return RecordMetadata(
            topic="",
            partition=0,
            topic_partition=None,
            offset=0,
            timestamp=timestamp,
            timestamp_type=0,
            log_start_offset=0,
        )


@patch
@delegates(AIOKafkaProducer.create_batch)
def create_batch(self: InMemoryProducer) -> "MockBatch":
    """
    Create a mock batch for the in-memory producer.

    Returns:
        A MockBatch instance.
    """
    return MockBatch()


@patch
@delegates(AIOKafkaProducer.send_batch)
async def send_batch(
    self: InMemoryProducer, batch: "MockBatch", topic: str, partition: Any
) -> None:
    """
    Send a batch of messages to the specified topic and partition.

    Args:
        batch: The MockBatch containing the messages to send.
        topic: The topic to send the batch of messages to.
        partition: The partition to send the batch of messages to.
    """
    for record in batch._batch:
        self.broker.write(
            bootstrap_server=self._bootstrap_servers,
            topic=topic,
            value=record[1],
            key=record[0],
            partition=partition,
        )

# %% ../../nbs/001_InMemoryBroker.ipynb 68
@patch
@contextmanager
def lifecycle(self: InMemoryBroker) -> Iterator[InMemoryBroker]:
    """
    Context manager for the lifecycle of the in-memory broker.

    Yields:
        An instance of the in-memory broker.
    """
    logger.info(
        "InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!"
    )
    try:
        logger.info("InMemoryBroker starting")

        old_consumer = fastkafka._aiokafka_imports.AIOKafkaConsumer
        old_producer = fastkafka._aiokafka_imports.AIOKafkaProducer

        fastkafka._aiokafka_imports.AIOKafkaConsumer = InMemoryConsumer(self)
        fastkafka._aiokafka_imports.AIOKafkaProducer = InMemoryProducer(self)

        self.is_started = True
        yield self
    finally:
        logger.info("InMemoryBroker stopping")

        fastkafka._aiokafka_imports.AIOKafkaConsumer = old_consumer
        fastkafka._aiokafka_imports.AIOKafkaProducer = old_producer

        self.is_started = False
