SIENTIAPDE-994
Remove unused utility files and update requirements.txt to include new dependencies for data processing and database interaction.
This commit is contained in:
17
laborious/activities/base.py
Normal file
17
laborious/activities/base.py
Normal file
@@ -0,0 +1,17 @@
|
|||||||
|
from temporalio import activity
|
||||||
|
from sientia_do.notifications.handlers import NotificationHandler
|
||||||
|
from logging import Logger
|
||||||
|
|
||||||
|
|
||||||
|
class BaseActivity:
|
||||||
|
def __init__(self, logger: Logger, notification_handler: NotificationHandler):
|
||||||
|
self.logger = logger
|
||||||
|
self.notification_handler = notification_handler
|
||||||
|
|
||||||
|
@activity.defn(name="prepare_notification_handler")
|
||||||
|
async def prepare_notification_handler(self, schedule_name: str,
|
||||||
|
model_name: str,
|
||||||
|
model_id: str):
|
||||||
|
self.notification_handler.base_notification.schedule_name = schedule_name
|
||||||
|
self.notification_handler.base_notification.model_name = model_name
|
||||||
|
self.notification_handler.base_notification.model_id = model_id
|
||||||
47
laborious/activities/gates.py
Normal file
47
laborious/activities/gates.py
Normal file
@@ -0,0 +1,47 @@
|
|||||||
|
from temporalio import activity, workflow
|
||||||
|
|
||||||
|
with workflow.unsafe.imports_passed_through():
|
||||||
|
from logging import Logger
|
||||||
|
from sientia_do.notifications.handlers import NotificationHandler
|
||||||
|
from laborious.activities.base import BaseActivity
|
||||||
|
from typing import Any
|
||||||
|
from laborious.utils.filters.conditional_filters import filter_empty_data, filter_specific_variables_null_values
|
||||||
|
from pandas import DataFrame
|
||||||
|
|
||||||
|
|
||||||
|
filter_functions = {
|
||||||
|
'SPECIFIC_VARIABLES_NULL_VALUES': filter_specific_variables_null_values,
|
||||||
|
'EMPTY_DATA': filter_empty_data
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
class Gates(BaseActivity):
|
||||||
|
def __init__(self, logger: Logger, notification_handler: NotificationHandler):
|
||||||
|
super().__init__(logger, notification_handler)
|
||||||
|
|
||||||
|
@activity.defn(name="input_gate")
|
||||||
|
async def input_gate(self, input_data: dict[str, Any]) -> tuple[str, int]:
|
||||||
|
"""
|
||||||
|
Filters the data based on the filters. The return value is a tuple with the first element
|
||||||
|
being the policy and the second element being the confidence status.
|
||||||
|
Args:
|
||||||
|
input_data (dict): The input data.
|
||||||
|
Returns:
|
||||||
|
tuple[str, int]: ('stop', -1) if some filter policy is 'stop', ('continue', 2)
|
||||||
|
if no filter policy is 'stop' and some filter policy is 'continue',
|
||||||
|
None if no filter is applied.
|
||||||
|
"""
|
||||||
|
filters = input_data['filters']
|
||||||
|
data = DataFrame(input_data['data'])
|
||||||
|
|
||||||
|
filter_output = []
|
||||||
|
for fil, config in filters.items():
|
||||||
|
if filter_functions[fil](data, config):
|
||||||
|
filter_output.append(config['POLICY'])
|
||||||
|
|
||||||
|
if 'stop' in filter_output:
|
||||||
|
return 'stop', -1
|
||||||
|
elif 'continue' in filter_output:
|
||||||
|
return 'continue', 2
|
||||||
|
|
||||||
|
return None, 0
|
||||||
0
laborious/activities/mlflow.py
Normal file
0
laborious/activities/mlflow.py
Normal file
0
laborious/activities/opc.py
Normal file
0
laborious/activities/opc.py
Normal file
162
laborious/activities/postgres.py
Normal file
162
laborious/activities/postgres.py
Normal file
@@ -0,0 +1,162 @@
|
|||||||
|
import traceback
|
||||||
|
from temporalio import workflow, activity
|
||||||
|
|
||||||
|
from laborious.activities.base import BaseActivity
|
||||||
|
with workflow.unsafe.imports_passed_through():
|
||||||
|
from psycopg2.pool import ThreadedConnectionPool
|
||||||
|
from pandas import read_sql_query, DataFrame
|
||||||
|
from logging import Logger
|
||||||
|
from sientia_do.notifications.handlers import NotificationHandler
|
||||||
|
from sientia_do.notifications.models import NotificationLevel
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
|
||||||
|
class Postgres(BaseActivity):
|
||||||
|
def __init__(self, host: str, port: int,
|
||||||
|
user: str, password: str, dbname: str,
|
||||||
|
min_connections: int, max_connections: int,
|
||||||
|
logger: Logger, notification_handler: NotificationHandler):
|
||||||
|
self.host = host
|
||||||
|
self.port = port
|
||||||
|
self.user = user
|
||||||
|
self.password = password
|
||||||
|
self.dbname = dbname
|
||||||
|
|
||||||
|
self.pool = ThreadedConnectionPool(
|
||||||
|
minconn=min_connections,
|
||||||
|
maxconn=max_connections,
|
||||||
|
host=self.host,
|
||||||
|
port=self.port,
|
||||||
|
user=self.user,
|
||||||
|
password=self.password,
|
||||||
|
dbname=self.dbname)
|
||||||
|
|
||||||
|
super().__init__(logger, notification_handler)
|
||||||
|
|
||||||
|
def close(self):
|
||||||
|
self.pool.closeall()
|
||||||
|
|
||||||
|
def __del__(self):
|
||||||
|
self.close()
|
||||||
|
|
||||||
|
@activity.defn(name="load_custom_query")
|
||||||
|
async def load_custom_query(self, query: str) -> dict[str, dict]:
|
||||||
|
"""
|
||||||
|
Loads data from a custom query.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
query (str): The query to load data from.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
dict[str, dict]: The data from the query.
|
||||||
|
"""
|
||||||
|
self.logger.info(f"Fetching data from query: {query}")
|
||||||
|
|
||||||
|
conn = self.pool.getconn()
|
||||||
|
try:
|
||||||
|
data = read_sql_query(query, conn)
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
trace = traceback.format_exc()
|
||||||
|
self.notification_handler.build_and_send_notification(
|
||||||
|
notification_id="ERROR_LOADING_CUSTOM_QUERY",
|
||||||
|
message=f"Error fetching data from query: {e}",
|
||||||
|
block="load_custom_query",
|
||||||
|
level=NotificationLevel.ERROR,
|
||||||
|
attachment_content=trace
|
||||||
|
)
|
||||||
|
|
||||||
|
self.logger.error(trace)
|
||||||
|
|
||||||
|
return {}
|
||||||
|
finally:
|
||||||
|
self.pool.putconn(conn)
|
||||||
|
|
||||||
|
self.logger.info(f"Fetched {len(data)} rows")
|
||||||
|
self.logger.debug(f"Data: {data.to_string()}")
|
||||||
|
|
||||||
|
return data.to_dict()
|
||||||
|
|
||||||
|
@activity.defn(name="repeat_last_prediction")
|
||||||
|
async def repeat_last_prediction(self, query_items: dict[str, str]):
|
||||||
|
"""
|
||||||
|
Repeats the last prediction for a given model.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
query_items (dict[str, str]): The query items.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
None
|
||||||
|
"""
|
||||||
|
schema = query_items["schema"]
|
||||||
|
table_name = query_items["table_name"]
|
||||||
|
model = query_items["model"]
|
||||||
|
|
||||||
|
repeat_query = f"""
|
||||||
|
INSERT INTO \"{schema}\".{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 \"{schema}\".{table_name}
|
||||||
|
WHERE model_id = {model}
|
||||||
|
ORDER BY timestamp DESC
|
||||||
|
LIMIT 1;
|
||||||
|
"""
|
||||||
|
self.logger.info(f"Repeating last prediction for model {model}")
|
||||||
|
self.logger.debug(f"Query: {repeat_query}")
|
||||||
|
|
||||||
|
conn = self.pool.getconn()
|
||||||
|
|
||||||
|
try:
|
||||||
|
cursor = conn.cursor()
|
||||||
|
cursor.execute(repeat_query)
|
||||||
|
conn.commit()
|
||||||
|
cursor.close()
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
trace = traceback.format_exc()
|
||||||
|
self.notification_handler.build_and_send_notification(
|
||||||
|
notification_id="ERROR_REPEATING_LAST_PREDICTION",
|
||||||
|
message=f"Error repeating last prediction: {e}",
|
||||||
|
block="repeat_last_prediction",
|
||||||
|
level=NotificationLevel.ERROR,
|
||||||
|
attachment_content=trace
|
||||||
|
)
|
||||||
|
|
||||||
|
self.logger.error(trace)
|
||||||
|
|
||||||
|
finally:
|
||||||
|
self.pool.putconn(conn)
|
||||||
|
|
||||||
|
@activity.defn(name="export_data_to_postgres")
|
||||||
|
async def export_data_to_postgres(self, input_data: dict[str, Any]):
|
||||||
|
"""
|
||||||
|
Exports data to a postgres table.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
input_data (dict[str, Any]): The data to export.
|
||||||
|
"""
|
||||||
|
|
||||||
|
schema = input_data["schema"]
|
||||||
|
table_name = input_data["table_name"]
|
||||||
|
data = DataFrame(input_data["data"])
|
||||||
|
|
||||||
|
conn = self.pool.getconn()
|
||||||
|
|
||||||
|
try:
|
||||||
|
data.to_sql(table_name, conn, schema=schema,
|
||||||
|
if_exists="append", index=False)
|
||||||
|
conn.commit()
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
trace = traceback.format_exc()
|
||||||
|
self.notification_handler.build_and_send_notification(
|
||||||
|
notification_id="ERROR_EXPORTING_DATA_TO_POSTGRES",
|
||||||
|
message=f"Error exporting data to postgres: {e}",
|
||||||
|
block="export_data_to_postgres",
|
||||||
|
level=NotificationLevel.ERROR,
|
||||||
|
attachment_content=trace
|
||||||
|
)
|
||||||
|
|
||||||
|
self.logger.error(trace)
|
||||||
|
|
||||||
|
finally:
|
||||||
|
self.pool.putconn(conn)
|
||||||
0
laborious/utils/__init__.py
Normal file
0
laborious/utils/__init__.py
Normal file
0
laborious/utils/filters/__init__.py
Normal file
0
laborious/utils/filters/__init__.py
Normal file
69
laborious/utils/filters/api_filters.py
Normal file
69
laborious/utils/filters/api_filters.py
Normal file
@@ -0,0 +1,69 @@
|
|||||||
|
import numpy as np
|
||||||
|
from pandas import DataFrame
|
||||||
|
from laborious.utils.filters.base_filter import Filter
|
||||||
|
|
||||||
|
|
||||||
|
class ApiErrorFilter(Filter):
|
||||||
|
def __init__(self, policy):
|
||||||
|
self.policy = policy
|
||||||
|
super().__init__('API_FILTER')
|
||||||
|
|
||||||
|
def method(self, response: dict, prediction_confidence: int):
|
||||||
|
"""
|
||||||
|
Processes the API response and determines the next action based on the response and prediction confidence.
|
||||||
|
Args:
|
||||||
|
response (dict): The API response to be processed.
|
||||||
|
prediction_confidence (int): The confidence level of the prediction.
|
||||||
|
Returns:
|
||||||
|
str: 'stop' if the policy is to stop on captured errors, 'continue' if the policy is to continue on captured errors.
|
||||||
|
Raises:
|
||||||
|
KeyError: If 'success' or 'content' keys are missing in the response dictionary.
|
||||||
|
"""
|
||||||
|
|
||||||
|
captured = False
|
||||||
|
if not response and prediction_confidence == 10:
|
||||||
|
self.warning('No valid response.')
|
||||||
|
captured = True
|
||||||
|
|
||||||
|
else:
|
||||||
|
if not response['success']:
|
||||||
|
message = response['content']["message"]
|
||||||
|
self.warning(
|
||||||
|
f'Model repository error: {message}')
|
||||||
|
captured = True
|
||||||
|
if captured and self.policy == 'stop':
|
||||||
|
return 'stop'
|
||||||
|
elif captured and self.policy == 'continue':
|
||||||
|
return 'continue'
|
||||||
|
|
||||||
|
|
||||||
|
class NaNValuesFilter(Filter):
|
||||||
|
def __init__(self, policy):
|
||||||
|
self.policy = policy
|
||||||
|
super().__init__('NAN_VALUES')
|
||||||
|
|
||||||
|
def method(self, predictions: DataFrame, prediction_confidence: int):
|
||||||
|
"""
|
||||||
|
Processes the given predictions DataFrame by replacing None values with NaN,
|
||||||
|
dropping the 'timestamp' column if it exists, and checking for NaN values.
|
||||||
|
Args:
|
||||||
|
predictions (pd.DataFrame): The DataFrame containing prediction data.
|
||||||
|
prediction_confidence (float): The confidence level of the predictions.
|
||||||
|
Returns:
|
||||||
|
float or int or bool: Returns the prediction confidence if the DataFrame
|
||||||
|
is not entirely NaN. If all values are NaN and the policy is 'stop',
|
||||||
|
returns False. If all values are NaN and the policy is 'continue',
|
||||||
|
returns 18.
|
||||||
|
"""
|
||||||
|
|
||||||
|
data = predictions.replace({None: np.nan}).drop(
|
||||||
|
columns=['timestamp'], errors='ignore')
|
||||||
|
|
||||||
|
if data.isna().all().all():
|
||||||
|
if self.policy == 'stop':
|
||||||
|
self.warning('All values are NaN.')
|
||||||
|
return False
|
||||||
|
elif self.policy == 'continue':
|
||||||
|
return 18
|
||||||
|
|
||||||
|
return prediction_confidence
|
||||||
13
laborious/utils/filters/base_filter.py
Normal file
13
laborious/utils/filters/base_filter.py
Normal file
@@ -0,0 +1,13 @@
|
|||||||
|
from pandas import DataFrame
|
||||||
|
|
||||||
|
class Filter:
|
||||||
|
def __init__(self, id: str):
|
||||||
|
self.id = id
|
||||||
|
self.warnings = []
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def method(df: DataFrame) -> DataFrame:
|
||||||
|
raise NotImplementedError
|
||||||
|
|
||||||
|
def warning(self, message: str):
|
||||||
|
self.warnings.append(f'[{self.id}] - {message}')
|
||||||
19
laborious/utils/filters/conditional_filters.py
Normal file
19
laborious/utils/filters/conditional_filters.py
Normal file
@@ -0,0 +1,19 @@
|
|||||||
|
from typing import List
|
||||||
|
|
||||||
|
from laborious.utils.filters.base_filter import Filter
|
||||||
|
from pandas import DataFrame
|
||||||
|
|
||||||
|
|
||||||
|
def filter_specific_variables_null_values(data: DataFrame, config: dict) -> bool:
|
||||||
|
"""
|
||||||
|
Returns True if the data is empty, False otherwise.
|
||||||
|
"""
|
||||||
|
return data[
|
||||||
|
data['variable'].isin(config['VARIABLES']) & data['value'].isna()].empty
|
||||||
|
|
||||||
|
|
||||||
|
def filter_empty_data(data: DataFrame, _config: dict) -> bool:
|
||||||
|
"""
|
||||||
|
Returns True if the data is empty, False otherwise.
|
||||||
|
"""
|
||||||
|
return data.empty
|
||||||
@@ -1,4 +1,4 @@
|
|||||||
GitPython
|
|
||||||
temporalio
|
temporalio
|
||||||
pytest
|
psycopg2-binary
|
||||||
python-dotenv
|
pandas
|
||||||
|
git+ssh://git@github.com/Aignosi/sientia-dataops-library.git
|
||||||
0
tests/__init__.py
Normal file
0
tests/__init__.py
Normal file
0
tests/laborious/__init__.py
Normal file
0
tests/laborious/__init__.py
Normal file
0
tests/laborious/activities/__init__.py
Normal file
0
tests/laborious/activities/__init__.py
Normal file
274
tests/laborious/activities/test_gates.py
Normal file
274
tests/laborious/activities/test_gates.py
Normal file
@@ -0,0 +1,274 @@
|
|||||||
|
from unittest.mock import MagicMock, patch
|
||||||
|
from pandas import DataFrame
|
||||||
|
from pytest import fixture, mark
|
||||||
|
|
||||||
|
from laborious.activities.gates import Gates
|
||||||
|
|
||||||
|
|
||||||
|
@fixture
|
||||||
|
def gates():
|
||||||
|
return Gates(
|
||||||
|
logger=MagicMock(),
|
||||||
|
notification_handler=MagicMock()
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch('laborious.activities.gates.filter_functions')
|
||||||
|
async def test_input_gate_specific_variables_null_values_with_stop_policy_only(
|
||||||
|
filter_functions_mock,
|
||||||
|
gates
|
||||||
|
):
|
||||||
|
specific_variables_null_values_mock = MagicMock(return_value=True)
|
||||||
|
empty_data_mock = MagicMock(return_value=False)
|
||||||
|
|
||||||
|
def functions_side_effect(x):
|
||||||
|
if x == 'SPECIFIC_VARIABLES_NULL_VALUES':
|
||||||
|
return specific_variables_null_values_mock
|
||||||
|
return empty_data_mock
|
||||||
|
|
||||||
|
filter_functions_mock.__getitem__.side_effect = functions_side_effect
|
||||||
|
|
||||||
|
input_data = {
|
||||||
|
'filters': {
|
||||||
|
'SPECIFIC_VARIABLES_NULL_VALUES': {
|
||||||
|
'POLICY': 'stop',
|
||||||
|
'VARIABLES': ['variable2']
|
||||||
|
}
|
||||||
|
},
|
||||||
|
'data': {
|
||||||
|
'variable': ['variable1', 'variable2'],
|
||||||
|
'value': [1, 2]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
result = await gates.input_gate(input_data)
|
||||||
|
assert result == ('stop', -1)
|
||||||
|
|
||||||
|
input_args = specific_variables_null_values_mock.call_args
|
||||||
|
assert input_args[0][0].equals(DataFrame(
|
||||||
|
{'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
|
||||||
|
assert input_args[0][1] == input_data['filters']['SPECIFIC_VARIABLES_NULL_VALUES']
|
||||||
|
|
||||||
|
empty_data_mock.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch('laborious.activities.gates.filter_functions')
|
||||||
|
async def test_input_gate_specific_variables_null_values_with_continue_policy_only(
|
||||||
|
filter_functions_mock,
|
||||||
|
gates
|
||||||
|
):
|
||||||
|
specific_variables_null_values_mock = MagicMock(return_value=True)
|
||||||
|
empty_data_mock = MagicMock(return_value=False)
|
||||||
|
|
||||||
|
def functions_side_effect(x):
|
||||||
|
if x == 'SPECIFIC_VARIABLES_NULL_VALUES':
|
||||||
|
return specific_variables_null_values_mock
|
||||||
|
return empty_data_mock
|
||||||
|
|
||||||
|
filter_functions_mock.__getitem__.side_effect = functions_side_effect
|
||||||
|
|
||||||
|
input_data = {
|
||||||
|
'filters': {
|
||||||
|
'SPECIFIC_VARIABLES_NULL_VALUES': {
|
||||||
|
'POLICY': 'continue',
|
||||||
|
'VARIABLES': ['variable2']
|
||||||
|
}
|
||||||
|
},
|
||||||
|
'data': {
|
||||||
|
'variable': ['variable1', 'variable2'],
|
||||||
|
'value': [1, 2]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
result = await gates.input_gate(input_data)
|
||||||
|
assert result == ('continue', 2)
|
||||||
|
|
||||||
|
input_args = specific_variables_null_values_mock.call_args
|
||||||
|
assert input_args[0][0].equals(DataFrame(
|
||||||
|
{'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
|
||||||
|
assert input_args[0][1] == input_data['filters']['SPECIFIC_VARIABLES_NULL_VALUES']
|
||||||
|
|
||||||
|
empty_data_mock.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch('laborious.activities.gates.filter_functions')
|
||||||
|
async def test_input_gate_specific_variables_null_values_no_filtered(
|
||||||
|
filter_functions_mock,
|
||||||
|
gates
|
||||||
|
):
|
||||||
|
specific_variables_null_values_mock = MagicMock(return_value=False)
|
||||||
|
empty_data_mock = MagicMock(return_value=False)
|
||||||
|
|
||||||
|
def functions_side_effect(x):
|
||||||
|
if x == 'SPECIFIC_VARIABLES_NULL_VALUES':
|
||||||
|
return specific_variables_null_values_mock
|
||||||
|
return empty_data_mock
|
||||||
|
|
||||||
|
filter_functions_mock.__getitem__.side_effect = functions_side_effect
|
||||||
|
|
||||||
|
input_data = {
|
||||||
|
'filters': {
|
||||||
|
'SPECIFIC_VARIABLES_NULL_VALUES': {
|
||||||
|
'POLICY': 'stop',
|
||||||
|
'VARIABLES': ['variable2']
|
||||||
|
}
|
||||||
|
},
|
||||||
|
'data': {
|
||||||
|
'variable': ['variable1', 'variable2'],
|
||||||
|
'value': [1, 2]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
result = await gates.input_gate(input_data)
|
||||||
|
assert result == (None, 0)
|
||||||
|
|
||||||
|
specific_variables_null_values_input_args = specific_variables_null_values_mock.call_args
|
||||||
|
|
||||||
|
assert specific_variables_null_values_input_args[0][0].equals(DataFrame(
|
||||||
|
{'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
|
||||||
|
assert specific_variables_null_values_input_args[0][
|
||||||
|
1] == input_data['filters']['SPECIFIC_VARIABLES_NULL_VALUES']
|
||||||
|
|
||||||
|
empty_data_mock.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch('laborious.activities.gates.filter_functions')
|
||||||
|
async def test_input_gate_one_stop_policy(
|
||||||
|
filter_functions_mock,
|
||||||
|
gates
|
||||||
|
):
|
||||||
|
specific_variables_null_values_mock = MagicMock(return_value=True)
|
||||||
|
empty_data_mock = MagicMock(return_value=True)
|
||||||
|
|
||||||
|
def functions_side_effect(x):
|
||||||
|
if x == 'SPECIFIC_VARIABLES_NULL_VALUES':
|
||||||
|
return specific_variables_null_values_mock
|
||||||
|
return empty_data_mock
|
||||||
|
|
||||||
|
filter_functions_mock.__getitem__.side_effect = functions_side_effect
|
||||||
|
|
||||||
|
input_data = {
|
||||||
|
'filters': {
|
||||||
|
'SPECIFIC_VARIABLES_NULL_VALUES': {
|
||||||
|
'POLICY': 'stop',
|
||||||
|
'VARIABLES': ['variable2']
|
||||||
|
},
|
||||||
|
'EMPTY_DATA': {
|
||||||
|
'POLICY': 'continue',
|
||||||
|
}
|
||||||
|
},
|
||||||
|
'data': {
|
||||||
|
'variable': ['variable1', 'variable2'],
|
||||||
|
'value': [1, 2]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
result = await gates.input_gate(input_data)
|
||||||
|
assert result == ('stop', -1)
|
||||||
|
|
||||||
|
specific_variables_null_values_input_args = specific_variables_null_values_mock.call_args
|
||||||
|
assert specific_variables_null_values_input_args[0][0].equals(DataFrame(
|
||||||
|
{'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
|
||||||
|
assert specific_variables_null_values_input_args[0][
|
||||||
|
1] == input_data['filters']['SPECIFIC_VARIABLES_NULL_VALUES']
|
||||||
|
|
||||||
|
empty_data_input_args = empty_data_mock.call_args
|
||||||
|
assert empty_data_input_args[0][0].equals(DataFrame(
|
||||||
|
{'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
|
||||||
|
assert empty_data_input_args[0][1] == input_data['filters']['EMPTY_DATA']
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch('laborious.activities.gates.filter_functions')
|
||||||
|
async def test_input_gate_one_continue_policy(
|
||||||
|
filter_functions_mock,
|
||||||
|
gates
|
||||||
|
):
|
||||||
|
specific_variables_null_values_mock = MagicMock(return_value=False)
|
||||||
|
empty_data_mock = MagicMock(return_value=True)
|
||||||
|
|
||||||
|
def functions_side_effect(x):
|
||||||
|
if x == 'SPECIFIC_VARIABLES_NULL_VALUES':
|
||||||
|
return specific_variables_null_values_mock
|
||||||
|
return empty_data_mock
|
||||||
|
|
||||||
|
filter_functions_mock.__getitem__.side_effect = functions_side_effect
|
||||||
|
|
||||||
|
input_data = {
|
||||||
|
'filters': {
|
||||||
|
'SPECIFIC_VARIABLES_NULL_VALUES': {
|
||||||
|
'POLICY': 'stop',
|
||||||
|
'VARIABLES': ['variable2']
|
||||||
|
},
|
||||||
|
'EMPTY_DATA': {
|
||||||
|
'POLICY': 'continue',
|
||||||
|
}
|
||||||
|
},
|
||||||
|
'data': {
|
||||||
|
'variable': ['variable1', 'variable2'],
|
||||||
|
'value': [1, 2]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
result = await gates.input_gate(input_data)
|
||||||
|
assert result == ('continue', 2)
|
||||||
|
|
||||||
|
specific_variables_null_values_input_args = specific_variables_null_values_mock.call_args
|
||||||
|
assert specific_variables_null_values_input_args[0][0].equals(DataFrame(
|
||||||
|
{'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
|
||||||
|
assert specific_variables_null_values_input_args[0][
|
||||||
|
1] == input_data['filters']['SPECIFIC_VARIABLES_NULL_VALUES']
|
||||||
|
|
||||||
|
empty_data_input_args = empty_data_mock.call_args
|
||||||
|
assert empty_data_input_args[0][0].equals(DataFrame(
|
||||||
|
{'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
|
||||||
|
assert empty_data_input_args[0][1] == input_data['filters']['EMPTY_DATA']
|
||||||
|
|
||||||
|
|
||||||
|
@mark.asyncio
|
||||||
|
@patch('laborious.activities.gates.filter_functions')
|
||||||
|
async def test_input_gate_no_filtered(
|
||||||
|
filter_functions_mock,
|
||||||
|
gates
|
||||||
|
):
|
||||||
|
specific_variables_null_values_mock = MagicMock(return_value=False)
|
||||||
|
empty_data_mock = MagicMock(return_value=False)
|
||||||
|
|
||||||
|
def functions_side_effect(x):
|
||||||
|
if x == 'SPECIFIC_VARIABLES_NULL_VALUES':
|
||||||
|
return specific_variables_null_values_mock
|
||||||
|
return empty_data_mock
|
||||||
|
|
||||||
|
filter_functions_mock.__getitem__.side_effect = functions_side_effect
|
||||||
|
|
||||||
|
input_data = {
|
||||||
|
'filters': {
|
||||||
|
'SPECIFIC_VARIABLES_NULL_VALUES': {
|
||||||
|
'POLICY': 'stop',
|
||||||
|
'VARIABLES': ['variable2']
|
||||||
|
},
|
||||||
|
'EMPTY_DATA': {
|
||||||
|
'POLICY': 'continue',
|
||||||
|
}
|
||||||
|
},
|
||||||
|
'data': {
|
||||||
|
'variable': ['variable1', 'variable2'],
|
||||||
|
'value': [1, 2]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
result = await gates.input_gate(input_data)
|
||||||
|
assert result == (None, 0)
|
||||||
|
|
||||||
|
specific_variables_null_values_input_args = specific_variables_null_values_mock.call_args
|
||||||
|
assert specific_variables_null_values_input_args[0][0].equals(DataFrame(
|
||||||
|
{'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
|
||||||
|
|
||||||
|
empty_data_input_args = empty_data_mock.call_args
|
||||||
|
assert empty_data_input_args[0][0].equals(DataFrame(
|
||||||
|
{'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
|
||||||
|
assert empty_data_input_args[0][1] == input_data['filters']['EMPTY_DATA']
|
||||||
132
tests/laborious/activities/test_postgres.py
Normal file
132
tests/laborious/activities/test_postgres.py
Normal file
@@ -0,0 +1,132 @@
|
|||||||
|
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()
|
||||||
0
tests/laborious/utils/__init__.py
Normal file
0
tests/laborious/utils/__init__.py
Normal file
0
tests/laborious/utils/filters/__init__.py
Normal file
0
tests/laborious/utils/filters/__init__.py
Normal file
26
tests/laborious/utils/filters/test_conditional_filters.py
Normal file
26
tests/laborious/utils/filters/test_conditional_filters.py
Normal file
@@ -0,0 +1,26 @@
|
|||||||
|
from pandas import DataFrame
|
||||||
|
|
||||||
|
from laborious.utils.filters.conditional_filters import filter_specific_variables_null_values, filter_empty_data
|
||||||
|
|
||||||
|
|
||||||
|
def test_filter_specific_variables_null_values():
|
||||||
|
assert filter_specific_variables_null_values(
|
||||||
|
DataFrame(
|
||||||
|
{'variable': ['variable1', 'variable2'], 'value': [1, 2]}),
|
||||||
|
variables=['variable2']) == True
|
||||||
|
|
||||||
|
|
||||||
|
def test_filter_specific_variables_null_values_with_null_values():
|
||||||
|
assert filter_specific_variables_null_values(
|
||||||
|
DataFrame(
|
||||||
|
{'variable': ['variable1', 'variable2'], 'value': [1, None]}),
|
||||||
|
variables=['variable2']) == False
|
||||||
|
|
||||||
|
|
||||||
|
def test_filter_empty_data():
|
||||||
|
assert filter_empty_data(DataFrame()) == True
|
||||||
|
|
||||||
|
|
||||||
|
def test_filter_empty_data_with_data():
|
||||||
|
assert filter_empty_data(
|
||||||
|
DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, 2]})) == False
|
||||||
0
values.yaml
Normal file
0
values.yaml
Normal file
Reference in New Issue
Block a user