from unittest.mock import ANY, AsyncMock, 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 @fixture(autouse=True) def _passthrough_from_dict(): with patch( 'laborious.activities.mlflow.MinioDataFramePayload.from_dict', side_effect=lambda x: x ): yield @patch('laborious.activities.mlflow.MLFlowRepository') @patch('laborious.activities.mlflow.MinioRepository') def test___init__(mock_minio_repository, mock_mlflow_repository): logger = MagicMock() notification_handler = MagicMock() metrics_controller = AsyncMock() minio_repo = mock_minio_repository( endpoint='localhost:9000', access_key='minio', secret_key='minio123', logger=logger, notification_handler=notification_handler, metrics_controller=metrics_controller, bucket='test', ) mlflow = MLFlow( mlflow_host='http://localhost', mlflow_port=5000, mlflow_username='admin', mlflow_password='admin', minio_repository=minio_repo, logger=logger, notification_handler=notification_handler, metrics_controller=metrics_controller, ) 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, ANY, ANY ) mock_minio_repository.assert_called_once_with( endpoint='localhost:9000', access_key='minio', secret_key='minio123', logger=ANY, notification_handler=ANY, metrics_controller=ANY, bucket='test', ) @fixture @patch('laborious.activities.mlflow.MLFlowRepository') @patch('laborious.activities.mlflow.MinioRepository') def mlflow(mock_minio_repository, mock_mlflow_repository): logger = MagicMock() notification_handler = MagicMock() metrics_controller = AsyncMock() minio_repo = mock_minio_repository( endpoint='localhost:9000', access_key='minio', secret_key='minio123', logger=logger, notification_handler=notification_handler, metrics_controller=metrics_controller, bucket='test', ) mlflow = MLFlow( mlflow_host='http://localhost:5000', mlflow_port=5000, mlflow_username='admin', mlflow_password='admin', minio_repository=minio_repo, logger=logger, notification_handler=notification_handler, metrics_controller=metrics_controller, ) mlflow.model_monitoring_repository = AsyncMock() mlflow.minio_repository = AsyncMock() mlflow.send_notification = MagicMock() mlflow.emit_metric = AsyncMock() mlflow.send_notification_async = AsyncMock() mlflow.error = MagicMock() mlflow.debug = MagicMock() mlflow.info = MagicMock() mlflow.warning = MagicMock() mlflow.critical = 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.MinioDataFramePayload.from_dataframe', new_callable=AsyncMock, ) async def test_request_transform_success(mock_from_dataframe, mlflow): data_mock = MagicMock() payload = AsyncMock() payload.retrieve = AsyncMock(return_value=data_mock) input_data = { **metadata, 'data': payload, 'model_name': 'test_model', 'model_config': {}, } transform_response = {'success': True, 'content': MagicMock()} mlflow.model_monitoring_repository.transform.return_value = transform_response data_mock.sort_values.return_value = data_mock data_mock.drop_duplicates.return_value = data_mock data_mock.pivot.return_value = data_mock response_data = await mlflow.request_transform(input_data) mlflow.model_monitoring_repository.transform.assert_called_once_with( 'test_model', data_mock, {}, metadata['metadata'] ) mock_from_dataframe.assert_called_once() assert response_data == mock_from_dataframe.return_value @mark.asyncio @patch( 'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe', new_callable=AsyncMock, ) async def test_request_transform_failure(mock_from_dataframe, mlflow): data_mock = MagicMock() payload = AsyncMock() payload.retrieve = AsyncMock(return_value=data_mock) input_data = { **metadata, 'data': payload, 'model_name': 'test_model', 'model_config': {}, } transform_response = {'success': False, 'message': 'Transform failed'} mlflow.model_monitoring_repository.transform.return_value = transform_response data_mock.sort_values.return_value = data_mock data_mock.drop_duplicates.return_value = data_mock data_mock.pivot.return_value = data_mock response_data = await mlflow.request_transform(input_data) mock_from_dataframe.assert_called_once_with( dataframe=None, minio_repo=mlflow.minio_repository, model_name='test_model', operation='transform', status=transform_response, workflow_metadata=metadata['metadata'], last_timestamp=payload.last_timestamp, ) assert response_data == mock_from_dataframe.return_value @mark.asyncio @patch( 'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe', new_callable=AsyncMock, ) @patch('laborious.activities.mlflow.to_datetime') async def test_request_predict(mock_to_datetime, mock_from_dataframe, mlflow): data_mock = MagicMock() payload = AsyncMock() payload.retrieve = AsyncMock(return_value=data_mock) input_data = { **metadata, 'data': payload, 'model_name': 'test_model', 'model_config': {}, } predict_response = {'success': True, 'content': MagicMock()} mlflow.model_monitoring_repository.predict.return_value = predict_response response_data = await mlflow.request_predict(input_data) data_mock.replace.assert_called_once_with(np.nan, None, inplace=True) mock_to_datetime.assert_called_once_with( data_mock.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ ) mock_to_datetime.return_value.dt.strftime.assert_called_once_with(DATETIME_FORMAT) mlflow.model_monitoring_repository.predict.assert_called_once_with( 'test_model', data_mock, {}, metadata['metadata'] ) mock_from_dataframe.assert_called_once() assert response_data == mock_from_dataframe.return_value @mark.asyncio @patch( 'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe', new_callable=AsyncMock, ) @patch('laborious.activities.mlflow.to_datetime') async def test_request_predict_failure(mock_to_datetime, mock_from_dataframe, mlflow): data_mock = MagicMock() payload = AsyncMock() payload.retrieve = AsyncMock(return_value=data_mock) input_data = { **metadata, 'data': payload, 'model_name': 'test_model', 'model_config': {}, } predict_response = {'success': False, 'message': 'Predict failed'} mlflow.model_monitoring_repository.predict.return_value = predict_response response_data = await mlflow.request_predict(input_data) mock_from_dataframe.assert_called_once_with( dataframe=None, minio_repo=mlflow.minio_repository, model_name='test_model', operation='predict', status=predict_response, workflow_metadata=metadata['metadata'], last_timestamp=payload.last_timestamp, ) assert response_data == mock_from_dataframe.return_value @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.', } raw_data = MagicMock(columns=['variable', 'timestamp', 'value']) payload = AsyncMock() payload.retrieve = AsyncMock(return_value=raw_data) response = await mlflow.retrain_model( { **metadata, 'data': payload, 'model_name': 'test_model', 'model_config': { 'target': 'target', 'transform_flavor': 'sklearn', 'predict_flavor': 'pyfunc', }, } ) timestamp = raw_data.__getitem__.return_value.max.return_value raw_data.sort_values.assert_not_called() raw_data.drop_duplicates.assert_called_once_with(subset=['variable', 'timestamp'], keep='first') raw_data = raw_data.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_with_payload_data(mock_to_datetime, mlflow): mlflow.model_monitoring_repository.retrain_model.return_value = { 'success': True, 'experiment': 'test_experiment', 'message': 'Model retrained successfully.', } raw_data = MagicMock(columns=['variable', 'timestamp', 'value', 'created_at']) payload = AsyncMock() payload.retrieve = AsyncMock(return_value=raw_data) response = await mlflow.retrain_model( { **metadata, 'data': payload, 'model_name': 'test_model', 'model_config': { 'target': 'target', 'transform_flavor': 'sklearn', 'predict_flavor': 'pyfunc', }, } ) assert response['success'] is True mlflow.minio_repository.download_file.assert_not_called() @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.', } raw_data = MagicMock(columns=['variable', 'timestamp', 'value', 'created_at']) payload = AsyncMock() payload.retrieve = AsyncMock(return_value=raw_data) response = await mlflow.retrain_model( { **metadata, 'data': payload, 'model_name': 'test_model', 'model_config': { 'target': 'target', 'transform_flavor': 'sklearn', 'predict_flavor': 'pyfunc', }, } ) 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_async.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): response = await mlflow.retrain_model( { **metadata, 'model_name': 'test_model', 'model_config': { 'target': 'target', 'transform_flavor': 'sklearn', 'predict_flavor': 'pyfunc', }, } ) assert response == { 'success': False, 'message': "Error loading retrain data: '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_async.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') @mark.asyncio @patch('laborious.activities.mlflow.to_datetime') async def test_get_reference_data_success(mock_to_datetime, mlflow): # Arrange input_data = { **metadata, 'model_name': 'test_model', } # Mock reference data DataFrame mock_reference_data = MagicMock() mock_reference_data.__getitem__.return_value = MagicMock() mock_to_datetime.return_value.dt.strftime.return_value = MagicMock() mock_reference_data.to_dict.return_value = [ {'timestamp': '2023-05-26 11:12:27', 'value': 1.0}, {'timestamp': '2023-05-26 11:12:28', 'value': 2.0}, ] mlflow.model_monitoring_repository.load_artifact_dataframe.return_value = mock_reference_data # Act result = await mlflow.get_reference_data(input_data) # Assert mlflow.model_monitoring_repository.load_artifact_dataframe.assert_called_once_with( model_name='test_model', artifact_path='evaluation_data.csv', metadata=metadata['metadata'], ) mock_to_datetime.assert_called_once_with(mock_reference_data.__getitem__.return_value) mock_reference_data.to_dict.assert_called_once_with(orient='records') assert result == mock_reference_data.to_dict.return_value @mark.asyncio async def test_get_reference_data_not_found(mlflow): # Arrange input_data = { **metadata, 'model_name': 'test_model', } mlflow.model_monitoring_repository.load_artifact_dataframe.return_value = None # Act result = await mlflow.get_reference_data(input_data) # Assert mlflow.model_monitoring_repository.load_artifact_dataframe.assert_called_once_with( model_name='test_model', artifact_path='evaluation_data.csv', metadata=metadata['metadata'], ) mlflow.warning.assert_called_once_with( 'Reference data not found for model test_model', metadata['metadata'] ) assert result is None @mark.asyncio async def test_get_reference_data_exception(mlflow): # Arrange input_data = { **metadata, 'model_name': 'test_model', } mlflow.model_monitoring_repository.load_artifact_dataframe.side_effect = Exception( 'Error loading artifact' ) # Act & Assert with raises(Exception) as e: await mlflow.get_reference_data(input_data) assert str(e.value) == 'Error loading artifact' mlflow.model_monitoring_repository.load_artifact_dataframe.assert_called_once_with( model_name='test_model', artifact_path='evaluation_data.csv', metadata=metadata['metadata'], )