from unittest.mock import ANY, MagicMock, call, patch import numpy as np from pytest import fixture, mark, raises from sientia_do.notifications.models import NotificationLevel from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ from laborious.activities.mlflow import MLFlow @patch('laborious.activities.mlflow.MLFlowRepository') @patch('laborious.activities.mlflow.MinioRepository') def test___init__(mock_minio_repository, mock_mlflow_repository): mlflow = MLFlow( mlflow_host='http://localhost', mlflow_port=5000, mlflow_username='admin', mlflow_password='admin', minio_config={ 'endpoint_url': 'http://localhost:9000', 'access_key': 'minio', 'secret_key': 'minio123', 'region_name': 'us-east-1', 'default_bucket': 'test', }, logger=MagicMock(), notification_handler=MagicMock(), ) assert mlflow.mlflow_host == 'http://localhost' assert mlflow.mlflow_port == 5000 assert mlflow.mlflow_username == 'admin' assert mlflow.mlflow_password == 'admin' mock_mlflow_repository.assert_called_once_with('http://localhost:5000', 'admin', 'admin', ANY) mock_minio_repository.assert_called_once_with( logger=ANY, notification_handler=ANY, minio_endpoint_url='http://localhost:9000', minio_access_key='minio', minio_secret_key='minio123', minio_region_name='us-east-1', minio_default_bucket='test', ) @fixture @patch('laborious.activities.mlflow.MLFlowRepository') @patch('laborious.activities.mlflow.MinioRepository') def mlflow(mock_minio_repository, mock_mlflow_repository): mlflow = MLFlow( mlflow_host='http://localhost:5000', mlflow_port=5000, mlflow_username='admin', mlflow_password='admin', minio_config={ 'endpoint_url': 'http://localhost:9000', 'access_key': 'minio', 'secret_key': 'minio123', 'region_name': 'us-east-1', 'default_bucket': 'test', }, logger=MagicMock(), notification_handler=MagicMock(), ) mlflow.send_notification = MagicMock() return mlflow metadata = { 'metadata': { 'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule', }, } @mark.asyncio @patch('laborious.activities.mlflow.DataFrame') @patch('laborious.activities.mlflow.max') async def test_request_transform_success(mock_max, mock_dataframe, mlflow): mock_max.return_value = '2024-01-02' # Mock input data input_data = { **metadata, 'data': [ { 'timestamp': '2024-01-01', 'variable': 'var1', 'value': 1.0, 'created_at': '2024-01-01 12:00:00', }, { 'timestamp': '2024-01-01', 'variable': 'var2', 'value': 2.0, 'created_at': '2024-01-01 12:00:00', }, { 'timestamp': '2024-01-02', 'variable': 'var1', 'value': 3.0, 'created_at': '2024-01-02 12:00:00', }, { 'timestamp': '2024-01-02', 'variable': 'var2', 'value': 4.0, 'created_at': '2024-01-02 12:00:00', }, { 'timestamp': '2024-01-02', 'variable': 'var1', 'value': 1.0, 'created_at': '2024-01-01 12:00:00', }, { 'timestamp': '2024-01-02', 'variable': 'var2', 'value': 1.0, 'created_at': '2024-01-01 12:00:00', }, ], 'model_name': 'test_model', 'model_config': {}, } # Mock the transform response expected_response = {'prediction': [0.5, 0.6], 'timestamp': ['2024-01-01', '2024-01-02']} mlflow.model_monitoring_repository.transform.return_value = expected_response mock_dataframe.return_value.sort_values.return_value = mock_dataframe.return_value mock_dataframe.return_value.drop_duplicates.return_value = mock_dataframe.return_value # Call the method response_data = await mlflow.request_transform(input_data) # Verify the data was correctly transformed mock_dataframe.assert_called_once_with(input_data['data']) mock_dataframe.return_value.pivot.assert_called_once_with( index='timestamp', columns='variable', values='value' ) mock_dataframe = mock_dataframe.return_value.pivot.return_value mock_dataframe.fillna.assert_called_once_with(np.nan, inplace=True) # mock_dataframe.reset_index.assert_called_once() mock_dataframe.columns.name = None # Verify the response assert response_data == expected_response # Verify the repository was called with correct arguments mlflow.model_monitoring_repository.transform.assert_called_once_with( 'test_model', mock_dataframe, {}, metadata['metadata'] ) @mark.asyncio @patch('laborious.activities.mlflow.DataFrame') @patch('laborious.activities.mlflow.to_datetime') @patch('laborious.activities.mlflow.max') async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflow): mock_max.return_value = '2024-01-02' # Mock input data input_data = { **metadata, 'data': { 'variable': { '2024-01-01': 'var1', '2024-01-02': 'var2', '2024-01-03': 'var1', '2024-01-04': 'var2', }, 'value': {'2024-01-01': 1.0, '2024-01-02': 2.0, '2024-01-03': 3.0, '2024-01-04': 4.0}, }, 'model_name': 'test_model', 'model_config': {}, } # Mock the predict response expected_response = {'prediction': [0.5, 0.6]} mlflow.model_monitoring_repository.predict.return_value = expected_response # Call the method response_data = await mlflow.request_predict(input_data) mock_dataframe.assert_called_once_with(input_data['data']) mock_dataframe.return_value.replace.assert_called_once_with(np.nan, None, inplace=True) mock_dataframe.return_value.__setitem__.assert_any_call( 'timestamp', mock_to_datetime.return_value.dt.strftime.return_value ) mock_dataframe.return_value.__setitem__.assert_any_call( 'timestamp', mock_to_datetime.return_value.dt.strftime.return_value ) mock_to_datetime.assert_called_once_with( mock_dataframe.return_value.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ ) mock_to_datetime.return_value.dt.strftime.assert_called_once_with( DATETIME_FORMAT ) mock_to_datetime.assert_called_once_with( mock_dataframe.return_value.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ ) mock_to_datetime.return_value.dt.strftime.assert_called_once_with(DATETIME_FORMAT) # Verify the response assert response_data == expected_response # Verify the repository was called with correct arguments mlflow.model_monitoring_repository.predict.assert_called_once_with( 'test_model', mock_dataframe.return_value, {}, metadata['metadata'] ) @mark.asyncio @patch('laborious.activities.mlflow.to_datetime') async def test_retrain_model_success_data_success_retrain(mock_to_datetime, mlflow): mlflow.model_monitoring_repository.retrain_model.return_value = { 'success': True, 'experiment': 'test_experiment', 'message': 'Model retrained successfully.', } response = await mlflow.retrain_model( { **metadata, 'object_key': 'test_object_key', 'model_name': 'test_model', 'model_config': { 'target': 'target', 'transform_flavor': 'sklearn', 'predict_flavor': 'pyfunc', }, } ) raw_data = mlflow.minio_repository.get_parquet_as_dataframe.return_value timestamp = raw_data.__getitem__.return_value.max.return_value raw_data.sort_values.assert_called_once_with('created_at', ascending=False) raw_data.sort_values.return_value.drop_duplicates.assert_called_once_with( subset=['variable', 'timestamp'], keep='first' ) raw_data = raw_data.sort_values.return_value.drop_duplicates.return_value raw_data.drop.assert_has_calls( [ call(columns=['model_id'], inplace=True, errors='ignore'), call(columns=['created_at'], inplace=True, errors='ignore'), ] ) raw_data.pivot.assert_called_once_with(index='timestamp', columns='variable', values='value') raw_data.pivot.return_value.fillna.assert_called_once_with(np.nan, inplace=True) raw_data = raw_data.pivot.return_value raw_data.__setitem__.assert_has_calls( [ call('timestamp', raw_data.index), call('timestamp', mock_to_datetime.return_value.dt.strftime.return_value), call('timestamp', mock_to_datetime.return_value), ] ) mock_to_datetime.assert_has_calls( [call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ)] ) mock_to_datetime.assert_has_calls( [call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT)] ) mlflow.model_monitoring_repository.retrain_model.assert_called_once_with( data=raw_data, model_name='test_model', model_config={ 'target': 'target', 'transform_flavor': 'sklearn', 'predict_flavor': 'pyfunc', }, metadata=metadata['metadata'], ) assert response == { 'success': True, 'experiment': 'test_experiment', 'message': 'Model retrained successfully.', 'timestamp': timestamp, } @mark.asyncio @patch('laborious.activities.mlflow.to_datetime') async def test_retrain_model_success_data_fail_retrain(mock_to_datetime, mlflow): mlflow.model_monitoring_repository.retrain_model.return_value = { 'success': False, 'traceback': 'test_traceback', 'message': 'Model retrained failed.', } response = await mlflow.retrain_model( { **metadata, 'object_key': 'test_object_key', 'model_name': 'test_model', 'model_config': { 'target': 'target', 'transform_flavor': 'sklearn', 'predict_flavor': 'pyfunc', }, } ) raw_data = mlflow.minio_repository.get_parquet_as_dataframe.return_value timestamp = raw_data.__getitem__.return_value.max.return_value raw_data.sort_values.assert_called_once_with('created_at', ascending=False) raw_data.sort_values.return_value.drop_duplicates.assert_called_once_with( subset=['variable', 'timestamp'], keep='first' ) raw_data = raw_data.sort_values.return_value.drop_duplicates.return_value raw_data.drop.assert_has_calls( [ call(columns=['model_id'], inplace=True, errors='ignore'), call(columns=['created_at'], inplace=True, errors='ignore'), ] ) raw_data.pivot.assert_called_once_with(index='timestamp', columns='variable', values='value') raw_data.pivot.return_value.fillna.assert_called_once_with(np.nan, inplace=True) raw_data = raw_data.pivot.return_value raw_data.__setitem__.assert_has_calls( [ call('timestamp', raw_data.index), call('timestamp', mock_to_datetime.return_value.dt.strftime.return_value), call('timestamp', mock_to_datetime.return_value), ] ) mock_to_datetime.assert_has_calls( [call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ)] ) mock_to_datetime.assert_has_calls( [call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT)] ) mlflow.model_monitoring_repository.retrain_model.assert_called_once_with( data=raw_data, model_name='test_model', model_config={ 'target': 'target', 'transform_flavor': 'sklearn', 'predict_flavor': 'pyfunc', }, metadata=metadata['metadata'], ) mlflow.send_notification.assert_called_once_with( metadata=metadata['metadata'], notification_id='RETRAIN_MODEL_ERROR', message='Error retraining model test_model: Model retrained failed.', block='retrain_model', level=NotificationLevel.ERROR, attachment_content=ANY, ) assert response == { 'success': False, 'traceback': 'test_traceback', 'message': 'Model retrained failed.', 'timestamp': timestamp, } @mark.asyncio async def test_retrain_model_data_error(mlflow): mlflow.minio_repository.get_parquet_as_dataframe.side_effect = Exception( 'Error loading retrain data' ) response = await mlflow.retrain_model( { **metadata, 'object_key': 'test_object_key', 'model_name': 'test_model', 'model_config': { 'target': 'target', 'transform_flavor': 'sklearn', 'predict_flavor': 'pyfunc', }, } ) assert response == { 'success': False, 'message': 'Error loading retrain data: Error loading retrain data', 'traceback': ANY, 'timestamp': ANY, } @mark.asyncio async def test_retrain_model_data_error_no_minio_repository(mlflow): mlflow.minio_repository = None with raises(ValueError) as e: await mlflow.retrain_model( { **metadata, 'object_key': 'test_object_key', 'model_name': 'test_model', 'model_config': { 'target': 'target', 'transform_flavor': 'sklearn', 'predict_flavor': 'pyfunc', }, } ) assert str(e.value) == 'Minio repository not initialized' @mark.asyncio async def test_update_production_model(mlflow): input_data = { **metadata, 'model_name': 'test_model', 'model_id': 1, 'experiment': 'test', 'timestamp': 2, 'status': 'success', } response = await mlflow.update_production_model(input_data) mlflow.model_monitoring_repository.update_production_model.assert_called_once_with( experiment='test', model_name='test_model', metadata=metadata['metadata'] ) assert response == mlflow.model_monitoring_repository.update_production_model.return_value @mark.asyncio async def test_update_production_model_error(mlflow): mlflow.model_monitoring_repository.update_production_model.side_effect = Exception( 'Error updating production model' ) input_data = { **metadata, 'model_name': 'test_model', 'model_id': 1, 'experiment': 'test', 'timestamp': 2, 'status': 'success', } try: await mlflow.update_production_model(input_data) except Exception as e: assert str(e) == 'Error updating production model' mlflow.send_notification.assert_called_once_with( metadata=metadata['metadata'], notification_id='UPDATE_PRODUCTION_MODEL_ERROR', message='Error updating production model test_model: Error updating production model', block='update_production_model', level=NotificationLevel.ERROR, attachment_content=ANY, ) else: raise AssertionError('No exception raised')