b203b7d22c
Implement workflows for fake data generation, scouter processing, and core scouter operations - Added `FakeData` workflow to generate random data and send it to a Kafka topic. - Implemented `Scouter` workflow to load data from Kafka and trigger the core scouter workflow. - Created `CoreScouter` workflow to process data through quality gates, aggregation, and export to PostgreSQL. - Developed comprehensive unit tests for activities and workflows, ensuring proper functionality and error handling. - Enhanced Redis and Postgres activities with robust testing for data handling and error notifications. - Introduced quality filters for data validation and implemented tests to verify their functionality.
91 lines
3.1 KiB
Python
91 lines
3.1 KiB
Python
from unittest.mock import ANY, MagicMock, patch
|
|
from pytest import fixture
|
|
from pytest import mark
|
|
from sientia_do.notifications.models import NotificationLevel
|
|
|
|
from scouter.activities.postgres import Postgres
|
|
|
|
|
|
@fixture
|
|
@patch("scouter.activities.postgres.create_engine")
|
|
@patch("scouter.activities.postgres.sessionmaker")
|
|
def postgres_client(mock_sessionmaker, mock_engine):
|
|
# Create a mock session
|
|
mock_session = MagicMock()
|
|
mock_session.commit = MagicMock()
|
|
mock_session.close = MagicMock()
|
|
|
|
# Configure the session to work with context management
|
|
mock_session.__enter__ = MagicMock(return_value=mock_session)
|
|
mock_session.__exit__ = MagicMock(return_value=None)
|
|
|
|
# Configure the sessionmaker to return our mock session
|
|
mock_sessionmaker.return_value = mock_session
|
|
|
|
# Configure the engine to return our mock sessionmaker
|
|
mock_engine.return_value = MagicMock()
|
|
mock_engine.return_value.dispose = MagicMock()
|
|
|
|
# Create the Postgres client
|
|
client = Postgres(
|
|
host="localhost",
|
|
port=5432,
|
|
user="postgres",
|
|
password="postgres",
|
|
dbname="postgres",
|
|
min_connections=1,
|
|
max_connections=10,
|
|
logger=MagicMock(),
|
|
notification_handler=MagicMock(),
|
|
)
|
|
|
|
# Set up the session factory
|
|
client.session_factory = mock_sessionmaker
|
|
|
|
return client
|
|
|
|
|
|
@mark.asyncio
|
|
@patch("scouter.activities.postgres.DataFrame")
|
|
async def test_export_data_to_postgres_success(mock_dataframe, postgres_client):
|
|
data = {"schema": "test", "table_name": "test",
|
|
"data": {"a": [1, 2, 3], "b": [4, 5, 6]}}
|
|
await postgres_client.export_data_to_postgres(data)
|
|
|
|
# Verify notification handler wasn't called
|
|
postgres_client.notification_handler.build_and_send_notification.assert_not_called()
|
|
|
|
# Verify session handling
|
|
mock_dataframe.assert_called_once_with(data["data"])
|
|
mock_dataframe.return_value.to_sql.assert_called_once_with(
|
|
data["table_name"],
|
|
postgres_client.engine,
|
|
schema=data["schema"],
|
|
if_exists="append",
|
|
index=False
|
|
)
|
|
postgres_client.session_factory.return_value.commit.assert_called_once()
|
|
postgres_client.session_factory.return_value.close.assert_called_once()
|
|
|
|
|
|
@mark.asyncio
|
|
@patch("scouter.activities.postgres.DataFrame", return_value=MagicMock(
|
|
to_sql=MagicMock(side_effect=Exception("Error exporting data to postgres"))
|
|
))
|
|
async def test_export_data_to_postgres_error(_mock_dataframe, postgres_client):
|
|
data = {"schema": "test", "table_name": "test",
|
|
"data": {"a": [1, 2, 3], "b": [4, 5, 6]}}
|
|
await postgres_client.export_data_to_postgres(data)
|
|
|
|
# Verify error notification was sent
|
|
postgres_client.notification_handler.build_and_send_notification.assert_called_once_with(
|
|
notification_id="ERROR_EXPORTING_DATA_TO_POSTGRES",
|
|
message="Error exporting data to postgres: Error exporting data to postgres",
|
|
block="export_data_to_postgres",
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=ANY
|
|
)
|
|
|
|
# Verify session handling
|
|
postgres_client.session_factory.return_value.close.assert_called_once()
|