From e7f214b14473a48d60b7bcac4fc91030400d0b95 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Wed, 7 May 2025 17:39:45 -0300 Subject: [PATCH] SIENTIAPDE-994 Remove unused utility files and update requirements.txt to include new dependencies for data processing and database interaction. --- laborious/util/__init__.py => README.md | 0 laborious/activities/base.py | 17 ++ laborious/activities/gates.py | 47 +++ laborious/activities/mlflow.py | 0 laborious/activities/opc.py | 0 laborious/activities/postgres.py | 162 +++++++++++ laborious/utils/__init__.py | 0 laborious/utils/filters/__init__.py | 0 laborious/utils/filters/api_filters.py | 69 +++++ laborious/utils/filters/base_filter.py | 13 + .../utils/filters/conditional_filters.py | 19 ++ laborious/{util => utils}/git_clone.py | 0 requirements.txt | 6 +- tests/__init__.py | 0 tests/laborious/__init__.py | 0 tests/laborious/activities/__init__.py | 0 tests/laborious/activities/test_gates.py | 274 ++++++++++++++++++ tests/laborious/activities/test_postgres.py | 132 +++++++++ tests/laborious/utils/__init__.py | 0 tests/laborious/utils/filters/__init__.py | 0 .../utils/filters/test_conditional_filters.py | 26 ++ values.yaml | 0 22 files changed, 762 insertions(+), 3 deletions(-) rename laborious/util/__init__.py => README.md (100%) create mode 100644 laborious/activities/base.py create mode 100644 laborious/activities/gates.py create mode 100644 laborious/activities/mlflow.py create mode 100644 laborious/activities/opc.py create mode 100644 laborious/activities/postgres.py create mode 100644 laborious/utils/__init__.py create mode 100644 laborious/utils/filters/__init__.py create mode 100644 laborious/utils/filters/api_filters.py create mode 100644 laborious/utils/filters/base_filter.py create mode 100644 laborious/utils/filters/conditional_filters.py rename laborious/{util => utils}/git_clone.py (100%) create mode 100644 tests/__init__.py create mode 100644 tests/laborious/__init__.py create mode 100644 tests/laborious/activities/__init__.py create mode 100644 tests/laborious/activities/test_gates.py create mode 100644 tests/laborious/activities/test_postgres.py create mode 100644 tests/laborious/utils/__init__.py create mode 100644 tests/laborious/utils/filters/__init__.py create mode 100644 tests/laborious/utils/filters/test_conditional_filters.py create mode 100644 values.yaml diff --git a/laborious/util/__init__.py b/README.md similarity index 100% rename from laborious/util/__init__.py rename to README.md diff --git a/laborious/activities/base.py b/laborious/activities/base.py new file mode 100644 index 0000000..c195f8a --- /dev/null +++ b/laborious/activities/base.py @@ -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 diff --git a/laborious/activities/gates.py b/laborious/activities/gates.py new file mode 100644 index 0000000..73cd49a --- /dev/null +++ b/laborious/activities/gates.py @@ -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 diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py new file mode 100644 index 0000000..e69de29 diff --git a/laborious/activities/opc.py b/laborious/activities/opc.py new file mode 100644 index 0000000..e69de29 diff --git a/laborious/activities/postgres.py b/laborious/activities/postgres.py new file mode 100644 index 0000000..0fc8cdd --- /dev/null +++ b/laborious/activities/postgres.py @@ -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) diff --git a/laborious/utils/__init__.py b/laborious/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/laborious/utils/filters/__init__.py b/laborious/utils/filters/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/laborious/utils/filters/api_filters.py b/laborious/utils/filters/api_filters.py new file mode 100644 index 0000000..4e4bdd2 --- /dev/null +++ b/laborious/utils/filters/api_filters.py @@ -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 diff --git a/laborious/utils/filters/base_filter.py b/laborious/utils/filters/base_filter.py new file mode 100644 index 0000000..f5d9948 --- /dev/null +++ b/laborious/utils/filters/base_filter.py @@ -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}') \ No newline at end of file diff --git a/laborious/utils/filters/conditional_filters.py b/laborious/utils/filters/conditional_filters.py new file mode 100644 index 0000000..d819a3d --- /dev/null +++ b/laborious/utils/filters/conditional_filters.py @@ -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 diff --git a/laborious/util/git_clone.py b/laborious/utils/git_clone.py similarity index 100% rename from laborious/util/git_clone.py rename to laborious/utils/git_clone.py diff --git a/requirements.txt b/requirements.txt index 72a1b86..0b537a6 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ -GitPython temporalio -pytest -python-dotenv \ No newline at end of file +psycopg2-binary +pandas +git+ssh://git@github.com/Aignosi/sientia-dataops-library.git \ No newline at end of file diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/laborious/__init__.py b/tests/laborious/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/laborious/activities/__init__.py b/tests/laborious/activities/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/laborious/activities/test_gates.py b/tests/laborious/activities/test_gates.py new file mode 100644 index 0000000..28f83ca --- /dev/null +++ b/tests/laborious/activities/test_gates.py @@ -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'] diff --git a/tests/laborious/activities/test_postgres.py b/tests/laborious/activities/test_postgres.py new file mode 100644 index 0000000..98e74e5 --- /dev/null +++ b/tests/laborious/activities/test_postgres.py @@ -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() diff --git a/tests/laborious/utils/__init__.py b/tests/laborious/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/laborious/utils/filters/__init__.py b/tests/laborious/utils/filters/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/laborious/utils/filters/test_conditional_filters.py b/tests/laborious/utils/filters/test_conditional_filters.py new file mode 100644 index 0000000..5e9b7a0 --- /dev/null +++ b/tests/laborious/utils/filters/test_conditional_filters.py @@ -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 diff --git a/values.yaml b/values.yaml new file mode 100644 index 0000000..e69de29