vitor-aignosi
ad00661516
Refactor tests for Postgres activities and improve error handling - Updated test_postgres.py to enhance the testing of load_custom_query method, including cases for None data and date conversion. - Refactored repeat_last_prediction tests to use mocks for SQLAlchemy session execution. - Added tests for export_data_to_postgres method, covering both success and error scenarios. - Improved the initialization tests for Activities class to ensure proper instantiation of dependencies. - Enhanced test coverage for OPC repository connection validation. - Updated tests for prediction workflows to streamline input handling and improve clarity. - Introduced tests for connectors configuration to validate environment variable handling for MLFlow, OPC, and Postgres. - Added tests for logger utility to ensure default settings are correctly applied.
150 lines
4.4 KiB
Python
150 lines
4.4 KiB
Python
from pytest import mark
|
|
from unittest.mock import patch, MagicMock, ANY
|
|
from laborious.activities.activities import Activities
|
|
from laborious.activities.postgres import Postgres
|
|
from laborious.activities.mlflow import MLFlow
|
|
from laborious.activities.gates import Gates
|
|
from laborious.activities.opc import OPC
|
|
|
|
|
|
@patch('laborious.activities.activities.Postgres.__init__')
|
|
@patch('laborious.activities.activities.MLFlow.__init__')
|
|
@patch('laborious.activities.activities.OPC.__init__')
|
|
@patch('laborious.activities.activities.Gates.__init__')
|
|
def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgres_init):
|
|
|
|
postgres_config = {
|
|
'host': 'localhost',
|
|
'port': 5432,
|
|
'user': 'postgres',
|
|
'password': 'postgres',
|
|
'dbname': 'postgres',
|
|
'min_connections': 1,
|
|
'max_connections': 10
|
|
}
|
|
|
|
mlflow_config = {
|
|
'host': 'localhost',
|
|
'port': 5000,
|
|
'username': 'mlflow',
|
|
'password': 'mlflow'
|
|
}
|
|
|
|
opc_config = {
|
|
'bootstrap_servers': 'localhost:9092',
|
|
'polling_time': 1000,
|
|
'group_id': 'test-group'
|
|
}
|
|
|
|
logger = MagicMock()
|
|
notification_handler = MagicMock()
|
|
|
|
activities = Activities(
|
|
postgres_config=postgres_config,
|
|
mlflow_config=mlflow_config,
|
|
opc_config=opc_config,
|
|
logger=logger,
|
|
notification_handler=notification_handler
|
|
)
|
|
|
|
assert isinstance(activities, Activities)
|
|
assert isinstance(activities, Postgres)
|
|
assert isinstance(activities, MLFlow)
|
|
assert isinstance(activities, OPC)
|
|
assert isinstance(activities, Gates)
|
|
|
|
mock_postgres_init.assert_called_once_with(
|
|
ANY,
|
|
host=postgres_config['host'],
|
|
port=postgres_config['port'],
|
|
user=postgres_config['user'],
|
|
password=postgres_config['password'],
|
|
dbname=postgres_config['dbname'],
|
|
min_connections=postgres_config['min_connections'],
|
|
max_connections=postgres_config['max_connections'],
|
|
logger=logger,
|
|
notification_handler=notification_handler
|
|
)
|
|
|
|
mock_mlflow_init.assert_called_once_with(
|
|
ANY,
|
|
mlflow_host=mlflow_config['host'],
|
|
mlflow_port=mlflow_config['port'],
|
|
mlflow_username=mlflow_config['username'],
|
|
mlflow_password=mlflow_config['password'],
|
|
logger=logger,
|
|
notification_handler=notification_handler
|
|
)
|
|
|
|
mock_opc_init.assert_called_once_with(
|
|
ANY,
|
|
opc_servers=opc_config,
|
|
logger=logger,
|
|
notification_handler=notification_handler
|
|
)
|
|
|
|
mock_gates_init.assert_called_once_with(
|
|
ANY,
|
|
logger=logger,
|
|
notification_handler=notification_handler
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.activities.activities.Postgres.__init__')
|
|
@patch('laborious.activities.activities.MLFlow.__init__')
|
|
@patch('laborious.activities.activities.OPC.__init__')
|
|
async def test_prepare_activity(_mock_opc_init,
|
|
_mock_mlflow_init, _mock_postgres_init):
|
|
postgres_config = {
|
|
'host': 'localhost',
|
|
'port': 5432,
|
|
'user': 'postgres',
|
|
'password': 'postgres',
|
|
'dbname': 'postgres',
|
|
'min_connections': 1,
|
|
'max_connections': 10
|
|
}
|
|
|
|
mlflow_config = {
|
|
'host': 'localhost',
|
|
'port': 5000,
|
|
'username': 'mlflow',
|
|
'password': 'mlflow'
|
|
}
|
|
|
|
opc_config = {
|
|
'bootstrap_servers': 'localhost:9092',
|
|
'polling_time': 1000,
|
|
'group_id': 'test-group'
|
|
}
|
|
|
|
logger = MagicMock()
|
|
notification_handler = MagicMock()
|
|
|
|
activities = Activities(
|
|
postgres_config=postgres_config,
|
|
mlflow_config=mlflow_config,
|
|
opc_config=opc_config,
|
|
logger=logger,
|
|
notification_handler=notification_handler
|
|
)
|
|
|
|
input_data = {
|
|
'workflow_name': 'test-workflow-name',
|
|
'schedule_name': 'test-schedule-name',
|
|
'model_name': 'test-model-name',
|
|
'model_id': 'test-model-id'
|
|
}
|
|
|
|
await activities.prepare_activity(input_data)
|
|
|
|
assert activities.notification_handler.base_notification.pipeline_name == input_data[
|
|
'workflow_name']
|
|
assert activities.notification_handler.base_notification.schedule_name == input_data[
|
|
'schedule_name']
|
|
assert activities.notification_handler.base_notification.model_name == input_data[
|
|
'model_name']
|
|
assert activities.notification_handler.base_notification.model_id == input_data[
|
|
'model_id']
|