e252ca7962
Full coverage Refactor code structure and improve logging configuration; update tests for activities and connectors
83 lines
2.1 KiB
Python
83 lines
2.1 KiB
Python
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
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_load_from_kafka(kafka):
|
|
input_data = {"topic": "test-topic"}
|
|
|
|
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"}
|
|
|
|
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)
|