# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/999_Test_Utils.ipynb.

# %% auto 0
__all__ = ['logger', 'kafka_server_url', 'kafka_server_port', 'kafka_config', 'true_after', 'create_missing_topics',
           'create_testing_topic', 'create_and_fill_testing_topic', 'nb_safe_seed', 'mock_AIOKafkaProducer_send']

# %% ../nbs/999_Test_Utils.ipynb 1
from typing import List, Dict, Any, Optional, Callable, Tuple, Generator
from os import environ
from contextlib import contextmanager, asynccontextmanager
import random
from datetime import datetime, timedelta
import time
import asyncio
import hashlib

import unittest

from confluent_kafka.admin import AdminClient, NewTopic
from aiokafka import AIOKafkaProducer, AIOKafkaConsumer

from ._components.logger import get_logger

# %% ../nbs/999_Test_Utils.ipynb 3
logger = get_logger(__name__)

# %% ../nbs/999_Test_Utils.ipynb 5
kafka_server_url = (
    environ["KAFKA_HOSTNAME"] if "KAFKA_HOSTNAME" in environ else "localhost"
)
kafka_server_port = environ["KAFKA_PORT"] if "KAFKA_PORT" in environ else "9092"

kafka_config = {
    "bootstrap.servers": f"{kafka_server_url}:{kafka_server_port}",
    # "group.id": f"{kafka_server_url}:{kafka_server_port}_group"
}

# %% ../nbs/999_Test_Utils.ipynb 6
def true_after(seconds: float):
    """Function returning True after a given number of seconds"""
    t = datetime.now()

    def _true_after(seconds=seconds, t=t):
        return (datetime.now() - t) > timedelta(seconds=seconds)

    return _true_after

# %% ../nbs/999_Test_Utils.ipynb 8
## TODO: Check if replication num is <= of number of brokers
## TODO: Add tests for:
#             - Replication factor (less than and greater than number of brokers)
#             - Num partitions


def create_missing_topics(
    admin: AdminClient,
    topic_names: List[str],
    *,
    num_partitions: Optional[int] = None,
    replication_factor: Optional[int] = None,
    **kwargs,
) -> None:
    if not replication_factor:
        replication_factor = len(admin.list_topics().brokers)
    if not num_partitions:
        num_partitions = replication_factor
    existing_topics = list(admin.list_topics().topics.keys())
    logger.debug(
        f"create_missing_topics({topic_names}): existing_topics={existing_topics}, num_partitions={num_partitions}, replication_factor={replication_factor}"
    )
    new_topics = [
        NewTopic(
            topic,
            num_partitions=num_partitions,
            replication_factor=replication_factor,
            **kwargs,
        )
        for topic in topic_names
        if topic not in existing_topics
    ]
    if len(new_topics):
        logger.info(f"create_missing_topics({topic_names}): new_topics = {new_topics}")
        fs = admin.create_topics(new_topics)
        while not set(topic_names).issubset(set(admin.list_topics().topics.keys())):
            time.sleep(1)

# %% ../nbs/999_Test_Utils.ipynb 10
@contextmanager
def create_testing_topic(
    kafka_config: Dict[str, Any], topic_prefix: str, seed: Optional[int] = None
) -> Generator[str, None, None]:
    # create random topic name
    random.seed(seed)
    topic = topic_prefix + str(random.randint(0, 10**10)).zfill(3)

    # delete topic if it already exists
    admin = AdminClient(kafka_config)
    existing_topics = admin.list_topics().topics.keys()
    if topic in existing_topics:
        logger.warning(f"topic {topic} exists, deleting it...")
        fs = admin.delete_topics(topics=[topic])
        results = {k: f.result() for k, f in fs.items()}
        while topic in admin.list_topics().topics.keys():
            time.sleep(1)
    try:
        # create topic if needed
        create_missing_topics(admin, [topic])
        while topic not in admin.list_topics().topics.keys():
            time.sleep(1)
        yield topic

    finally:
        pass
        # cleanup if needed again
        fs = admin.delete_topics(topics=[topic])
        while topic in admin.list_topics().topics.keys():
            time.sleep(1)

# %% ../nbs/999_Test_Utils.ipynb 12
@asynccontextmanager
async def create_and_fill_testing_topic(
    msgs: List[bytes], kafka_config: Dict[str, str] = kafka_config, *, seed: int
) -> Generator[str, None, None]:

    with create_testing_topic(kafka_config, "my_topic_", seed=seed) as topic:

        producer = AIOKafkaProducer(bootstrap_servers=kafka_config["bootstrap.servers"])
        logger.info(f"Producer {producer} created.")

        await producer.start()
        logger.info(f"Producer {producer} started.")
        try:
            fx = [
                producer.send(
                    topic,
                    msg,
                    key=f"{i % 17}".encode("utf-8"),
                )
                for i, msg in enumerate(msgs)
            ]
            await producer.flush()
            sent_msgs = [await f for f in fx]
            msg_statuses = [await s for s in sent_msgs]
            logger.info(f"Sent messages: len(sent_msgs)={len(sent_msgs)}")

            yield topic
        finally:
            await producer.stop()
            logger.info(f"Producer {producer} stoped.")

# %% ../nbs/999_Test_Utils.ipynb 15
def nb_safe_seed(s: str) -> Callable[[int], int]:
    """Gets a unique seed function for a notebook

    Params:
        s: name of the notebook used to initialize the seed function

    Returns:
        A unique seed function
    """
    init_seed = int(hashlib.sha1(s.encode("utf-8")).hexdigest(), 16) % (10**8)

    def _get_seed(x: int = 0, *, init_seed: int = init_seed) -> int:
        return init_seed + x

    return _get_seed

# %% ../nbs/999_Test_Utils.ipynb 17
@contextmanager
def mock_AIOKafkaProducer_send():
    """Mocks **send** method of **AIOKafkaProducer**"""
    with unittest.mock.patch("__main__.AIOKafkaProducer.send") as mock:

        async def _f():
            pass

        mock.return_value = asyncio.create_task(_f())

        yield mock
