Merge pull request #4 from Aignosi/feature/SIENTIAPDE-1250

SIENTIAPDE-1250: Implement Experiment Tracking Activity and Refactor Test Structure

.github/workflows/quality-gate.yml

@@ -201,7 +201,7 @@ jobs:
         uses: actions/cache@v3
         with:
           path: ~/.cache/pip
-          key: ${{ runner.os }}-pip-${{ hashFiles(steps.prepare-requirements.outputs.PROCESSED_REQUIREMENTS_FILE) }}
+          key: ${{ runner.os }}-pip-${{ hashFiles('requirements.txt', 'requirements-dev.txt') }}
           restore-keys: |
             ${{ runner.os }}-pip-
 

README.md

@@ -149,13 +149,20 @@ The Model Manager system uses a Temporal-based workflow architecture with clear
 
 #### **Activities (`model_manager/activities/`)**
 - **Activities**: Main activity orchestrator combining all functionality through multiple inheritance
+- **ExperimentTracking**: ML experiment lifecycle tracking and database operations (extends Postgres)
+  - Unified `update_experiment_run()` method for all experiment status updates
+  - Support for three update types: STATUS, STATUS_WITH_ERROR, MODEL_SAVED
+  - Automatic error message truncation (1024 chars)
+  - Connection pooling and retry logic via Postgres base class
 - **Gates**: Data quality validation and filtering mechanisms
 - **MLFlow**: Model transformation and prediction operations
+- **MinIO**: Object storage operations for file management
 - **Key Features**:
   - Multiple inheritance pattern for unified activity interface
   - Configurable filter policies and validation rules
   - MLFlow model serving integration with configurable flavors
   - Comprehensive error handling and notification integration
+  - Experiment tracking with automatic status management
 
 #### **Data Services (`model_manager/utils/`)**
 - **Connectors Config**: Environment variable-based configuration management
@@ -799,11 +806,10 @@ The `.github/workflows/quality-gate.yml` workflow automatically runs all validat
 ### Test Structure
 ```
 tests/
-├── laborious/
-│   ├── activities/      # Activity implementation tests
-│   ├── workflows/       # Workflow orchestration tests
-│   ├── utils/          # Utility function tests
-│   └── worker/         # Worker tests
+├── activities/      # Activity implementation tests
+├── workflows/       # Workflow orchestration tests
+├── utils/          # Utility function tests
+└── worker/         # Worker tests
 ```
 
 ### Test Execution
@@ -815,8 +821,8 @@ pip install pytest pytest-cov pytest-asyncio
 pytest --cov=model_manager --cov-report=html
 
 # Run specific test modules
-pytest tests/laborious/activities/test_gates.py
-pytest tests/laborious/workflows/test_predictions_batch.py
+pytest tests/activities/test_gates.py
+pytest tests/workflows/test_predictions_batch.py
 ```
 
 ## Monitoring and Metrics
@@ -981,7 +987,7 @@ The project maintains **99%+ code coverage** with comprehensive unit and integra
 pytest tests/ --cov=model_manager --cov-report=term-missing --cov-report=xml --cov-report=html
 
 # Run specific test file
-pytest tests/laborious/activities/test_gates.py -v
+pytest tests/activities/test_gates.py -v
 
 # Run with coverage visualization
 pytest tests/ --cov=model_manager --cov-report=xml
@@ -1089,11 +1095,10 @@ model_manager/
 - **Test structure**:
   ```
   tests/
-  ├── laborious/
-  │   ├── activities/      # Activity tests
-  │   ├── workflows/       # Workflow tests
-  │   ├── utils/          # Utility tests
-  │   └── worker/         # Worker tests
+  ├── activities/      # Activity tests
+  ├── workflows/       # Workflow tests
+  ├── utils/          # Utility tests
+  └── worker/         # Worker tests
   ```
 
 ## Troubleshooting

model_manager/activities/activities.py

@@ -5,14 +5,14 @@ with workflow.unsafe.imports_passed_through():
 
     from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
     from sientia_do.observability.logger import Logger
-    from sientia_do.temporal.activities.postgres import Postgres
 
+    from model_manager.activities.experiment_tracking import ExperimentTracking
     from model_manager.activities.gates import Gates
     from model_manager.activities.minio import MinIO
     from model_manager.activities.mlflow import MLFlow
 
 
-class Activities(Postgres, MLFlow, MinIO, Gates):
+class Activities(ExperimentTracking, MLFlow, MinIO, Gates):
     """
     Main activities orchestrator for the Model Manager system.
 
@@ -21,7 +21,7 @@ class Activities(Postgres, MLFlow, MinIO, Gates):
     MLFlow model interactions, MinIO storage operations, and data quality validation.
 
     The class implements multiple inheritance to combine specialized functionality:
-    - Postgres: Database operations and data persistence
+    - ExperimentTracking: ML experiment lifecycle tracking and database operations (extends Postgres)
     - MLFlow: Model inference and transformation operations
     - MinIO: Object storage operations (file upload/download/delete)
     - Gates: Data quality validation and filtering mechanisms
@@ -62,7 +62,7 @@ class Activities(Postgres, MLFlow, MinIO, Gates):
             Exception: If any parent class initialization fails
         """
         # Initialize parent classes
-        Postgres.__init__(
+        ExperimentTracking.__init__(
             self,
             host=postgres_config['host'],
             port=postgres_config['port'],
@@ -107,10 +107,10 @@ class Activities(Postgres, MLFlow, MinIO, Gates):
         Gracefully shutdown all activities and clean up resources.
 
         This method ensures proper cleanup of all resources including:
-        - PostgreSQL connection pools
+        - PostgreSQL connection pools (via ExperimentTracking)
         - Any other resources that need explicit cleanup
 
         The method should be called before the application terminates to ensure
         proper resource cleanup and prevent resource leaks.
         """
-        Postgres.close(self)
+        ExperimentTracking.close(self)

model_manager/activities/experiment_tracking.py

@@ -0,0 +1,224 @@
+"""
+Experiment tracking activities for managing ML experiment lifecycle.
+
+This module provides activities for tracking and updating experiment run status
+in the PostgreSQL database, extending the base Postgres activity with specialized
+methods for experiment management.
+"""
+
+from temporalio import activity, workflow
+
+with workflow.unsafe.imports_passed_through():
+    import traceback
+    from datetime import UTC, datetime
+    from enum import Enum
+    from typing import Any
+
+    from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
+    from sientia_do.notifications.models import NotificationLevel
+    from sientia_do.observability.logger import Logger
+    from sientia_do.temporal.activities.postgres import Postgres
+
+
+class UpdateType(str, Enum):
+    """Types of experiment run updates."""
+
+    STATUS = 'status'
+    STATUS_WITH_ERROR = 'status_with_error'
+    MODEL_SAVED = 'model_saved'
+
+
+class ExperimentTracking(Postgres):
+    """
+    Activity for tracking ML experiment lifecycle and status updates.
+
+    This activity extends the Postgres activity to provide specialized methods
+    for managing experiment runs, including status updates, error tracking, and
+    model registration. It maintains the experiment lifecycle from initialization
+    through training, model saving, and cleanup.
+
+    The activity uses the existing Postgres connection pool and adds experiment-specific
+    operations with proper error handling and notifications.
+
+    Attributes:
+        logger (Logger): Logger instance for observability
+        notification_handler (NotificationHandler): Handler for sending notifications
+    """
+
+    def __init__(
+        self,
+        host: str,
+        port: int,
+        user: str,
+        password: str,
+        dbname: str,
+        min_connections: int,
+        max_connections: int,
+        logger: Logger,
+        notification_handler: NotificationHandler,
+    ):
+        """
+        Initialize ExperimentTracking activity with database configuration.
+
+        Args:
+            host: PostgreSQL server hostname
+            port: PostgreSQL server port
+            user: Database user
+            password: Database password
+            dbname: Database name
+            min_connections: Minimum connections in pool
+            max_connections: Maximum connections in pool
+            logger: Logger instance for observability
+            notification_handler: Notification handler for alerts
+
+        Raises:
+            ConnectionError: If database connection cannot be established
+        """
+        super().__init__(
+            host=host,
+            port=port,
+            user=user,
+            password=password,
+            dbname=dbname,
+            min_connections=min_connections,
+            max_connections=max_connections,
+            logger=logger,
+            notification_handler=notification_handler,
+        )
+        self.logger = logger
+        self.notification_handler = notification_handler
+
+    @activity.defn(name='update_experiment_run')
+    async def update_experiment_run(self, input_data: dict[str, Any]) -> None:
+        """
+        Update experiment run with status, errors, or model information.
+
+        This activity provides a unified interface for all experiment run updates,
+        supporting different update types through a single method. It automatically
+        selects the appropriate SQL query based on the update type and parameters.
+
+        Args:
+            input_data: Configuration for experiment run update operation
+                Required keys:
+                    - metadata (dict): Workflow execution metadata
+                    - experiment_run_id (int): Unique identifier for the experiment run
+                    - update_type (str): Type of update (status, status_with_error, model_saved)
+                Optional keys:
+                    - status (str): New status for the experiment run
+                    - error_message (str): Error message if update failed
+                    - run_name (str): MLFlow run name if model was saved
+
+        Raises:
+            ValueError: If required parameters are missing for the update type
+            RuntimeError: If update operation fails
+
+        Example:
+            # Update status only
+            await update_experiment_run({
+                'metadata': {},
+                'experiment_run_id': 123,
+                'update_type': 'status',
+                'status': 'TRAINING_SUCCESS'
+            })
+
+            # Update with error
+            await update_experiment_run({
+                'metadata': {},
+                'experiment_run_id': 123,
+                'update_type': 'status_with_error',
+                'status': 'TRAINING_ERROR',
+                'error_message': 'Model training failed: insufficient data'
+            })
+
+            # Update with model saved
+            await update_experiment_run({
+                'metadata': {},
+                'experiment_run_id': 123,
+                'update_type': 'model_saved',
+                'run_name': 'experiment-model-5'
+            })
+        """
+        metadata = input_data.get('metadata', {})
+        experiment_run_id = input_data['experiment_run_id']
+        update_type = input_data['update_type']
+        status = input_data.get('status')
+        error_message = input_data.get('error_message')
+        run_name = input_data.get('run_name')
+
+        try:
+            self.info(
+                f'Updating experiment run {experiment_run_id} with type: {update_type}', metadata
+            )
+
+            # Validate parameters based on update type
+            query_params: tuple[Any, ...]
+            if update_type == UpdateType.STATUS:
+                if not status:
+                    raise ValueError('status is required for STATUS update type')
+                sql_query = """
+                    UPDATE experiment_run
+                    SET status = %s, updated_at = %s
+                    WHERE id = %s
+                """
+                query_params = (status, datetime.now(UTC), experiment_run_id)
+
+            elif update_type == UpdateType.STATUS_WITH_ERROR:
+                if not status or not error_message:
+                    raise ValueError(
+                        'status and error_message are required for STATUS_WITH_ERROR update type'
+                    )
+
+                # Truncate the error_message to 1024 characters if necessary
+                if len(error_message) > 1024:
+                    error_message = error_message[:1024]
+
+                sql_query = """
+                    UPDATE experiment_run
+                    SET status = %s, error_message = %s, updated_at = %s
+                    WHERE id = %s
+                """
+                query_params = (
+                    status,
+                    error_message,
+                    datetime.now(UTC),
+                    experiment_run_id,
+                )
+
+            elif update_type == UpdateType.MODEL_SAVED:
+                if not run_name:
+                    raise ValueError('run_name is required for MODEL_SAVED update type')
+                sql_query = """
+                    UPDATE experiment_run
+                    SET run_name = %s, status = %s, updated_at = %s
+                    WHERE id = %s
+                """
+                query_params = (run_name, status, datetime.now(UTC), experiment_run_id)
+
+            else:
+                raise ValueError(f'Invalid update_type: {update_type}')
+
+            # Execute update query
+            result = await self.execute_query(sql_query, query_params)
+
+            if result.get('rowcount', 0) == 0:
+                error_msg = f'No rows updated for experiment run {experiment_run_id}'
+                raise ValueError(error_msg)
+
+            self.info(
+                f'Successfully updated experiment run {experiment_run_id} with type {update_type}',
+                metadata,
+            )
+
+        except Exception as e:  # noqa: BLE001
+            error_msg = f'Error updating experiment run - ID: {experiment_run_id}, Type: {update_type}, Error: {str(e)}'
+            trace = traceback.format_exc()
+            self.send_notification(
+                metadata=metadata,
+                notification_id='UPDATE_EXPERIMENT_RUN_ERROR',
+                message=error_msg,
+                block='update_experiment_run',
+                level=NotificationLevel.ERROR,
+                attachment_content=trace,
+            )
+            self.error(trace, metadata=metadata)
+            raise RuntimeError(error_msg) from e

tests/activities/test_activities.py

@@ -1,18 +1,20 @@
 from unittest.mock import ANY, MagicMock, patch
 
 from pytest import mark
-from sientia_do.temporal.activities.postgres import Postgres
 
 from model_manager.activities.activities import Activities
+from model_manager.activities.experiment_tracking import ExperimentTracking
 from model_manager.activities.gates import Gates
 from model_manager.activities.mlflow import MLFlow
 
 
-@patch('model_manager.activities.activities.Postgres.__init__')
+@patch('model_manager.activities.activities.ExperimentTracking.__init__')
 @patch('model_manager.activities.activities.MLFlow.__init__')
 @patch('model_manager.activities.activities.MinIO.__init__')
 @patch('model_manager.activities.activities.Gates.__init__')
-def test___init__(mock_gates_init, mock_minio_init, mock_mlflow_init, mock_postgres_init):
+def test___init__(
+    mock_gates_init, mock_minio_init, mock_mlflow_init, mock_experiment_tracking_init
+):
     postgres_config = {
         'host': 'localhost',
         'port': 5432,
@@ -49,11 +51,11 @@ def test___init__(mock_gates_init, mock_minio_init, mock_mlflow_init, mock_postg
     )
 
     assert isinstance(activities, Activities)
-    assert isinstance(activities, Postgres)
+    assert isinstance(activities, ExperimentTracking)
     assert isinstance(activities, MLFlow)
     assert isinstance(activities, Gates)
 
-    mock_postgres_init.assert_called_once_with(
+    mock_experiment_tracking_init.assert_called_once_with(
         ANY,
         host=postgres_config['host'],
         port=postgres_config['port'],
@@ -97,9 +99,9 @@ def test___init__(mock_gates_init, mock_minio_init, mock_mlflow_init, mock_postg
 
 
 @mark.asyncio
-@patch('model_manager.activities.activities.Postgres', return_value=MagicMock())
+@patch('model_manager.activities.activities.ExperimentTracking', return_value=MagicMock())
 @patch('model_manager.activities.activities.MLFlow', return_value=MagicMock())
-async def test_shutdown(_mock_mlflow_init, mock_postgres_init):
+async def test_shutdown(_mock_mlflow_init, mock_experiment_tracking_init):
     postgres_config = {
         'host': 'localhost',
         'port': 5432,
@@ -136,4 +138,4 @@ async def test_shutdown(_mock_mlflow_init, mock_postgres_init):
     )
 
     await activities.shutdown()
-    mock_postgres_init.close.assert_called_once()
+    mock_experiment_tracking_init.close.assert_called_once()

tests/activities/test_experiment_tracking.py

@@ -0,0 +1,385 @@
+"""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'