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.
160 lines
4.9 KiB
Python
160 lines
4.9 KiB
Python
from unittest.mock import MagicMock, patch
|
|
from pytest import fixture, mark
|
|
import pandas as pd
|
|
from laborious.activities.postgres import Postgres
|
|
|
|
|
|
@fixture
|
|
@patch("laborious.activities.postgres.create_engine")
|
|
def postgres_activity(_mock_create_engine):
|
|
return Postgres(
|
|
host="localhost",
|
|
port=5432,
|
|
user="test_user",
|
|
password="test_password",
|
|
dbname="test_db",
|
|
min_connections=1,
|
|
max_connections=5,
|
|
logger=MagicMock(),
|
|
notification_handler=MagicMock()
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
@patch("laborious.activities.postgres.read_sql_query")
|
|
async def test_load_custom_query_none_data(mock_read_sql_query, postgres_activity):
|
|
query = "SELECT * FROM test_table LIMIT 1"
|
|
mock_read_sql_query.return_value = None
|
|
|
|
result = await postgres_activity.load_custom_query(query)
|
|
|
|
assert isinstance(result, dict)
|
|
assert len(result) == 0
|
|
|
|
|
|
@mark.asyncio
|
|
@patch("laborious.activities.postgres.read_sql_query")
|
|
async def test_load_custom_query_date_converted(mock_read_sql_query, postgres_activity):
|
|
query = "SELECT * FROM test_table LIMIT 1"
|
|
mock_data = pd.DataFrame({"column1": [1], "column2": ["test"]})
|
|
mock_data['date'] = pd.to_datetime('2022-01-01')
|
|
|
|
mock_read_sql_query.return_value = mock_data
|
|
|
|
result = await postgres_activity.load_custom_query(query)
|
|
|
|
assert isinstance(result, dict)
|
|
assert len(result) == 3
|
|
assert "column1" in result
|
|
assert "column2" in result
|
|
assert "date" in result
|
|
assert result['date'] == {0: '2022-01-01 00:00:00'}
|
|
|
|
|
|
@mark.asyncio
|
|
@patch("laborious.activities.postgres.read_sql_query")
|
|
async def test_load_custom_query_success(mock_read_sql_query, postgres_activity):
|
|
query = "SELECT * FROM test_table LIMIT 1"
|
|
mock_data = pd.DataFrame({"column1": [1], "column2": ["test"]})
|
|
|
|
mock_read_sql_query.return_value = mock_data
|
|
|
|
result = await postgres_activity.load_custom_query(query)
|
|
|
|
assert isinstance(result, dict)
|
|
assert len(result) == 2
|
|
assert "column1" in result
|
|
assert "column2" in result
|
|
postgres_activity.logger.info.assert_called()
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_load_custom_query_error(postgres_activity):
|
|
query = "SELECT * FROM non_existent_table"
|
|
error_msg = "Table not found"
|
|
|
|
with patch("laborious.activities.postgres.read_sql_query", side_effect=ValueError(error_msg)):
|
|
result = await postgres_activity.load_custom_query(query)
|
|
|
|
assert isinstance(result, dict)
|
|
assert len(result) == 0
|
|
postgres_activity.notification_handler.build_and_send_notification.assert_called_once()
|
|
postgres_activity.logger.error.assert_called()
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_repeat_last_prediction_success(postgres_activity):
|
|
query_items = {
|
|
"schema": "public",
|
|
"table_name": "predictions",
|
|
"model": 1
|
|
}
|
|
|
|
with patch("sqlalchemy.orm.session.Session.execute") as mock_execute:
|
|
await postgres_activity.repeat_last_prediction(query_items)
|
|
|
|
mock_execute.assert_called_once()
|
|
postgres_activity.logger.info.assert_called()
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_repeat_last_prediction_error(postgres_activity):
|
|
query_items = {
|
|
"schema": "public",
|
|
"table_name": "predictions",
|
|
"model": 1
|
|
}
|
|
error_msg = "Database error"
|
|
|
|
with patch("sqlalchemy.orm.session.Session.execute", side_effect=ValueError(error_msg)):
|
|
await postgres_activity.repeat_last_prediction(query_items)
|
|
|
|
postgres_activity.notification_handler.build_and_send_notification.assert_called_once()
|
|
postgres_activity.logger.error.assert_called()
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_export_data_to_postgres_success(postgres_activity):
|
|
input_data = {
|
|
"schema": "public",
|
|
"table_name": "test_table",
|
|
"data": pd.DataFrame({"column1": [1, 2], "column2": ["a", "b"]})
|
|
}
|
|
|
|
with patch("laborious.activities.postgres.DataFrame.to_sql") as mock_to_sql:
|
|
await postgres_activity.export_data_to_postgres(input_data)
|
|
|
|
mock_to_sql.assert_called_once()
|
|
postgres_activity.logger.debug.assert_called()
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_export_data_to_postgres_error(postgres_activity):
|
|
input_data = {
|
|
"schema": "public",
|
|
"table_name": "test_table",
|
|
"data": pd.DataFrame({"column1": [1, 2], "column2": ["a", "b"]})
|
|
}
|
|
error_msg = "Export failed"
|
|
|
|
with patch("laborious.activities.postgres.DataFrame.to_sql", side_effect=ValueError(error_msg)):
|
|
await postgres_activity.export_data_to_postgres(input_data)
|
|
|
|
postgres_activity.notification_handler.build_and_send_notification.assert_called_once()
|
|
postgres_activity.logger.error.assert_called()
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_close(postgres_activity):
|
|
postgres_activity.close()
|
|
|
|
postgres_activity.engine.dispose.assert_called_once()
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_del(postgres_activity):
|
|
postgres_activity.close = MagicMock()
|
|
postgres_activity.__del__()
|
|
|
|
postgres_activity.close.assert_called_once()
|