vitor-aignosi
e7f214b144
Remove unused utility files and update requirements.txt to include new dependencies for data processing and database interaction.
133 lines
5.5 KiB
Python
133 lines
5.5 KiB
Python
from unittest.mock import ANY, MagicMock, patch
|
|
from pandas import DataFrame
|
|
from pytest import fixture
|
|
from pytest import mark
|
|
from sientia_do.notifications.models import NotificationLevel
|
|
|
|
from laborious.activities.postgres import Postgres
|
|
|
|
|
|
@fixture
|
|
@patch("laborious.activities.postgres.ThreadedConnectionPool")
|
|
def postgres_client(mock_pool):
|
|
return Postgres(
|
|
host="localhost",
|
|
port=5432,
|
|
user="postgres",
|
|
password="postgres",
|
|
dbname="postgres",
|
|
min_connections=1,
|
|
max_connections=10,
|
|
logger=MagicMock(),
|
|
notification_handler=MagicMock(),
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
@patch("laborious.activities.postgres.read_sql_query",
|
|
return_value=DataFrame([{"a": 1, "b": 2}]))
|
|
async def test_load_custom_query_success(mock_read_sql_query, postgres_client):
|
|
query = "SELECT * FROM test"
|
|
result = await postgres_client.load_custom_query(query)
|
|
assert result is not None
|
|
assert len(result) > 0
|
|
assert result == {'a': {0: 1}, 'b': {0: 2}}
|
|
postgres_client.notification_handler.build_and_send_notification.assert_not_called()
|
|
|
|
|
|
@mark.asyncio
|
|
@patch("laborious.activities.postgres.read_sql_query",
|
|
side_effect=Exception("Error fetching data from query"))
|
|
async def test_load_custom_query_error(mock_read_sql_query, postgres_client):
|
|
query = "SELECT * FROM test"
|
|
result = await postgres_client.load_custom_query(query)
|
|
assert result == {}
|
|
postgres_client.notification_handler.build_and_send_notification.assert_called_once_with(
|
|
notification_id="ERROR_LOADING_CUSTOM_QUERY",
|
|
message="Error fetching data from query: Error fetching data from query",
|
|
block="load_custom_query",
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=ANY
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_repeat_last_prediction_success(postgres_client):
|
|
query_items = {"schema": "test", "table_name": "test", "model": 1}
|
|
await postgres_client.repeat_last_prediction(query_items)
|
|
postgres_client.notification_handler.build_and_send_notification.assert_not_called()
|
|
postgres_client.pool.getconn.assert_called_once()
|
|
postgres_client.pool.putconn.assert_called_once()
|
|
|
|
postgres_client.pool.getconn.return_value.cursor.assert_called_once()
|
|
postgres_client.pool.getconn.return_value.cursor.return_value.execute.assert_called_once_with(
|
|
f"""
|
|
INSERT INTO \"{query_items['schema']}\".{query_items['table_name']} (model_id, prediction, timestamp, response_time, prediction_status, prediction_confidence, created_at)
|
|
SELECT model_id, prediction, timestamp, response_time, prediction_status, prediction_confidence, NOW()
|
|
FROM \"{query_items['schema']}\".{query_items['table_name']}
|
|
WHERE model_id = {query_items['model']}
|
|
ORDER BY timestamp DESC
|
|
LIMIT 1;
|
|
"""
|
|
)
|
|
postgres_client.pool.getconn.return_value.commit.assert_called_once()
|
|
postgres_client.pool.getconn.return_value.cursor.return_value.close.assert_called_once()
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_repeat_last_prediction_error(postgres_client):
|
|
postgres_client.pool.getconn.return_value.cursor.return_value.execute.side_effect = Exception(
|
|
"Error repeating last prediction")
|
|
query_items = {"schema": "test", "table_name": "test", "model": 1}
|
|
await postgres_client.repeat_last_prediction(query_items)
|
|
postgres_client.notification_handler.build_and_send_notification.assert_called_once_with(
|
|
notification_id="ERROR_REPEATING_LAST_PREDICTION",
|
|
message="Error repeating last prediction: Error repeating last prediction",
|
|
block="repeat_last_prediction",
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=ANY
|
|
)
|
|
postgres_client.pool.getconn.assert_called_once()
|
|
postgres_client.pool.putconn.assert_called_once()
|
|
|
|
|
|
@mark.asyncio
|
|
@patch("laborious.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)
|
|
postgres_client.notification_handler.build_and_send_notification.assert_not_called()
|
|
postgres_client.pool.getconn.assert_called_once()
|
|
postgres_client.pool.putconn.assert_called_once()
|
|
|
|
mock_dataframe.assert_called_once_with(data["data"])
|
|
mock_dataframe.return_value.to_sql.assert_called_once_with(
|
|
data["table_name"],
|
|
postgres_client.pool.getconn.return_value,
|
|
schema=data["schema"],
|
|
if_exists="append",
|
|
index=False
|
|
)
|
|
postgres_client.pool.getconn.return_value.commit.assert_called_once()
|
|
|
|
|
|
@mark.asyncio
|
|
@patch("laborious.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)
|
|
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
|
|
)
|
|
|
|
postgres_client.pool.getconn.assert_called_once()
|
|
postgres_client.pool.putconn.assert_called_once()
|