from unittest.mock import MagicMock, patch, ANY
from pytest import fixture, mark
from pandas import DataFrame
from scouter.activities.kafka import Kafka


@fixture
@patch("scouter.activities.kafka.KafkaConsumer")
def kafka(_kafka_consumer):
    return Kafka(
        bootstrap_servers="localhost:9092",
        polling_time=1000,
        group_id="test-group",
        logger=MagicMock(),
        notification_handler=MagicMock()
    )


@patch("scouter.activities.kafka.KafkaConsumer")
def test___init__(kafka_consumer):
    kafka = Kafka(
        bootstrap_servers="localhost:9092",
        polling_time=1000,
        group_id="test-group",
        logger=MagicMock(),
        notification_handler=MagicMock()
    )

    assert kafka.polling_time == 1000
    assert kafka.kafka_connector == kafka_consumer.return_value

    kafka_consumer.assert_called_once_with(
        bootstrap_servers="localhost:9092",
        auto_offset_reset="earliest",
        enable_auto_commit=True,
        group_id="test-group",
        value_deserializer=ANY
    )


metadata = {
    'metadata': {
        'model_id': 'test_model_id',
        'model_name': 'test_model',
        'schedule_name': 'test_schedule',
        'workflow_name': 'scouter'
    }
}


@mark.asyncio
async def test_load_from_kafka(kafka):
    input_data = {"topic": "test-topic", **metadata}

    data = [
        ("test-topic", [
            MagicMock(
                value=f"test-value-{i}"
            ) for i in range(10)
        ])
    ]

    kafka.kafka_connector.poll.return_value = MagicMock(
        items=MagicMock(return_value=data)
    )

    expected = DataFrame([d.value for d in data[0][1]]).to_dict()

    result = await kafka.load_from_kafka(input_data)

    assert result == expected

    kafka.kafka_connector.subscribe.assert_called_once_with(["test-topic"])

    kafka.kafka_connector.poll.assert_called_once_with(timeout_ms=1000)


@mark.asyncio
async def test_load_from_kafka_empty(kafka):
    input_data = {"topic": "test-topic", **metadata}

    kafka.kafka_connector.poll.return_value = MagicMock(
        items=MagicMock(return_value=[])
    )

    result = await kafka.load_from_kafka(input_data)

    assert result == {}

    kafka.kafka_connector.subscribe.assert_called_once_with(["test-topic"])

    kafka.kafka_connector.poll.assert_called_once_with(timeout_ms=1000)