4a90b77786
Update sientia-dataops-library version to 1.2.0 in requirements.txt; refactor logging in activities to use a unified Logger instance and include metadata in log messages across various activities.
93 lines
2.3 KiB
Python
93 lines
2.3 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
|
|
)
|
|
|
|
|
|
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)
|