from unittest.mock import ANY, MagicMock, patch import numpy as np from pytest import fixture, mark from sientia_do.notifications.models import NotificationLevel from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ from model_manager.activities.mlflow import MLFlow @patch('model_manager.activities.mlflow.MLFlowRepository') def test___init__(mock_mlflow_repository): mlflow = MLFlow( mlflow_host='http://localhost', mlflow_port=5000, mlflow_username='admin', mlflow_password='admin', 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) @fixture @patch('model_manager.activities.mlflow.MLFlowRepository') def mlflow(mock_mlflow_repository): mlflow = MLFlow( mlflow_host='http://localhost:5000', mlflow_port=5000, mlflow_username='admin', mlflow_password='admin', 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('model_manager.activities.mlflow.DataFrame') @patch('model_manager.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('model_manager.activities.mlflow.DataFrame') @patch('model_manager.activities.mlflow.to_datetime') @patch('model_manager.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_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 async def test_retrain_model(mlflow): data = { 'model_id': [4, 5, 6, 7], 'created_at': [1, 2, 3, 4], 'timestamp': [1, 1, 2, 2], 'variable': ['var1', 'var2', 'var1', 'var2'], 'value': [1, 2, 3, 4], } mlflow.model_monitoring_repository.retrain_model.return_value = ( 'Model retrained successfully', 'test', ) response = await mlflow.retrain_model({**metadata, 'data': data, 'model_name': 'test_model'}) mlflow.model_monitoring_repository.retrain_model.assert_called_once() assert response == { 'status': 'Model retrained successfully', 'timestamp': 2, 'experiment': 'test', } @mark.asyncio async def test_retrain_model_error(mlflow): mlflow.model_monitoring_repository.retrain_model.side_effect = Exception( 'Error retraining model' ) data = { 'model_id': [4, 5, 6, 7], 'created_at': [1, 2, 3, 4], 'timestamp': [1, 1, 2, 2], 'variable': ['var1', 'var2', 'var1', 'var2'], 'value': [1, 2, 3, 4], } try: await mlflow.retrain_model({**metadata, 'data': data, 'model_name': 'test_model'}) except Exception as e: # noqa: BLE001 assert str(e) == 'Error retraining model' mlflow.send_notification.assert_called_once_with( metadata=metadata['metadata'], notification_id='RETRAIN_MODEL_ERROR', message='Error retraining model test_model: Error retraining model', block='retrain_model', level=NotificationLevel.ERROR, attachment_content=ANY, ) else: raise AssertionError('No exception raised') @mark.asyncio async def test_update_production_model(mlflow): mlflow.model_monitoring_repository.update_production_model.return_value = { 'data1': 1, 'data2': 2, } 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' ) assert response == { 'data1': {0: 1}, 'data2': {0: 2}, 'model_id': {0: 1}, 'model_name': {0: 'test_model'}, 'timestamp': {0: 2}, 'status': {0: 'success'}, } @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: # noqa: BLE001 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')