gitea_admin/sientia-dataops-scouter_temporal
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

SIENTIAPDE-1110

Browse Source 
Enhance Gates and Redis activities by adding metadata parameter to apply_aggregation and notification methods. Refactor notification handling to use send_notification for improved consistency. Update tests to reflect changes in notification method calls and ensure proper functionality with new metadata integration.
This commit is contained in:
vitor-aignosi
2025-07-04 11:01:20 -03:00
parent d08d1b1337
commit 08240efc8b
6 changed files with 425 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
scouter/activities/gates.py
View File

@@ -16,7 +16,8 @@ quality_gate_filters = {
class Gates(BaseActivity):
def apply_aggregation(self, group: DataFrame, aggr_function: str) -> float | None | str:
def apply_aggregation(self, group: DataFrame, aggr_function: str,
metadata: dict[str, Any]) -> float | None | str:
"""
Apply aggregation function to a group of data.
@@ -48,7 +49,8 @@ class Gates(BaseActivity):
elif aggr_function == 'min':
return group['value'].min()
else:
self.notification_handler.build_and_send_notification(
self.send_notification(
metadata=metadata,
notification_id="AGGREGATION_ISSUES",
message=f"Invalid aggregation function: {aggr_function}",
block="aggregate_data",
@@ -100,7 +102,8 @@ class Gates(BaseActivity):
# Get the latest timestamp
latest_timestamp = group['timestamp'].max()
aggr_value = self.apply_aggregation(group, aggr_function)
aggr_value = self.apply_aggregation(
group, aggr_function, metadata)
if aggr_value == 'continue':
continue
@@ -145,7 +148,8 @@ class Gates(BaseActivity):
except Exception as e:
trace = traceback.format_exc()
self.notification_handler.build_and_send_notification(
self.send_notification(
metadata=metadata,
notification_id="AGGREGATION_ISSUES",
message=f"Error aggregating data: {e}",
block="aggregate_data",
@@ -202,7 +206,8 @@ class Gates(BaseActivity):
except Exception as e:
trace = traceback.format_exc()
self.notification_handler.build_and_send_notification(
self.send_notification(
metadata=metadata,
notification_id="DATA_QUALITY_GATE_ISSUES",
message=f"Error applying filter {filter_name}: {e}",
block="data_quality_gate",
@@ -219,7 +224,8 @@ class Gates(BaseActivity):
message = f"{len(filtered_data)} rows has quality issues: {filter_name}: {policy}"
attachment = filtered_data.to_string()
self.notification_handler.build_and_send_notification(
self.send_notification(
metadata=metadata,
notification_id=f"DATA_QUALITY_GATE_ISSUES__{filter_name}",
message=message,
block="data_quality_gate",

81
scouter/activities/redis.py
View File

@@ -1,8 +1,10 @@
import traceback
from temporalio import workflow, activity
with workflow.unsafe.imports_passed_through():
from logging import Logger
from sientia_do.notifications.handlers import NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from sientia_do.temporal.activities.redis_base import Redis as RedisBase
from sientia_do.temporal.utils.logger import Logger
from typing import Any
@@ -26,7 +28,18 @@ class Redis(RedisBase):
metadata = input_data['metadata']
key = f"last_data_timestamp_{input_data['workflow_name']}_{input_data['schedule_name']}"
data_hold = self.get(key)
try:
data_hold = self.get(key)
except Exception as e:
self.send_notification(
metadata=metadata,
notification_id="REDIS_GET_ERROR",
message=f"Error getting last data timestamp: {e}",
block="get_last_data_timestamp",
level=NotificationLevel.ERROR,
attachment_content=traceback.format_exc()
)
raise e
self.debug(
f"Last collected timestamp: {data_hold}",
@@ -48,6 +61,12 @@ class Redis(RedisBase):
data = DataFrame(input_data['data'])
if data.empty:
self.warning("No data to insert",
metadata=metadata
)
return None
last_data_timestamp = data['inserted_at'].max()
self.debug(
@@ -55,7 +74,19 @@ class Redis(RedisBase):
metadata=metadata
)
self.set(key, last_data_timestamp, ttl=None)
try:
self.set(key, last_data_timestamp, ttl=None)
except Exception as e:
self.send_notification(
metadata=metadata,
notification_id="REDIS_SET_ERROR",
message=f"Error setting last data timestamp: {e}",
block="put_last_data_timestamp",
level=NotificationLevel.ERROR,
attachment_content=traceback.format_exc()
)
raise e
return last_data_timestamp
@@ -83,7 +114,18 @@ class Redis(RedisBase):
key = f"held_data_{input_data['workflow_name']}_{input_data['schedule_name']}"
data_hold = self.get(key)
try:
data_hold = self.get(key)
except Exception as e:
self.send_notification(
metadata=metadata,
notification_id="REDIS_GET_ERROR",
message=f"Error getting held data: {e}",
block="group_and_hold_data",
level=NotificationLevel.ERROR,
attachment_content=traceback.format_exc()
)
raise e
if not data_hold:
data_hold = {}
@@ -93,22 +135,33 @@ class Redis(RedisBase):
)
return data_hold
for _, row in data.iterrows():
value = row['value']
try:
for _, row in data.iterrows():
value = row['value']
data_hold[row['name']] = value
data_hold[row['name']] = value
data_hold['timestamp'] = data['timestamp'].max() if not data.empty else \
datetime.now().strftime("%Y-%m-%d %H:%M:%S")
data_hold['timestamp'] = data['timestamp'].max() if not data.empty else \
datetime.now().strftime("%Y-%m-%d %H:%M:%S")
self.set(key, data_hold, ttl=retention_time)
self.set(key, data_hold, ttl=retention_time)
data_hold_df = DataFrame(data_hold, index=[0])
data_hold_melted = data_hold_df.melt(
id_vars='timestamp', var_name='variable', value_name='value')
data_hold_melted['model_id'] = input_data['model_id']
data_hold_df = DataFrame(data_hold, index=[0])
data_hold_melted = data_hold_df.melt(
id_vars='timestamp', var_name='variable', value_name='value')
data_hold_melted['model_id'] = input_data['model_id']
data_hold_melted.reset_index(drop=True, inplace=True)
data_hold_melted.reset_index(drop=True, inplace=True)
except Exception as e:
self.send_notification(
metadata=metadata,
notification_id="REDIS_SET_ERROR",
message=f"Error setting held data: {e}",
block="group_and_hold_data",
level=NotificationLevel.ERROR,
attachment_content=traceback.format_exc()
)
raise e
self.debug(
f"Data grouped and held successfully:\n {data_hold_melted.to_string()}",

32
tests/activities/test_gates.py
View File

@@ -11,7 +11,9 @@ def gates_fixture():
"""Fixture to create a Gates instance with mocked dependencies."""
logger = Mock()
notification_handler = MagicMock()
return Gates(logger=logger, notification_handler=notification_handler)
gates = Gates(logger=logger, notification_handler=notification_handler)
gates.send_notification = MagicMock()
return gates
metadata = {
@@ -51,7 +53,7 @@ async def test_data_quality_gate_with_null_values_filter_discard(gates_fixture):
# Verify
assert len(result['tag']) == 2
assert 'tag2' not in result['tag']
gates_fixture.notification_handler.build_and_send_notification.assert_called_once()
gates_fixture.send_notification.assert_called_once()
@pytest.mark.asyncio
@@ -84,7 +86,7 @@ async def test_data_quality_gate_with_out_of_bounds_filter_keep(gates_fixture):
# Verify data is kept but notification is sent
assert len(result['tag']) == 3 # All rows kept
gates_fixture.notification_handler.build_and_send_notification.assert_called_once()
gates_fixture.send_notification.assert_called_once()
@pytest.mark.asyncio
@@ -117,7 +119,7 @@ async def test_data_quality_gate_with_multiple_filters(gates_fixture):
assert result == {'tag': {0: 'tag1', 3: 'tag4'}, 'name': {0: 'tag1', 3: 'tag4'}, 'value': {
0: 1.0, 3: 4.0}, 'timestamp': {0: '2023-01-01', 3: '2023-01-04'}}
# Should be called twice (once for each filter)
assert gates_fixture.notification_handler.build_and_send_notification.call_count == 2
assert gates_fixture.send_notification.call_count == 2
@pytest.mark.asyncio
@@ -184,8 +186,8 @@ async def test_data_quality_gate_with_filter_error(gates_fixture):
# Verify error notification is sent and data is unchanged
assert len(result['tag']) == 1
gates_fixture.notification_handler.build_and_send_notification.assert_called_once()
call_args = gates_fixture.notification_handler.build_and_send_notification.call_args[1]
gates_fixture.send_notification.assert_called_once()
call_args = gates_fixture.send_notification.call_args[1]
assert call_args['notification_id'] == "DATA_QUALITY_GATE_ISSUES"
assert call_args['level'] == NotificationLevel.ERROR
assert "Filter error" in call_args['message']
@@ -214,7 +216,7 @@ async def test_data_quality_gate_with_empty_data(gates_fixture):
# Verify empty result and no notifications
assert len(result['tag']) == 0
gates_fixture.notification_handler.build_and_send_notification.assert_not_called()
gates_fixture.send_notification.assert_not_called()
@pytest.mark.asyncio
@@ -240,7 +242,7 @@ async def test_data_quality_gate_with_no_filters(gates_fixture):
# Verify data is unchanged and no notifications
assert len(result['tag']) == 1
gates_fixture.notification_handler.build_and_send_notification.assert_not_called()
gates_fixture.send_notification.assert_not_called()
@pytest.mark.parametrize(
@@ -265,14 +267,15 @@ async def test_data_quality_gate_with_no_filters(gates_fixture):
)
def test_apply_aggregation(gates_fixture, group_data, aggr_function, expected_result):
"""Test apply_aggregation method with various scenarios."""
result = gates_fixture.apply_aggregation(group_data, aggr_function)
result = gates_fixture.apply_aggregation(
group_data, aggr_function, metadata)
assert result == expected_result
# Check notification was sent for invalid function
if aggr_function == 'invalid':
gates_fixture.notification_handler.build_and_send_notification.assert_called_once()
gates_fixture.send_notification.assert_called_once()
else:
gates_fixture.notification_handler.build_and_send_notification.assert_not_called()
gates_fixture.send_notification.assert_not_called()
@pytest.mark.asyncio
@@ -308,7 +311,7 @@ async def test_aggregate_data(gates_fixture):
# Verify
assert result == expected_result
gates_fixture.notification_handler.build_and_send_notification.assert_not_called()
gates_fixture.send_notification.assert_not_called()
@pytest.mark.asyncio
@@ -342,7 +345,7 @@ async def test_aggregate_data_with_continue(gates_fixture):
# Verify
assert result == expected_result
gates_fixture.notification_handler.build_and_send_notification.assert_not_called()
gates_fixture.send_notification.assert_not_called()
@pytest.mark.asyncio
@@ -373,7 +376,8 @@ async def test_aggregate_data_raise_exception(gates_fixture):
await gates_fixture.aggregate_data(input_data)
except Exception as e:
assert str(e) == "Test exception"
gates_fixture.notification_handler.build_and_send_notification.assert_called_once_with(
gates_fixture.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id="AGGREGATION_ISSUES",
message="Error aggregating data: Test exception",
block="aggregate_data",

192
tests/activities/test_mongo.py Normal file
View File

@@ -0,0 +1,192 @@
from datetime import datetime
from unittest.mock import ANY, MagicMock, patch
from pytest import fixture, mark
from sientia_do.notifications.models import NotificationLevel
from scouter.activities.mongodb import MongoDB, clear_mongo_id
def test_clear_mongo_id():
"""Test clear_mongo_id"""
data = [
{'_id': '1', 'name': 'test1'},
{'_id': '2', 'name': [{
'_id': '3',
'name': 'test3'
}]}
]
result = clear_mongo_id(data)
assert result == [{'name': 'test1'}, {'name': [{'name': 'test3'}]}]
@patch('scouter.activities.mongodb.MongoClient')
def test_mongodb___init__(mock_mongo_client):
"""Test MongoDB __init__"""
mongo = MongoDB(
connection_string='mongodb://localhost:27017',
database_name='test_db',
logger=MagicMock(),
notification_handler=MagicMock()
)
mock_mongo_client.assert_called_once_with(
'mongodb://localhost:27017',
serverSelectionTimeoutMS=5000
)
mock_mongo_client.return_value.server_info.assert_called_once()
mock_mongo_client.return_value.__getitem__.assert_called_once_with(
'test_db')
assert mongo.client is not None
assert mongo.database is not None
@fixture
@patch('scouter.activities.mongodb.MongoClient')
def mongodb_activity(mock_mongo_client):
"""Test MongoDB activity"""
mongo = MongoDB(
connection_string='mongodb://localhost:27017',
database_name='test_db',
logger=MagicMock(),
notification_handler=MagicMock()
)
return mongo
def test_shutdown_success(mongodb_activity):
"""Test shutdown"""
mongodb_activity.shutdown()
mongodb_activity.client.close.assert_called_once()
def test_shutdown_error(mongodb_activity):
"""Test shutdown"""
mongodb_activity.client.close = MagicMock(side_effect=Exception('test'))
mongodb_activity.shutdown()
mongodb_activity.client.close.assert_called_once()
@mark.asyncio
async def test_load_latest_data_none_last_data_timestamp(mongodb_activity):
"""Test load_latest_data"""
collection = MagicMock()
mongodb_activity.database.__getitem__.return_value = collection
collection.find.return_value = [
{
'name': 'test1',
'value': 1,
'inserted_at': datetime.strptime(
'2023-01-01 12:00:00.000000', '%Y-%m-%d %H:%M:%S.%f')
}
]
result = await mongodb_activity.load_latest_data({
'metadata': {'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'},
'collection_name': 'test_collection',
'last_data_timestamp': None
})
mongodb_activity.database.__getitem__.assert_called_once_with(
'test_collection')
collection.find.assert_called_once_with(
{},
{"_id": 0}
)
assert result == {
'name': {
0: 'test1'
},
'value': {
0: 1
},
'inserted_at': {
0: '2023-01-01 12:00:00.000000'
}
}
@mark.asyncio
async def test_load_latest_data_not_none_last_data_timestamp(mongodb_activity):
"""Test load_latest_data"""
collection = MagicMock()
mongodb_activity.database.__getitem__.return_value = collection
collection.find.return_value = [
{
'name': 'test1',
'value': 1,
'inserted_at': datetime.strptime(
'2023-01-01 12:00:00.000000', '%Y-%m-%d %H:%M:%S.%f')
}
]
result = await mongodb_activity.load_latest_data({
'metadata': {'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'},
'collection_name': 'test_collection',
'last_data_timestamp': '2023-01-01 12:00:00.000000'
})
mongodb_activity.database.__getitem__.assert_called_once_with(
'test_collection')
collection.find.assert_called_once_with(
{
'inserted_at': {
'$gt': datetime.strptime(
'2023-01-01 12:00:00.000000', '%Y-%m-%d %H:%M:%S.%f')
}
},
{"_id": 0}
)
assert result == {
'name': {
0: 'test1'
},
'value': {
0: 1
},
'inserted_at': {
0: '2023-01-01 12:00:00.000000'
}
}
@mark.asyncio
async def test_load_latest_data_error(mongodb_activity):
"""Test load_latest_data"""
collection = MagicMock()
mongodb_activity.send_notification = MagicMock()
mongodb_activity.database.__getitem__.return_value = collection
collection.find.side_effect = Exception('test')
try:
await mongodb_activity.load_latest_data({
'metadata': {'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'},
'collection_name': 'test_collection',
'last_data_timestamp': '2023-01-01 12:00:00.000000'
})
except Exception as e:
assert str(e) == 'test'
mongodb_activity.send_notification.assert_called_once_with(
metadata={'workflow_name': 'test_pipeline',
'schedule_name': 'test_schedule'},
notification_id='MONGO_LOAD_ERROR',
message='Error loading data from MongoDB: test',
block='load_latest_data',
level=NotificationLevel.ERROR,
attachment_content=ANY
)

84
tests/activities/test_redis.py
View File

@@ -51,6 +51,90 @@ metadata = {
}
@pytest.mark.asyncio
async def test_get_last_data_timestamp_none(redis_activity):
"""Test get_last_data_timestamp"""
test_data = {
**metadata,
'workflow_name': 'test_pipeline',
'schedule_name': 'test_schedule'
}
redis_activity.get = MagicMock(return_value=None)
result = await redis_activity.get_last_data_timestamp(test_data)
assert result is None
@pytest.mark.asyncio
async def test_get_last_data_timestamp_not_none(redis_activity):
"""Test get_last_data_timestamp"""
test_data = {
**metadata,
'workflow_name': 'test_pipeline',
'schedule_name': 'test_schedule'
}
redis_activity.get = MagicMock(return_value='2023-01-01 12:00:00')
result = await redis_activity.get_last_data_timestamp(test_data)
redis_activity.get.assert_called_once_with(
'last_data_timestamp_test_pipeline_test_schedule'
)
assert result == '2023-01-01 12:00:00'
@pytest.mark.asyncio
async def test_put_last_data_timestamp_empty_dataframe(redis_activity):
"""Test put_last_data_timestamp with empty dataframe"""
test_data = {
**metadata,
'workflow_name': 'test_pipeline',
'schedule_name': 'test_schedule',
'data': DataFrame(columns=['name', 'value', 'timestamp']).to_dict('records')
}
redis_activity.set = MagicMock()
result = await redis_activity.put_last_data_timestamp(test_data)
assert result is None
redis_activity.set.assert_not_called()
@pytest.mark.asyncio
async def test_put_last_data_timestamp_not_empty_dataframe(redis_activity):
"""Test put_last_data_timestamp with not empty dataframe"""
data = DataFrame({
'name': ['sensor1', 'sensor2'],
'value': [25.5, 30.0],
'inserted_at': ['2023-01-01 12:00:00', '2023-01-01 12:00:01']
})
test_data = {
**metadata,
'workflow_name': 'test_pipeline',
'schedule_name': 'test_schedule',
'data': data.to_dict('records')
}
redis_activity.set = MagicMock()
result = await redis_activity.put_last_data_timestamp(test_data)
assert result == '2023-01-01 12:00:01'
redis_activity.set.assert_called_once_with(
'last_data_timestamp_test_pipeline_test_schedule',
'2023-01-01 12:00:01',
ttl=None
)
@pytest.mark.asyncio
async def test_group_and_hold_data_new_key(redis_activity):
"""Test group_and_hold_data with a new key"""

52
tests/utils/test_connectors_config.py
View File

@@ -2,6 +2,8 @@ import os
from unittest.mock import patch
import pytest
from scouter.utils.connectors_config import (
build_druid_config,
build_mongodb_config,
build_postgres_config,
build_kafka_config,
build_redis_config
@@ -113,3 +115,53 @@ def test_build_redis_config_with_env_vars():
'username': 'test',
'password': 'test'
}
def test_build_mongodb_config_defaults():
"""Test that build_mongodb_config returns default values when no env vars are set"""
config = build_mongodb_config()
assert config == {
'connection_string': 'mongodb://sientia:sientia@localhost:27017', # NOSONAR
'database_name': 'sientia'
}
def test_build_mongodb_config_with_env_vars():
"""Test that build_mongodb_config uses env vars when set"""
with patch.dict(os.environ, {
'MONGODB_URL': 'mongodb.example.com:27017',
'MONGODB_DATABASE_NAME': 'test_db',
'MONGODB_USERNAME': 'test',
'MONGODB_PASSWORD': 'test'
}):
config = build_mongodb_config()
assert config == {
'connection_string': 'mongodb://test:test@mongodb.example.com:27017',
'database_name': 'test_db'
}
def test_build_druid_config_defaults():
"""Test that build_druid_config returns default values when no env vars are set"""
config = build_druid_config()
assert config == {
'host': 'localhost',
'port': 8082
}
def test_build_druid_config_with_env_vars():
"""Test that build_druid_config uses env vars when set"""
with patch.dict(os.environ, {
'DRUID_HOST': 'druid.example.com',
'DRUID_PORT': '8083'
}):
config = build_druid_config()
assert config == {
'host': 'druid.example.com',
'port': 8083
}