"""Unit tests for ExperimentTracking activity.""" from unittest.mock import AsyncMock, MagicMock, patch from pytest import mark, raises from model_manager.activities.experiment_tracking import ExperimentTracking, UpdateType @mark.asyncio @patch('model_manager.activities.experiment_tracking.Postgres.__init__') async def test_update_experiment_run_status_success(mock_postgres_init): """Test successful status update.""" mock_postgres_init.return_value = None # Create instance tracking = ExperimentTracking( host='localhost', port=5432, user='test', password='test', dbname='test', min_connections=1, max_connections=10, logger=MagicMock(), notification_handler=MagicMock(), ) # Mock methods tracking.info = MagicMock() tracking.execute_query = AsyncMock(return_value={'rowcount': 1}) # Test data input_data = { 'metadata': {'workflow_id': 'test-123'}, 'experiment_run_id': 456, 'update_type': UpdateType.STATUS, 'status': 'TRAINING_SUCCESS', } # Execute await tracking.update_experiment_run(input_data) # Assertions tracking.info.assert_called() tracking.execute_query.assert_called_once() call_args = tracking.execute_query.call_args assert 'UPDATE experiment_run' in call_args[0][0] assert 'SET status = %s' in call_args[0][0] @mark.asyncio @patch('model_manager.activities.experiment_tracking.Postgres.__init__') async def test_update_experiment_run_status_with_error_success(mock_postgres_init): """Test successful status update with error message.""" mock_postgres_init.return_value = None tracking = ExperimentTracking( host='localhost', port=5432, user='test', password='test', dbname='test', min_connections=1, max_connections=10, logger=MagicMock(), notification_handler=MagicMock(), ) tracking.info = MagicMock() tracking.execute_query = AsyncMock(return_value={'rowcount': 1}) input_data = { 'metadata': {'workflow_id': 'test-123'}, 'experiment_run_id': 456, 'update_type': UpdateType.STATUS_WITH_ERROR, 'status': 'TRAINING_ERROR', 'error_message': 'Model training failed due to insufficient data', } await tracking.update_experiment_run(input_data) tracking.info.assert_called() tracking.execute_query.assert_called_once() call_args = tracking.execute_query.call_args assert 'UPDATE experiment_run' in call_args[0][0] assert 'SET status = %s, error_message = %s' in call_args[0][0] @mark.asyncio @patch('model_manager.activities.experiment_tracking.Postgres.__init__') async def test_update_experiment_run_error_message_truncation(mock_postgres_init): """Test that error messages longer than 1024 chars are truncated.""" mock_postgres_init.return_value = None tracking = ExperimentTracking( host='localhost', port=5432, user='test', password='test', dbname='test', min_connections=1, max_connections=10, logger=MagicMock(), notification_handler=MagicMock(), ) tracking.info = MagicMock() tracking.execute_query = AsyncMock(return_value={'rowcount': 1}) # Create error message longer than 1024 characters long_error = 'A' * 2000 input_data = { 'metadata': {}, 'experiment_run_id': 456, 'update_type': UpdateType.STATUS_WITH_ERROR, 'status': 'TRAINING_ERROR', 'error_message': long_error, } await tracking.update_experiment_run(input_data) # Check that error message was truncated to 1024 chars call_args = tracking.execute_query.call_args query_params = call_args[0][1] assert len(query_params[1]) == 1024 # error_message is second param @mark.asyncio @patch('model_manager.activities.experiment_tracking.Postgres.__init__') async def test_update_experiment_run_model_saved_success(mock_postgres_init): """Test successful model saved update.""" mock_postgres_init.return_value = None tracking = ExperimentTracking( host='localhost', port=5432, user='test', password='test', dbname='test', min_connections=1, max_connections=10, logger=MagicMock(), notification_handler=MagicMock(), ) tracking.info = MagicMock() tracking.execute_query = AsyncMock(return_value={'rowcount': 1}) input_data = { 'metadata': {}, 'experiment_run_id': 456, 'update_type': UpdateType.MODEL_SAVED, 'run_name': 'experiment-model-123', 'status': 'MLFLOW_SENT', } await tracking.update_experiment_run(input_data) tracking.info.assert_called() tracking.execute_query.assert_called_once() call_args = tracking.execute_query.call_args assert 'UPDATE experiment_run' in call_args[0][0] assert 'SET run_name = %s, status = %s' in call_args[0][0] @mark.asyncio @patch('model_manager.activities.experiment_tracking.Postgres.__init__') async def test_update_experiment_run_missing_status_raises_error(mock_postgres_init): """Test that missing status parameter raises ValueError.""" mock_postgres_init.return_value = None tracking = ExperimentTracking( host='localhost', port=5432, user='test', password='test', dbname='test', min_connections=1, max_connections=10, logger=MagicMock(), notification_handler=MagicMock(), ) tracking.send_notification = MagicMock() tracking.error = MagicMock() input_data = { 'metadata': {}, 'experiment_run_id': 456, 'update_type': UpdateType.STATUS, # Missing 'status' parameter } with raises(RuntimeError, match='Error updating experiment run'): await tracking.update_experiment_run(input_data) tracking.send_notification.assert_called_once() @mark.asyncio @patch('model_manager.activities.experiment_tracking.Postgres.__init__') async def test_update_experiment_run_missing_error_message_raises_error(mock_postgres_init): """Test that missing error_message parameter raises ValueError.""" mock_postgres_init.return_value = None tracking = ExperimentTracking( host='localhost', port=5432, user='test', password='test', dbname='test', min_connections=1, max_connections=10, logger=MagicMock(), notification_handler=MagicMock(), ) tracking.send_notification = MagicMock() tracking.error = MagicMock() input_data = { 'metadata': {}, 'experiment_run_id': 456, 'update_type': UpdateType.STATUS_WITH_ERROR, 'status': 'TRAINING_ERROR', # Missing 'error_message' parameter } with raises(RuntimeError, match='Error updating experiment run'): await tracking.update_experiment_run(input_data) tracking.send_notification.assert_called_once() @mark.asyncio @patch('model_manager.activities.experiment_tracking.Postgres.__init__') async def test_update_experiment_run_missing_run_name_raises_error(mock_postgres_init): """Test that missing run_name parameter raises ValueError.""" mock_postgres_init.return_value = None tracking = ExperimentTracking( host='localhost', port=5432, user='test', password='test', dbname='test', min_connections=1, max_connections=10, logger=MagicMock(), notification_handler=MagicMock(), ) tracking.send_notification = MagicMock() tracking.error = MagicMock() input_data = { 'metadata': {}, 'experiment_run_id': 456, 'update_type': UpdateType.MODEL_SAVED, # Missing 'run_name' parameter } with raises(RuntimeError, match='Error updating experiment run'): await tracking.update_experiment_run(input_data) tracking.send_notification.assert_called_once() @mark.asyncio @patch('model_manager.activities.experiment_tracking.Postgres.__init__') async def test_update_experiment_run_invalid_update_type_raises_error(mock_postgres_init): """Test that invalid update_type raises ValueError.""" mock_postgres_init.return_value = None tracking = ExperimentTracking( host='localhost', port=5432, user='test', password='test', dbname='test', min_connections=1, max_connections=10, logger=MagicMock(), notification_handler=MagicMock(), ) tracking.send_notification = MagicMock() tracking.error = MagicMock() input_data = { 'metadata': {}, 'experiment_run_id': 456, 'update_type': 'INVALID_TYPE', 'status': 'TRAINING_SUCCESS', } with raises(RuntimeError, match='Error updating experiment run'): await tracking.update_experiment_run(input_data) tracking.send_notification.assert_called_once() @mark.asyncio @patch('model_manager.activities.experiment_tracking.Postgres.__init__') async def test_update_experiment_run_no_rows_updated_raises_error(mock_postgres_init): """Test that zero rows updated raises ValueError.""" mock_postgres_init.return_value = None tracking = ExperimentTracking( host='localhost', port=5432, user='test', password='test', dbname='test', min_connections=1, max_connections=10, logger=MagicMock(), notification_handler=MagicMock(), ) tracking.info = MagicMock() tracking.execute_query = AsyncMock(return_value={'rowcount': 0}) tracking.send_notification = MagicMock() tracking.error = MagicMock() input_data = { 'metadata': {}, 'experiment_run_id': 999, # Non-existent ID 'update_type': UpdateType.STATUS, 'status': 'TRAINING_SUCCESS', } with raises(RuntimeError, match='Error updating experiment run'): await tracking.update_experiment_run(input_data) tracking.send_notification.assert_called_once() @mark.asyncio @patch('model_manager.activities.experiment_tracking.Postgres.__init__') async def test_update_experiment_run_sends_notification_on_error(mock_postgres_init): """Test that notification is sent when update fails.""" mock_postgres_init.return_value = None tracking = ExperimentTracking( host='localhost', port=5432, user='test', password='test', dbname='test', min_connections=1, max_connections=10, logger=MagicMock(), notification_handler=MagicMock(), ) tracking.info = MagicMock() tracking.execute_query = AsyncMock(side_effect=Exception('Database error')) tracking.send_notification = MagicMock() tracking.error = MagicMock() input_data = { 'metadata': {'workflow_id': 'test-123'}, 'experiment_run_id': 456, 'update_type': UpdateType.STATUS, 'status': 'TRAINING_SUCCESS', } with raises(RuntimeError): await tracking.update_experiment_run(input_data) # Verify notification was sent tracking.send_notification.assert_called_once() call_args = tracking.send_notification.call_args assert call_args[1]['notification_id'] == 'UPDATE_EXPERIMENT_RUN_ERROR' assert call_args[1]['metadata'] == {'workflow_id': 'test-123'} def test_update_type_enum_values(): """Test UpdateType enum has correct values.""" assert UpdateType.STATUS == 'status' assert UpdateType.STATUS_WITH_ERROR == 'status_with_error' assert UpdateType.MODEL_SAVED == 'model_saved' # Tests for __del__ method - 100% coverage @patch('model_manager.activities.experiment_tracking.Postgres.__init__') def test___del___with_engine_and_parent_del_exists(mock_postgres_init): """Test __del__ calls parent destructor when engine exists and parent has __del__.""" mock_postgres_init.return_value = None tracking = ExperimentTracking( host='localhost', port=5432, user='test', password='test', dbname='test', min_connections=1, max_connections=10, logger=MagicMock(), notification_handler=MagicMock(), ) # Add engine attribute to simulate Postgres initialization tracking.engine = MagicMock() # Track if parent __del__ was called parent_del_called = [] def mock_parent_del(self): """Mock parent __del__ that tracks when it's called.""" parent_del_called.append(True) # Patch parent class to have __del__ method with patch.object(type(tracking).__bases__[0], '__del__', mock_parent_del, create=True): # Trigger __del__ - should call parent __del__ (line 103) tracking.__del__() # Verify parent __del__ was called (line 103 executed) assert len(parent_del_called) == 1, 'Parent __del__ should have been called once' @patch('model_manager.activities.experiment_tracking.Postgres.__init__') def test___del___without_engine(mock_postgres_init): """Test __del__ does not call parent destructor when engine attribute is missing.""" mock_postgres_init.return_value = None tracking = ExperimentTracking( host='localhost', port=5432, user='test', password='test', dbname='test', min_connections=1, max_connections=10, logger=MagicMock(), notification_handler=MagicMock(), ) # Ensure engine attribute does not exist if hasattr(tracking, 'engine'): delattr(tracking, 'engine') # Mock parent __del__ to track if it's called mock_parent_del = MagicMock() with patch.object(type(tracking).__bases__[0], '__del__', mock_parent_del, create=True): # Trigger __del__ - should NOT call parent __del__ (line 100 is False) tracking.__del__() # Verify parent __del__ was NOT called mock_parent_del.assert_not_called() @patch('model_manager.activities.experiment_tracking.Postgres.__init__') def test___del___when_parent_has_no_del_method(mock_postgres_init): """Test __del__ handles case when parent class has no __del__ method - covers line 102 false branch.""" mock_postgres_init.return_value = None tracking = ExperimentTracking( host='localhost', port=5432, user='test', password='test', dbname='test', min_connections=1, max_connections=10, logger=MagicMock(), notification_handler=MagicMock(), ) # Add engine attribute to pass the first hasattr check (line 100) tracking.engine = MagicMock() # Create a mock class without __del__ method class MockSuperWithoutDel: """Mock class that explicitly does not have __del__ method.""" pass # Patch super() to return an instance without __del__ mock_super_instance = MockSuperWithoutDel() with patch('builtins.super', return_value=mock_super_instance): # Trigger __del__ - should handle the case when hasattr(super(), '__del__') is False (line 102) try: tracking.__del__() # Test passes - the false branch of line 102 was executed without error except Exception as e: raise AssertionError( f'__del__ should handle super() without __del__ method, but raised: {e}' ) from e @patch('model_manager.activities.experiment_tracking.Postgres.__init__') def test___del___handles_exception_from_parent_del(mock_postgres_init): """Test __del__ handles exceptions from parent destructor gracefully - covers line 104.""" mock_postgres_init.return_value = None tracking = ExperimentTracking( host='localhost', port=5432, user='test', password='test', dbname='test', min_connections=1, max_connections=10, logger=MagicMock(), notification_handler=MagicMock(), ) # Add engine attribute tracking.engine = MagicMock() # Mock parent __del__ to raise an exception mock_parent_del = MagicMock(side_effect=RuntimeError('Cleanup failed')) with patch.object(type(tracking).__bases__[0], '__del__', mock_parent_del, create=True): # Trigger __del__ - should catch exception and not propagate it (line 104-106) try: tracking.__del__() # Test passes if no exception is raised except Exception as e: raise AssertionError( f'__del__ should handle exceptions gracefully, but raised: {e}' ) from e @patch('model_manager.activities.experiment_tracking.Postgres.__init__') def test___del___handles_attribute_error_from_parent_del(mock_postgres_init): """Test __del__ handles AttributeError from parent destructor - covers line 104 exception handling.""" mock_postgres_init.return_value = None tracking = ExperimentTracking( host='localhost', port=5432, user='test', password='test', dbname='test', min_connections=1, max_connections=10, logger=MagicMock(), notification_handler=MagicMock(), ) # Add engine attribute tracking.engine = MagicMock() # Mock parent __del__ to raise AttributeError mock_parent_del = MagicMock(side_effect=AttributeError('engine not found')) with patch.object(type(tracking).__bases__[0], '__del__', mock_parent_del, create=True): # Trigger __del__ - should catch AttributeError and not propagate it try: tracking.__del__() # Test passes if no exception is raised except Exception as e: raise AssertionError( f'__del__ should handle AttributeError gracefully, but raised: {e}' ) from e