SIENTIAPDE-1094

Refactor activity imports and remove unused base and logger files; update requirements for library versioning

laborious/activities/activities.py

@@ -1,13 +1,13 @@
 from temporalio import activity, workflow
 
 with workflow.unsafe.imports_passed_through():
-    from laborious.activities.postgres import Postgres
+    from sientia_do.temporal.activities.postgres import Postgres
+    from sientia_do.notifications.handlers import NotificationHandler
     from laborious.activities.mlflow import MLFlow
     from laborious.activities.gates import Gates
     from laborious.activities.opc import OPC
     from typing import Any
     from logging import Logger
-    from sientia_do.notifications.handlers import NotificationHandler
 
 
 class Activities(Postgres, MLFlow, Gates, OPC):

laborious/activities/base.py

@@ -1,26 +0,0 @@
-from typing import Any
-from logging import Logger
-from temporalio import activity
-from sientia_do.notifications.handlers import NotificationHandler
-
-
-class BaseActivity:
-    def __init__(self, logger: Logger, notification_handler: NotificationHandler):
-        self.logger = logger
-        self.notification_handler = notification_handler
-
-    @activity.defn(name="prepare_activity")
-    async def prepare_activity(self, input_data: dict[str, Any]):
-        """
-        Prepare the activity for the notification handler.
-
-        Args:
-            workflow_name (str): The name of the workflow.
-            schedule_name (str): The name of the schedule.
-            model_name (str): The name of the model.
-            model_id (str): The id of the model.
-        """
-        self.notification_handler.base_notification.pipeline_name = input_data['workflow_name']
-        self.notification_handler.base_notification.schedule_name = input_data['schedule_name']
-        self.notification_handler.base_notification.model_name = input_data['model_name']
-        self.notification_handler.base_notification.model_id = input_data['model_id']

laborious/activities/gates.py

@@ -6,7 +6,7 @@ with workflow.unsafe.imports_passed_through():
     from logging import Logger
     from sientia_do.notifications.handlers import NotificationHandler
     from sientia_do.notifications.models import NotificationLevel
-    from laborious.activities.base import BaseActivity
+    from sientia_do.temporal.activities.base import BaseActivity
     from laborious.utils.filters.mlflow_filters import nan_values_filter, api_error_filter
     from typing import Any
     from laborious.utils.filters.conditional_filters import (

laborious/activities/mlflow.py

@@ -4,11 +4,11 @@ from temporalio import activity, workflow
 
 
 with workflow.unsafe.imports_passed_through():
-    from laborious.activities.base import BaseActivity
+    from sientia_do.temporal.activities.base import BaseActivity
+    from sientia_do.notifications.handlers import NotificationHandler
     from laborious.utils.repository.model_repository import MLFlowRepository
     from typing import Any
     from logging import Logger
-    from sientia_do.notifications.handlers import NotificationHandler
 
 
 class MLFlow(BaseActivity):

laborious/activities/opc.py

@@ -5,7 +5,7 @@ with workflow.unsafe.imports_passed_through():
     from logging import Logger
     from sientia_do.notifications.handlers import NotificationHandler
     from sientia_do.notifications.models import NotificationLevel
-    from laborious.activities.base import BaseActivity
+    from sientia_do.temporal.activities.base import BaseActivity
     from laborious.utils.repository.opc_repository import OpcRepository
     from typing import Any
     import traceback

laborious/activities/postgres.py

@@ -1,181 +0,0 @@
-import traceback
-from temporalio import workflow, activity
-
-from laborious.activities.base import BaseActivity
-with workflow.unsafe.imports_passed_through():
-    from sqlalchemy import create_engine
-    from sqlalchemy.orm import sessionmaker
-    from sqlalchemy.pool import QueuePool
-    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
-
-        # Create SQLAlchemy engine with connection pooling
-        self.engine = create_engine(
-            f'postgresql://{user}:{password}@{host}:{port}/{dbname}',
-            poolclass=QueuePool,
-            pool_size=min_connections,
-            max_overflow=max_connections - min_connections,
-            pool_pre_ping=True
-        )
-        self.session_factory = sessionmaker(bind=self.engine)
-
-        BaseActivity.__init__(self, logger, notification_handler)
-
-    def close(self):
-        self.engine.dispose()
-
-    def __del__(self):
-        self.close()
-
-    @activity.defn(name="load_custom_query")
-    async def load_custom_query(self, query: str) -> dict[str, Any]:
-        """
-        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}")
-
-        data = None
-        with self.session_factory() as session:
-            try:
-                data = read_sql_query(query, self.engine)
-
-            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:
-                session.close()
-
-        if data is None:
-            return {}
-
-        # Converts any datetime datatype columns to string
-        for col in data.select_dtypes(include=['datetime64']).columns:
-            data[col] = data[col].dt.strftime('%Y-%m-%d %H:%M:%S')
-
-        self.logger.info(f"Fetched {len(data)} rows")
-        self.logger.debug(f"Data: \n{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. Contains:
-                schema (str): The schema of the table.
-                table_name (str): The name of the table.
-                model (int): The model to repeat the prediction for.
-
-        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}")
-
-        with self.session_factory() as session:
-            try:
-                session.execute(repeat_query)
-                session.commit()
-
-            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:
-                session.close()
-
-    @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. Contains:
-                schema (str): The schema of the table.
-                table_name (str): The name of the table.
-                data (DataFrame): The data to export.
-        """
-
-        self.logger.debug(
-            f"Exporting data to postgres: {input_data['data']}")
-
-        schema = input_data["schema"]
-        table_name = input_data["table_name"]
-        data = DataFrame(input_data["data"])
-
-        with self.session_factory() as session:
-            try:
-                data.to_sql(table_name, self.engine, schema=schema,
-                            if_exists="append", index=False)
-                session.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)
-
-            else:
-                self.logger.debug("Data exported to postgres")
-            finally:
-                session.close()

laborious/utils/logger.py

@@ -1,22 +0,0 @@
-from os import getenv
-import logging
-import sys
-
-
-def get_logger(name: str):
-    log_level = getenv('LOG_LEVEL', 'INFO').upper()
-
-    logger = logging.getLogger(name)
-    logger.setLevel(log_level)
-    stream_handler = logging.StreamHandler(sys.stdout)
-    stream_handler.setLevel(log_level)
-
-    stream_handler.setFormatter(
-        logging.Formatter(
-            '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
-        )
-    )
-
-    logger.addHandler(stream_handler)
-
-    return logger

laborious/utils/policies.py

@@ -1,9 +0,0 @@
-from datetime import timedelta
-from temporalio.common import RetryPolicy
-
-retry_policy = RetryPolicy(
-    initial_interval=timedelta(seconds=1),
-    backoff_coefficient=2.0,
-    maximum_interval=timedelta(minutes=1),
-    maximum_attempts=1
-)

laborious/utils/repository/model_repository.py

@@ -1,16 +1,17 @@
 """
 Model Monitoring Repository
 
-This module contains the ModelMonitoringRepository class, which is responsible for handling the communication with the Model Monitoring API.
+This module contains the ModelMonitoringRepository class, which is responsible
+for handling the communication with the Model Monitoring API.
 
-It includes the methods that are used to answer ModelMonitoringService requests using the Model Monitoring API functions.
+It includes the methods that are used to answer ModelMonitoringService requests using
+the Model Monitoring API functions.
 
 By Monitoring we mean the evaluation of the performance of models, the generation of reports.
 
 """
 from datetime import datetime
 import traceback
-import mlflow
 import pandas as pd
 from sientia.ModelServing import ModelServing
 
@@ -21,240 +22,6 @@ class MLFlowRepository():
         self.model_serving = ModelServing(tracking_uri=host,
                                           username=username, password=password)
 
-    def get_current_data_df(self, current_data: pd.DataFrame, model_name: str, target: str):
-        """
-        Get the current data as a DataFrame and update the prediction and target columns
-
-        Parameters:
-        current_data (pd.DataFrame): the current data
-        model_name (str): the name of the model
-        target (str): the target column
-
-        Returns:
-        DataFrame: the current data as a DataFrame
-
-
-        """
-        predictions = current_data['prediction']
-
-        target = current_data[target]
-        current_data = self.model_serving.get_transformed_data(
-            model_name, current_data, by='model')
-        current_data['prediction'] = predictions
-        current_data['target'] = target
-
-        return pd.DataFrame(current_data).dropna()
-
-    def get_artifact(self, destination: str, search_by: str, run_id: str = None,
-                     model_name: str = None, artifact_name: str = None) -> None:
-        """
-        Get an artifact in MLflow by experiment or model and save it to a destination path using API.
-        If the artifact is searched by model, the latest production version will be used.
-
-        Args:
-            destination: The destination path to save the artifact.
-            search_by: The way to search for the artifact ('experiment' or 'model').
-            run_id: The run ID of the experiment (if search_by is "experiment").
-            model_name: The name of the model (if search_by is "model").
-            artifact_name: The path of the artifact to download.
-
-        Returns:
-            artifact: The artifact(.csv) downloaded from MLflow.
-        """
-
-        self.model_serving.get_artifact(destination=destination, search_by=search_by,
-                                        run_id=run_id, model_name=model_name, artifact_name=artifact_name)
-
-    def calculate_model_metrics(self, real_data, predictions, flag):
-        """
-        Function to calculate the metrics of a model using API
-
-        Parameters:
-        real_data (array): the real data
-        predictions (array): the predictions
-
-        Returns:
-        dict: the metrics of the model including MSE and R2
-        """
-        return self.model_serving.get_model_metrics(reference_data=None, real_data=real_data, predictions=predictions, type_flag=flag)
-
-    def get_experiment_by_run_id(self, run_id: str) -> dict:
-        # Get the run information using the run_id
-        run = mlflow.get_run(run_id)
-
-        # Extract the experiment ID from the run
-        experiment_id = run.info.experiment_id
-
-        # Get the experiment details using the experiment ID
-        experiment = mlflow.get_experiment(experiment_id)
-        experiment_name = experiment.name
-        return experiment_name
-
-    def get_next_run_name(self, model_name: str) -> str:
-        """ 
-        Function to get the next run number of a specific model
-
-        Parameters:
-        model_name (str): the name of the model
-
-        Returns:
-        str: the next run number
-        """
-
-        runs = mlflow.search_runs(
-            experiment_names=[model_name], order_by=["start_time desc"])
-        next_run_number = len(runs) + 1
-        return f"{model_name}-{next_run_number}"
-
-    def retrain_model(self, data: pd.DataFrame, model_name: str) -> tuple:
-        """
-        Retrain a model with new data.
-
-        Parameters:
-            data (pandas.DataFrame): The new data to use for retraining.
-            model_name (str): The name of the model to retrain.
-            metrics_list (list): The metrics to be used to compare the models.
-            compare_metrics (bool): If True, the retrain will only be considered if the new model is better than the current one.
-                                If False, the retrain will always be considered.
-            split_dataset (bool): If True, the data will be split into X and Y and into training and testing sets.
-                If False, the data will be used as a unique block for retraining.
-            update_report (bool): If True, a report will be created with the data of the retrained model.
-            update_transformation (bool): If True, the model will be updated in the MLflow tracking server.
-            update_prediction (bool): If True, the prediction model will be updated in the MLflow tracking server.
-            shuffle_data (bool): If True, the data will be shuffled before splitting.
-            model_type (str): The type of model to get metrics for. Ex: 'regression', 'classification'.
-
-
-        Returns:
-            mlflow.sklearn.Model: The retrained prediction model.
-            mlflow.sklearn.Model: The retrained data model.
-            mse (float): The mean squared error of the retrained model.
-            r2 (float): The R-squared score of the retrained model.
-        """
-
-        # load predictor model
-        predictor_uri = f"models:/{model_name}/production"
-        # load transform model
-        latest_production_id = self.model_serving.get_model_run_id(
-            model_name, stage="Production"
-        )
-        transform_uri = self.model_serving.get_model_uri(
-            latest_production_id, prediction=False
-        )
-        # load
-        data_model = mlflow.sklearn.load_model(transform_uri)
-        prediction_model = mlflow.sklearn.load_model(predictor_uri)
-        data_model = data_model.fit(data)
-        treated_data = data_model.predict(data)
-        # align target column with treated_data
-        target_name = data_model.target_variable
-        y = data[target_name]
-        treated_data = pd.merge(
-            treated_data, y, left_index=True, right_index=True)
-        prediction_model = prediction_model.fit(treated_data)
-        # Example usage
-        experiment = self.get_experiment_by_run_id(latest_production_id)
-        pred_model_atributes = vars(prediction_model)  # load class attributes
-        data_model_atributes = vars(data_model)  # load class attributes
-        mlflow.set_experiment(experiment)
-        experiment_description = "Retrain model {model_name} with new data"
-        current_run_name = self.get_next_run_name(experiment)
-        with mlflow.start_run(
-            run_name=current_run_name, description=experiment_description
-        ) as _run:
-            # update transfomation model
-            # fixed parameters
-            for name_atribute, val_atribute in pred_model_atributes.items():
-                if name_atribute != "model":
-                    mlflow.log_param(name_atribute, val_atribute)
-            # update prediction model
-            for name_atribute, val_atribute in data_model_atributes.items():
-                if name_atribute != "model":
-                    mlflow.log_param(name_atribute, val_atribute)
-            # dynamic parameters, including model itself
-            mlflow.sklearn.log_model(data_model, "data_model")
-            file_path = f"laborious/data/raw_data_{model_name}.csv"
-            data.to_csv(
-                f"laborious/data/raw_data_{model_name}.csv", index=True)
-            # log the data raw
-            mlflow.log_artifact(file_path)
-
-            # dynamic parameters, including model itself
-            mlflow.sklearn.log_model(prediction_model, "prediction_model")
-            mlflow.log_param("retrain", True)
-
-            return "Model retrained successfully", experiment
-
-    def get_experiment(self, experiment_name: str) -> int:
-        experiment = mlflow.get_experiment_by_name(experiment_name)
-
-        if experiment is None:
-            raise ValueError(f'Experiment {experiment_name} not found')
-
-        return int(experiment.experiment_id)
-
-    def get_experiment_last_run(self, experiment_id: int) -> str:
-        runs = mlflow.search_runs(
-            experiment_ids=[experiment_id],
-            filter_string="",  # Sem filtro no MLflow ainda
-            output_format="pandas"
-        )
-
-        # Filtrar apenas as runs onde params.retrain == True
-        filtered_runs = runs[runs["params.retrain"] == 'True']
-
-        # Converter a coluna 'end_time' para datetime
-        filtered_runs['end_time'] = pd.to_datetime(filtered_runs['end_time'])
-
-        # Ordenar o DataFrame de forma descendente pela coluna 'end_time'
-        filtered_runs = filtered_runs.sort_values(
-            by='end_time', ascending=False)
-
-        # Pegar a última run_id do DataFrame filtrado e ordenado
-        latest_run_id = filtered_runs.iloc[0]['run_id']
-
-        return latest_run_id
-
-    def update_production_model_by_run_id(self, run_id: str, model_name: str) -> dict:
-        # Registrar o modelo
-        # Aqui estamos assumindo que você já tem um modelo salvo, caso contrário você precisará treiná-lo e salvá-lo primeiro.
-        # Se o modelo já está registrado, você pode usar o método register_model() ou pyfunc.load_model() para isso.
-        mlflow.register_model(
-            f"runs:/{run_id}/prediction_model", model_name)
-
-        # Colocar a versão do modelo em produção
-        # Depois de registrar o modelo, precisamos pegar a versão mais recente do modelo e movê-lo para o estágio 'Production'
-        client = mlflow.tracking.MlflowClient()
-
-        # Obter a versão mais recente registrada do modelo
-        model_versions = client.get_registered_model(
-            model_name).latest_versions
-        max_version = max(model_versions, key=lambda x: int(x.version)).version
-
-        # Mover a versão mais recente do modelo para o estágio de 'Production'
-        client.transition_model_version_stage(
-            name=model_name,
-            version=max_version,
-            stage="Production",
-            archive_existing_versions=True
-        )
-
-        return {
-            'model_name': model_name,
-            'version': max_version,
-            'mlflow_run_id': run_id
-        }
-
-    def update_production_model(self, experiment: str, model_name: str) -> dict:
-
-        experiment_id = self.get_experiment(experiment)
-        run_id = self.get_experiment_last_run(experiment_id)
-        metadata = self.update_production_model_by_run_id(run_id, model_name)
-
-        metadata['mlflow_experiment_id'] = experiment_id
-
-        return metadata
-
     def transform(self, model_name: str, data: pd.DataFrame, model_retention: int):
         try:
             return {

laborious/worker/worker.py

@@ -1,22 +1,23 @@
 from temporalio import workflow, client
 from temporalio.worker import Worker
-import sys
+
 
 with workflow.unsafe.imports_passed_through():
     import os
+    import sys
     import asyncio
     from laborious.workflows.predictions_batch import PredictionsBatch
     from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
     from laborious.workflows.sub_workflows.format_and_export_prediction import \
         FormatAndExportPrediction
     from laborious.activities.activities import Activities
-    from laborious.utils.logger import get_logger
     from laborious.utils.connectors_config import (
         build_postgres_config,
         build_mlflow_config,
         build_opc_config
     )
     from sientia_do.notifications.handlers import NotificationHandler
+    from sientia_do.temporal.utils.logger import get_logger
 
 
 async def main():

laborious/workflows/predictions_batch.py

@@ -3,7 +3,7 @@ from temporalio import workflow
 with workflow.unsafe.imports_passed_through():
     from laborious.activities.activities import Activities
     from typing import Any
-    from laborious.utils.policies import retry_policy
+    from sientia_do.temporal.utils.policies import retry_policy
     from datetime import timedelta
 
 

laborious/workflows/sub_workflows/format_and_export_prediction.py

@@ -4,7 +4,7 @@ with workflow.unsafe.imports_passed_through():
     from laborious.activities.activities import Activities
     from typing import Any
     from datetime import timedelta
-    from laborious.utils.policies import retry_policy
+    from sientia_do.temporal.utils.policies import retry_policy
 
 
 @workflow.defn(name="format_and_export_prediction")

laborious/workflows/sub_workflows/prediction_process.py

@@ -3,7 +3,7 @@ from temporalio import workflow
 with workflow.unsafe.imports_passed_through():
     from laborious.activities.activities import Activities
     from typing import Any
-    from laborious.utils.policies import retry_policy
+    from sientia_do.temporal.utils.policies import retry_policy
     from datetime import timedelta
 
 

requirements.txt

@@ -3,5 +3,5 @@ psycopg2-binary
 sqlalchemy
 asyncua
 redis
-git+ssh://git@github.com/Aignosi/sientia-dataops-library.git
+git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.1.14
 git+ssh://git@github.com/Aignosi/sientia-mlops-library.git

tests/laborious/activities/test_activities.py

@@ -1,7 +1,7 @@
 from pytest import mark
 from unittest.mock import patch, MagicMock, ANY
+from sientia_do.temporal.activities.postgres import Postgres
 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
@@ -139,9 +139,9 @@ async def test_prepare_activity(_mock_opc_init,
 
     await activities.prepare_activity(input_data)
 
-    assert activities.notification_handler.base_notification.pipeline_name == input_data[
+    assert activities.notification_handler.base_notification.pipeline == input_data[
         'workflow_name']
-    assert activities.notification_handler.base_notification.schedule_name == input_data[
+    assert activities.notification_handler.base_notification.trigger == input_data[
         'schedule_name']
     assert activities.notification_handler.base_notification.model_name == input_data[
         'model_name']

tests/laborious/activities/test_base.py

@@ -1,35 +0,0 @@
-from unittest.mock import MagicMock
-from laborious.activities.base import BaseActivity
-from pytest import fixture, mark
-from sientia_do.notifications.models import Notification
-
-
-@fixture
-def base_activity():
-    return BaseActivity(
-        logger=MagicMock(),
-        notification_handler=MagicMock(),
-    )
-
-
-@mark.asyncio
-async def test_prepare_activity(base_activity):
-    base_activity.notification_handler.base_notification = Notification(
-        project="project",
-        pipeline="pipeline",
-        trigger="-",
-        model_name="-",
-        model_id="-",
-    )
-
-    await base_activity.prepare_activity({
-        'workflow_name': 'test_workflow',
-        'schedule_name': 'test_schedule',
-        'model_name': 'test_model',
-        'model_id': 'test_model_id'
-    })
-
-    assert base_activity.notification_handler.base_notification.schedule_name == "test_schedule"
-    assert base_activity.notification_handler.base_notification.model_name == "test_model"
-    assert base_activity.notification_handler.base_notification.model_id == "test_model_id"
-    assert base_activity.notification_handler.base_notification.pipeline_name == "test_workflow"

tests/laborious/activities/test_postgres.py

@@ -1,159 +0,0 @@
-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()

tests/laborious/utils/repository/test_model_repository.py

@@ -1,14 +1,14 @@
 from unittest.mock import ANY, MagicMock, patch
 import numpy as np
-from pandas import DataFrame
 import pytest
 from laborious.utils.repository.model_repository import MLFlowRepository
 
 
 @pytest.fixture
 def mlflow_repository():
-    with patch('laborious.utils.repository.model_repository.ModelServing', autospec=True) as MockModelServing:
+    with patch('laborious.utils.repository.model_repository.ModelServing',
-        mock_instance = MockModelServing.return_value
+               autospec=True) as mock_model_serving:
+        mock_instance = mock_model_serving.return_value
         mock_instance.get_transformed_data = MagicMock()
 
         repo = MLFlowRepository(
@@ -19,190 +19,6 @@ def mlflow_repository():
         return repo
 
 
-def test_get_current_data_df(mlflow_repository):
-    current_data = {
-        'prediction': [1, 3],
-        'target': [1, 1],
-    }
-    mlflow_repository.model_serving.get_transformed_data.return_value = {
-        'var1': [1, 2],
-        'var2': [2, np.nan],
-    }
-    expected = DataFrame({
-        'var1': [1],
-        'var2': [2],
-        'prediction': [1],
-        'target': [1],
-    })
-    output = mlflow_repository.get_current_data_df(current_data,
-                                                   'model', 'target')
-
-    mlflow_repository.model_serving.get_transformed_data.assert_called_once_with(
-        'model', current_data, by='model')
-
-    diff = output.compare(expected)
-    assert diff.empty
-
-
-def test_get_artifact(mlflow_repository):
-    mlflow_repository.get_artifact(
-        'destination', 'search_by', 'run_id', 'model', 'artifact'
-    )
-    mlflow_repository.model_serving.get_artifact.assert_called_once_with(
-        destination='destination',
-        search_by='search_by',
-        run_id='run_id',
-        model_name='model',
-        artifact_name='artifact'
-    )
-
-
-def test_calculate_model_metrics(mlflow_repository):
-    mlflow_repository.model_serving.get_model_metrics.return_value = 'data'
-    real_data = 'real_data'
-    predictions = 'predictions'
-    flag = 'flag'
-    output = mlflow_repository.calculate_model_metrics(
-        real_data, predictions, flag
-    )
-    mlflow_repository.model_serving.get_model_metrics.assert_called_once_with(
-        reference_data=None,
-        real_data=real_data,
-        predictions=predictions,
-        type_flag=flag
-    )
-    assert output == 'data'
-
-
-@patch('laborious.utils.repository.model_repository.mlflow')
-def test_get_experiment_by_run_id(mlflow, mlflow_repository):
-    mlflow.get_run.return_value = MagicMock(
-        info=MagicMock(
-            experiment_id='0',
-        )
-    )
-    mlflow.get_experiment.return_value = MagicMock()
-    mlflow.get_experiment.return_value.name = 'test'
-
-    output = mlflow_repository.get_experiment_by_run_id('0')
-    assert output == 'test'
-    mlflow.get_run.assert_called_once_with('0')
-    mlflow.get_experiment.assert_called_once_with('0')
-
-
-@patch('laborious.utils.repository.model_repository.mlflow')
-def test_get_next_run_name(mlflow, mlflow_repository):
-    mlflow.search_runs.return_value = [1, 2, 3]
-    output = mlflow_repository.get_next_run_name('run')
-    assert output == 'run-4'
-    mlflow.search_runs.assert_called_once_with(
-        experiment_names=['run'],
-        order_by=['start_time desc'],
-    )
-
-
-@patch('laborious.utils.repository.model_repository.mlflow')
-def test_get_experiment_success(mlflow, mlflow_repository):
-    mlflow.get_experiment_by_name.return_value = MagicMock(
-        experiment_id='0')
-
-    output = mlflow_repository.get_experiment('test')
-
-    assert output == 0
-
-
-@patch('laborious.utils.repository.model_repository.mlflow')
-def test_get_experiment_error(mlflow, mlflow_repository):
-    mlflow.get_experiment_by_name.return_value = None
-
-    try:
-        mlflow_repository.get_experiment('test')
-    except ValueError as e:
-        assert str(e) == 'Experiment test not found'
-    else:
-        assert False
-
-
-@patch('laborious.utils.repository.model_repository.mlflow')
-def test_get_experiment_last_run(mlflow, mlflow_repository):
-    mlflow.search_runs.return_value = DataFrame({
-        'params.retrain': ['True', 'False', 'True', 'False'],
-        'end_time': ['2021-01-01', '2021-01-02', '2021-01-03', '2021-01-04'],
-        'run_id': ['0', '1', '2', '3'],
-    })
-
-    output = mlflow_repository.get_experiment_last_run(0)
-
-    mlflow.search_runs.assert_called_once_with(
-        experiment_ids=[0],
-        filter_string="",
-        output_format="pandas",
-    )
-
-    assert output == '2'
-
-
-@patch('laborious.utils.repository.model_repository.mlflow')
-def test_update_production_model_by_run_id(mlflow, mlflow_repository):
-    client_mock = MagicMock()
-    mlflow.tracking.MlflowClient.return_value = client_mock
-
-    client_mock.get_registered_model.return_value = MagicMock(
-        latest_versions=[
-            MagicMock(version='1'),
-            MagicMock(version='2'),
-            MagicMock(version='3'),
-        ]
-    )
-    output = mlflow_repository.update_production_model_by_run_id('0', 'test')
-
-    mlflow.register_model.assert_called_once_with(
-        "runs:/0/prediction_model",
-        'test',
-    )
-
-    mlflow.tracking.MlflowClient.assert_called_once()
-    client_mock.get_registered_model.assert_called_once_with('test')
-    client_mock.transition_model_version_stage.assert_called_once_with(
-        name='test',
-        version='3',
-        stage='Production',
-        archive_existing_versions=True,
-    )
-
-    assert output == {
-        'model_name': 'test',
-        'version': '3',
-        'mlflow_run_id': '0',
-    }
-
-
-def test_update_production_model(mlflow_repository):
-    connector = mlflow_repository
-
-    with patch.object(connector, 'get_experiment',
-                      return_value='0') as get_experiment:
-        with patch.object(connector, 'get_experiment_last_run',
-                          return_value='2') as get_experiment_last_run:
-            with patch.object(connector, 'update_production_model_by_run_id',
-                              return_value={'model_name': 'test', 'version': '3',
-                                            'mlflow_run_id': '0'}) as update_production_model_by_run_id:
-
-                output = connector.update_production_model('0', 'test')
-
-                get_experiment.assert_called_once_with('0')
-                get_experiment_last_run.assert_called_once_with('0')
-                update_production_model_by_run_id.assert_called_once_with(
-                    '2', 'test')
-
-                assert output == {
-                    'model_name': 'test',
-                    'version': '3',
-                    'mlflow_run_id': '0',
-                    'mlflow_experiment_id': '0',
-                }
-
-
 def test_transform_success(mlflow_repository):
     data = 'data'
     model_name = 'model'

tests/laborious/utils/test_logger.py

@@ -1,37 +0,0 @@
-import os
-from unittest.mock import patch
-import logging
-import pytest
-from laborious.utils.logger import get_logger
-
-
-@pytest.fixture
-def mock_env_vars():
-    with patch.dict(os.environ, {}, clear=True):
-        yield
-
-
-@pytest.mark.usefixtures("mock_env_vars")
-@patch('laborious.utils.logger.logging.Formatter')
-@patch('laborious.utils.logger.logging.StreamHandler')
-def test_get_logger_defaults(mock_stream_handler, mock_formatter):
-    """Test logger creation with default settings"""
-    # Mock the StreamHandler and Formatter
-
-    logger = get_logger('test_logger')
-
-    # Verify logger settings
-    assert logger.name == 'test_logger'
-    assert logger.level == logging.INFO
-
-    # Verify handler configuration
-    mock_stream_handler.return_value.setLevel.assert_called_once_with('INFO')
-    mock_stream_handler.return_value.setFormatter.assert_called_once()
-
-    # Verify formatter configuration
-    mock_formatter.assert_called_once_with(
-        '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
-    )
-
-    # Verify handler was added to logger
-    assert len(logger.handlers) == 1