SIENTIAPDE-1273

SIENTIAPDE-1273
Enhance security analysis and SQL injection handling

- Added skip for potential SQL injection false positives in Bandit configuration.
- Updated validate.sh to use the pyproject.toml configuration for Bandit security analysis.
- Refactored code to replace ensure_dataframe utility with direct DataFrame usage in multiple activities, improving clarity and reducing dependencies.
- Removed the deprecated dataframe_utils module to streamline the codebase.

tests/laborious/activities/test_model_metrics.py

@@ -1,4 +1,4 @@
-from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
+from unittest.mock import ANY, AsyncMock, MagicMock, patch
 
 from pandas import DataFrame
 from pytest import fixture, mark
@@ -22,7 +22,13 @@ def model_metrics_activity():
     model_metrics.send_notification = MagicMock()
     model_metrics.send_notification_async = AsyncMock()
     model_metrics.emit_metric = AsyncMock()
-    model_metrics.get_core_labels = MagicMock(return_value={'pod_id': 'test_pod', 'model_name': 'test_model', 'workflow_name': 'test_workflow'})
+    model_metrics.get_core_labels = MagicMock(
+        return_value={
+            'pod_id': 'test_pod',
+            'model_name': 'test_model',
+            'workflow_name': 'test_workflow',
+        }
+    )
     model_metrics.observe_lag = AsyncMock()
     model_metrics.pod_id = 'test_pod'
     return model_metrics
@@ -60,7 +66,7 @@ async def test_calculate_drift_invalid_chunk_period(model_metrics_activity):
     try:
         await model_metrics_activity.calculate_drift(input_data)
     except ValueError as e:
-        assert str(e) == "Invalid chunk period: invalid, must be \"min\" or \"s\""
+        assert str(e) == 'Invalid chunk period: invalid, must be "min" or "s"'
         model_metrics_activity.error.assert_called_once_with(
             'Invalid chunk period: invalid', metadata['metadata']
         )
@@ -76,26 +82,41 @@ async def test_calculate_drift_with_reference_data(
 ):
     # Arrange
     mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27'
-    mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00'
-    
+    mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
+        '2023-05-26 11:12:27+00:00'
+    )
+
     mock_drift_df = MagicMock()
     mock_drift_df.empty = False
     mock_drift_df.drop.return_value = mock_drift_df
     mock_drift_df.__getitem__.return_value.isin.return_value = [True]
     mock_drift_df.__getitem__.return_value = mock_drift_df
-    mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00'
+    mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
+        '2023-05-26 11:12:27+00:00'
+    )
     mock_drift_df.rename.return_value = mock_drift_df
     mock_drift_df.drop_duplicates.return_value = mock_drift_df
-    mock_drift_df.to_dict.return_value = {'method': ['ks_test'], 'value': [0.5], 'feature': ['feature1'], 'timestamp': ['2023-05-26 11:12:27+00:00'], 'model_id': ['test_model_id'], 'accurate': [True]}
-    
+    mock_drift_df.to_dict.return_value = [
+        {
+            'method': 'ks_test',
+            'value': 0.5,
+            'feature': 'feature1',
+            'timestamp': '2023-05-26 11:12:27+00:00',
+            'model_id': 'test_model_id',
+            'accurate': True,
+        }
+    ]
+
     model_metrics_activity.get_drift_metrics = AsyncMock(return_value=mock_drift_df)
 
-    reference_data = DataFrame({
-        'timestamp': ['2023-05-26 11:12:27'],
-        'target': [1.0],
-        'feature1': [1.0],
-    })
-    
+    reference_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'target': [1.0],
+            'feature1': [1.0],
+        }
+    )
+
     mock_target_df = MagicMock()
     mock_target_df.pivot.return_value = mock_target_df
     mock_target_df.index = ['2023-05-26 11:12:27']
@@ -124,16 +145,20 @@ async def test_calculate_drift_with_reference_data(
     result = await model_metrics_activity.calculate_drift(input_data)
 
     # Assert
-    assert isinstance(result, dict)
-    assert result == mock_drift_df.to_dict.return_value
+    assert isinstance(result, list)
+    assert result == mock_drift_df.to_dict.return_value  # type: ignore[comparison-overlap]
     model_metrics_activity.info.assert_called()
     model_metrics_activity.get_drift_metrics.assert_called_once()
     # Verify transformations were called
-    mock_drift_df.drop.assert_called_once_with(columns=['p_value', 'chunk_start_date', 'chunk_end_date'], inplace=True)
+    mock_drift_df.drop.assert_called_once_with(columns=['p_value'], inplace=True)
     mock_drift_df.__getitem__.assert_called()
-    mock_drift_df.rename.assert_called_once_with(columns={'metric': 'method', 'statistic': 'value'}, inplace=True)
-    mock_drift_df.drop_duplicates.assert_called_once_with(subset=['timestamp', 'method', 'feature'], keep='first', inplace=True)
-    mock_drift_df.to_dict.assert_called_once()
+    mock_drift_df.rename.assert_called_once_with(
+        columns={'metric': 'method', 'statistic': 'value'}, inplace=True
+    )
+    mock_drift_df.drop_duplicates.assert_called_once_with(
+        subset=['timestamp', 'method', 'feature'], keep='first', inplace=True
+    )
+    mock_drift_df.to_dict.assert_called_once_with(orient='records')
 
 
 @mark.asyncio
@@ -144,18 +169,31 @@ async def test_calculate_drift_without_reference_data(
 ):
     # Arrange
     mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27'
-    mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00'
-    
+    mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
+        '2023-05-26 11:12:27+00:00'
+    )
+
     mock_drift_df = MagicMock()
     mock_drift_df.empty = False
     mock_drift_df.drop.return_value = mock_drift_df
     mock_drift_df.__getitem__.return_value.isin.return_value = [True]
     mock_drift_df.__getitem__.return_value = mock_drift_df
-    mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00'
+    mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
+        '2023-05-26 11:12:27+00:00'
+    )
     mock_drift_df.rename.return_value = mock_drift_df
     mock_drift_df.drop_duplicates.return_value = mock_drift_df
-    mock_drift_df.to_dict.return_value = {'method': ['ks_test'], 'value': [0.5], 'feature': ['feature1'], 'timestamp': ['2023-05-26 11:12:27+00:00'], 'model_id': ['test_model_id'], 'accurate': [False]}
-    
+    mock_drift_df.to_dict.return_value = [
+        {
+            'method': 'ks_test',
+            'value': 0.5,
+            'feature': 'feature1',
+            'timestamp': '2023-05-26 11:12:27+00:00',
+            'model_id': 'test_model_id',
+            'accurate': False,
+        }
+    ]
+
     model_metrics_activity.get_drift_metrics = AsyncMock(return_value=mock_drift_df)
 
     target_data_dict = {
@@ -163,19 +201,25 @@ async def test_calculate_drift_without_reference_data(
         'variable': ['feature1', 'feature1', 'feature1'],
         'value': [1.0, 2.0, 3.0],
     }
-    
+
     mock_target_df = MagicMock()
     mock_target_df.pivot.return_value = mock_target_df
     mock_target_df.index = ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29']
     mock_target_df.reset_index.return_value = mock_target_df
     mock_target_df.dropna.return_value = mock_target_df
     mock_target_df.sort_values.return_value = mock_target_df
-    mock_target_df.head.return_value = DataFrame({'timestamp': ['2023-05-26 11:12:27'], 'feature1': [1.0]})
-    mock_target_df.__getitem__.return_value.apply.return_value = ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29']
+    mock_target_df.head.return_value = DataFrame(
+        {'timestamp': ['2023-05-26 11:12:27'], 'feature1': [1.0]}
+    )
+    mock_target_df.__getitem__.return_value.apply.return_value = [
+        '2023-05-26 11:12:27',
+        '2023-05-26 11:12:28',
+        '2023-05-26 11:12:29',
+    ]
     mock_target_df.drop.return_value.columns = ['feature1']
     mock_dataframe.return_value = mock_target_df
     mock_dataframe.side_effect = lambda x=None: mock_target_df if x is not None else mock_target_df
-    
+
     input_data = {
         **metadata,
         'model_name': 'test_model',
@@ -191,7 +235,7 @@ async def test_calculate_drift_without_reference_data(
     result = await model_metrics_activity.calculate_drift(input_data)
 
     # Assert
-    assert isinstance(result, dict)
+    assert isinstance(result, list)
     assert result == mock_drift_df.to_dict.return_value
     model_metrics_activity.warning.assert_called()
     model_metrics_activity.send_notification_async.assert_called_once_with(
@@ -203,11 +247,15 @@ async def test_calculate_drift_without_reference_data(
         attachment_content=ANY,
     )
     # Verify transformations were called
-    mock_drift_df.drop.assert_called_once_with(columns=['p_value', 'chunk_start_date', 'chunk_end_date'], inplace=True)
+    mock_drift_df.drop.assert_called_once_with(columns=['p_value'], inplace=True)
     mock_drift_df.__getitem__.assert_called()
-    mock_drift_df.rename.assert_called_once_with(columns={'metric': 'method', 'statistic': 'value'}, inplace=True)
-    mock_drift_df.drop_duplicates.assert_called_once_with(subset=['timestamp', 'method', 'feature'], keep='first', inplace=True)
-    mock_drift_df.to_dict.assert_called_once()
+    mock_drift_df.rename.assert_called_once_with(
+        columns={'metric': 'method', 'statistic': 'value'}, inplace=True
+    )
+    mock_drift_df.drop_duplicates.assert_called_once_with(
+        subset=['timestamp', 'method', 'feature'], keep='first', inplace=True
+    )
+    mock_drift_df.to_dict.assert_called_once_with(orient='records')
 
 
 @mark.asyncio
@@ -218,15 +266,17 @@ async def test_calculate_drift_empty_drift_df(
 ):
     # Arrange
     mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27'
-    
+
     model_metrics_activity.get_drift_metrics = AsyncMock(return_value=DataFrame())
 
-    reference_data = DataFrame({
-        'timestamp': ['2023-05-26 11:12:27'],
-        'target': [1.0],
-        'feature1': [1.0],
-    })
-    
+    reference_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'target': [1.0],
+            'feature1': [1.0],
+        }
+    )
+
     mock_target_df = MagicMock()
     mock_target_df.pivot.return_value = mock_target_df
     mock_target_df.index = ['2023-05-26 11:12:27']
@@ -235,7 +285,7 @@ async def test_calculate_drift_empty_drift_df(
     mock_target_df.__getitem__.return_value.apply.return_value = ['2023-05-26 11:12:27']
     mock_target_df.drop.return_value.columns = ['feature1']
     mock_dataframe.return_value = mock_target_df
-  
+
     input_data = {
         **metadata,
         'model_name': 'test_model',
@@ -255,8 +305,10 @@ async def test_calculate_drift_empty_drift_df(
     result = await model_metrics_activity.calculate_drift(input_data)
 
     # Assert
-    assert result == {}
-    model_metrics_activity.warning.assert_called_with('No drift metrics found', metadata['metadata'])
+    assert result == []
+    model_metrics_activity.warning.assert_called_with(
+        'No drift metrics found', metadata['metadata']
+    )
 
 
 @mark.asyncio
@@ -267,35 +319,37 @@ async def test_calculate_drift_empty_after_timestamp_filter(
 ):
     # Arrange
     mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27'
-    
+
     mock_drift_df = MagicMock()
     mock_drift_df.empty = False
     mock_drift_df.drop.return_value = mock_drift_df
-    
+
     # Set up __getitem__ to handle filtering - timestamp access returns series with isin=False
     # and filtering returns empty DataFrame
     mock_timestamp_series = MagicMock()
     mock_timestamp_series.isin.return_value = [False]
     mock_empty_df = MagicMock()
     mock_empty_df.empty = True
-    
+
     def getitem_side_effect(key):
         if key == 'timestamp':
             return mock_timestamp_series
         else:
             # This is the filtering operation - return empty DataFrame
             return mock_empty_df
-    
+
     mock_drift_df.__getitem__.side_effect = getitem_side_effect
 
     model_metrics_activity.get_drift_metrics = AsyncMock(return_value=mock_drift_df)
 
-    reference_data = DataFrame({
-        'timestamp': ['2023-05-26 11:12:27'],
-        'target': [1.0],
-        'feature1': [1.0],
-    })
-    
+    reference_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'target': [1.0],
+            'feature1': [1.0],
+        }
+    )
+
     mock_target_df = MagicMock()
     mock_target_df.pivot.return_value = mock_target_df
     mock_target_df.index = ['2023-05-26 11:12:27']
@@ -324,13 +378,13 @@ async def test_calculate_drift_empty_after_timestamp_filter(
     result = await model_metrics_activity.calculate_drift(input_data)
 
     # Assert
-    assert result == {}
+    assert result == []
     model_metrics_activity.warning.assert_called_with(
         'No drift metrics found after dropping rows where timestamp is not in target data',
-        metadata['metadata']
+        metadata['metadata'],
     )
     # Verify transformations were called
-    mock_drift_df.drop.assert_called_once_with(columns=['p_value', 'chunk_start_date', 'chunk_end_date'], inplace=True)
+    mock_drift_df.drop.assert_called_once_with(columns=['p_value'], inplace=True)
     mock_drift_df.__getitem__.assert_called()
 
 
@@ -342,26 +396,41 @@ async def test_calculate_drift_success_min(
 ):
     # Arrange
     mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27'
-    mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00'
-    
+    mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
+        '2023-05-26 11:12:27+00:00'
+    )
+
     mock_drift_df = MagicMock()
     mock_drift_df.empty = False
     mock_drift_df.drop.return_value = mock_drift_df
     mock_drift_df.__getitem__.return_value.isin.return_value = [True]
     mock_drift_df.__getitem__.return_value = mock_drift_df
-    mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00'
+    mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
+        '2023-05-26 11:12:27+00:00'
+    )
     mock_drift_df.rename.return_value = mock_drift_df
     mock_drift_df.drop_duplicates.return_value = mock_drift_df
-    mock_drift_df.to_dict.return_value = {'method': ['ks_test'], 'value': [0.5], 'feature': ['feature1'], 'timestamp': ['2023-05-26 11:12:27+00:00'], 'model_id': ['test_model_id'], 'accurate': [True]}
-    
+    mock_drift_df.to_dict.return_value = [
+        {
+            'method': 'ks_test',
+            'value': 0.5,
+            'feature': 'feature1',
+            'timestamp': '2023-05-26 11:12:27+00:00',
+            'model_id': 'test_model_id',
+            'accurate': True,
+        }
+    ]
+
     model_metrics_activity.get_drift_metrics = AsyncMock(return_value=mock_drift_df)
 
-    reference_data = DataFrame({
-        'timestamp': ['2023-05-26 11:12:27'],
-        'target': [1.0],
-        'feature1': [1.0],
-    })
-    
+    reference_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'target': [1.0],
+            'feature1': [1.0],
+        }
+    )
+
     mock_target_df = MagicMock()
     mock_target_df.pivot.return_value = mock_target_df
     mock_target_df.index = ['2023-05-26 11:12:27']
@@ -390,46 +459,63 @@ async def test_calculate_drift_success_min(
     result = await model_metrics_activity.calculate_drift(input_data)
 
     # Assert
-    assert isinstance(result, dict)
-    assert result == mock_drift_df.to_dict.return_value
+    assert isinstance(result, list)
+    assert result == mock_drift_df.to_dict.return_value  # type: ignore[comparison-overlap]
     model_metrics_activity.info.assert_called()
     model_metrics_activity.get_drift_metrics.assert_called_once()
     # Verify transformations were called
-    mock_drift_df.drop.assert_called_once_with(columns=['p_value', 'chunk_start_date', 'chunk_end_date'], inplace=True)
+    mock_drift_df.drop.assert_called_once_with(columns=['p_value'], inplace=True)
     mock_drift_df.__getitem__.assert_called()
-    mock_drift_df.rename.assert_called_once_with(columns={'metric': 'method', 'statistic': 'value'}, inplace=True)
-    mock_drift_df.drop_duplicates.assert_called_once_with(subset=['timestamp', 'method', 'feature'], keep='first', inplace=True)
-    mock_drift_df.to_dict.assert_called_once()
+    mock_drift_df.rename.assert_called_once_with(
+        columns={'metric': 'method', 'statistic': 'value'}, inplace=True
+    )
+    mock_drift_df.drop_duplicates.assert_called_once_with(
+        subset=['timestamp', 'method', 'feature'], keep='first', inplace=True
+    )
+    mock_drift_df.to_dict.assert_called_once_with(orient='records')
 
 
 @mark.asyncio
 @patch('laborious.activities.model_metrics.DataFrame')
 @patch('laborious.activities.model_metrics.to_datetime')
-async def test_calculate_drift_success_s(
-    mock_to_datetime, mock_dataframe, model_metrics_activity
-):
+async def test_calculate_drift_success_s(mock_to_datetime, mock_dataframe, model_metrics_activity):
     # Arrange
     mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27'
-    mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00'
-    
+    mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
+        '2023-05-26 11:12:27+00:00'
+    )
+
     mock_drift_df = MagicMock()
     mock_drift_df.empty = False
     mock_drift_df.drop.return_value = mock_drift_df
     mock_drift_df.__getitem__.return_value.isin.return_value = [True]
     mock_drift_df.__getitem__.return_value = mock_drift_df
-    mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00'
+    mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
+        '2023-05-26 11:12:27+00:00'
+    )
     mock_drift_df.rename.return_value = mock_drift_df
     mock_drift_df.drop_duplicates.return_value = mock_drift_df
-    mock_drift_df.to_dict.return_value = {'method': ['ks_test'], 'value': [0.5], 'feature': ['feature1'], 'timestamp': ['2023-05-26 11:12:27+00:00'], 'model_id': ['test_model_id'], 'accurate': [True]}
-    
+    mock_drift_df.to_dict.return_value = [
+        {
+            'method': 'ks_test',
+            'value': 0.5,
+            'feature': 'feature1',
+            'timestamp': '2023-05-26 11:12:27+00:00',
+            'model_id': 'test_model_id',
+            'accurate': True,
+        }
+    ]
+
     model_metrics_activity.get_drift_metrics = AsyncMock(return_value=mock_drift_df)
 
-    reference_data = DataFrame({
-        'timestamp': ['2023-05-26 11:12:27'],
-        'target': [1.0],
-        'feature1': [1.0],
-    })
-    
+    reference_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'target': [1.0],
+            'feature1': [1.0],
+        }
+    )
+
     mock_target_df = MagicMock()
     mock_target_df.pivot.return_value = mock_target_df
     mock_target_df.index = ['2023-05-26 11:12:27']
@@ -458,16 +544,20 @@ async def test_calculate_drift_success_s(
     result = await model_metrics_activity.calculate_drift(input_data)
 
     # Assert
-    assert isinstance(result, dict)
-    assert result == mock_drift_df.to_dict.return_value
+    assert isinstance(result, list)
+    assert result == mock_drift_df.to_dict.return_value  # type: ignore[comparison-overlap]
     model_metrics_activity.info.assert_called()
     model_metrics_activity.get_drift_metrics.assert_called_once()
     # Verify transformations were called
-    mock_drift_df.drop.assert_called_once_with(columns=['p_value', 'chunk_start_date', 'chunk_end_date'], inplace=True)
+    mock_drift_df.drop.assert_called_once_with(columns=['p_value'], inplace=True)
     mock_drift_df.__getitem__.assert_called()
-    mock_drift_df.rename.assert_called_once_with(columns={'metric': 'method', 'statistic': 'value'}, inplace=True)
-    mock_drift_df.drop_duplicates.assert_called_once_with(subset=['timestamp', 'method', 'feature'], keep='first', inplace=True)
-    mock_drift_df.to_dict.assert_called_once()
+    mock_drift_df.rename.assert_called_once_with(
+        columns={'metric': 'method', 'statistic': 'value'}, inplace=True
+    )
+    mock_drift_df.drop_duplicates.assert_called_once_with(
+        subset=['timestamp', 'method', 'feature'], keep='first', inplace=True
+    )
+    mock_drift_df.to_dict.assert_called_once_with(orient='records')
 
 
 @mark.asyncio
@@ -478,15 +568,19 @@ async def test_calculate_drift_get_drift_metrics_error(
 ):
     # Arrange
     mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27'
-    
-    model_metrics_activity.get_drift_metrics = AsyncMock(side_effect=Exception('Get drift metrics error'))
 
-    reference_data = DataFrame({
-        'timestamp': ['2023-05-26 11:12:27'],
-        'target': [1.0],
-        'feature1': [1.0],
-    })
-    
+    model_metrics_activity.get_drift_metrics = AsyncMock(
+        side_effect=Exception('Get drift metrics error')
+    )
+
+    reference_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'target': [1.0],
+            'feature1': [1.0],
+        }
+    )
+
     mock_target_df = MagicMock()
     mock_target_df.pivot.return_value = mock_target_df
     mock_target_df.index = ['2023-05-26 11:12:27']
@@ -515,10 +609,9 @@ async def test_calculate_drift_get_drift_metrics_error(
     result = await model_metrics_activity.calculate_drift(input_data)
 
     # Assert
-    assert result == {}
+    assert result == []
     model_metrics_activity.error.assert_called_once_with(
-        'Error getting drift metrics: Get drift metrics error',
-        metadata['metadata']
+        'Error getting drift metrics: Get drift metrics error', metadata['metadata']
     )
     model_metrics_activity.send_notification_async.assert_called_once_with(
         metadata=metadata['metadata'],
@@ -540,30 +633,36 @@ async def test_get_drift_metrics_success(
 ):
     # Arrange
     mock_time.return_value = 1000.0
-    
-    mock_drift_df = DataFrame({
-        'timestamp': ['2023-05-26 11:12:27'],
-        'metric': ['ks_test'],
-        'statistic': [0.5],
-        'feature': ['feature1'],
-    })
-    
+
+    mock_drift_df = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'metric': ['ks_test'],
+            'statistic': [0.5],
+            'feature': ['feature1'],
+        }
+    )
+
     mock_model_analysis.return_value.detect_univariate_drift.return_value = MagicMock()
     mock_model_analysis.return_value.detect_multivariate_drift.return_value = MagicMock()
     mock_model_analysis.return_value.get_drift_metrics_dataframe.return_value = mock_drift_df
 
-    reference_data = DataFrame({
-        'timestamp': ['2023-05-26 11:12:27'],
-        'target': [1.0],
-        'feature1': [1.0],
-    })
-    
-    target_data = DataFrame({
-        'timestamp': ['2023-05-26 11:12:27'],
-        'target': [1.0],
-        'feature1': [1.0],
-    })
-    
+    reference_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'target': [1.0],
+            'feature1': [1.0],
+        }
+    )
+
+    target_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'target': [1.0],
+            'feature1': [1.0],
+        }
+    )
+
     reference_columns = reference_data.drop(
         columns=['target', 'timestamp'], errors='ignore'
     ).columns
@@ -596,21 +695,27 @@ async def test_get_drift_metrics_univariate_error(
 ):
     # Arrange
     mock_time.return_value = 1000.0
-    
-    mock_model_analysis.return_value.detect_univariate_drift.side_effect = Exception('Univariate drift error')
 
-    reference_data = DataFrame({
-        'timestamp': ['2023-05-26 11:12:27'],
-        'target': [1.0],
-        'feature1': [1.0],
-    })
-    
-    target_data = DataFrame({
-        'timestamp': ['2023-05-26 11:12:27'],
-        'target': [1.0],
-        'feature1': [1.0],
-    })
-    
+    mock_model_analysis.return_value.detect_univariate_drift.side_effect = Exception(
+        'Univariate drift error'
+    )
+
+    reference_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'target': [1.0],
+            'feature1': [1.0],
+        }
+    )
+
+    target_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27'],
+            'target': [1.0],
+            'feature1': [1.0],
+        }
+    )
+
     reference_columns = reference_data.drop(
         columns=['target', 'timestamp'], errors='ignore'
     ).columns
@@ -629,28 +734,26 @@ async def test_get_drift_metrics_univariate_error(
     except Exception as e:
         assert str(e) == 'Univariate drift error'
         model_metrics_activity.error.assert_called_once_with(
-            'Error detecting univariate drift: Univariate drift error',
-            metadata['metadata']
+            'Error detecting univariate drift: Univariate drift error', metadata['metadata']
         )
         model_metrics_activity.emit_metric.assert_called_with(
-            metric_object=mock_metrics.MODEL_ANALYZE_ERROR_COUNT,
-            tags=ANY
+            metric_object=mock_metrics.MODEL_ANALYZE_ERROR_COUNT, tags=ANY
         )
     else:
         raise AssertionError('Expected Exception')
 
 
 @mark.asyncio
-async def test_calculate_simple_metrics_success_all_metrics(
-    model_metrics_activity
-):
+async def test_calculate_simple_metrics_success_all_metrics(model_metrics_activity):
     # Arrange
-    target_data = DataFrame({
-        'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'],
-        'target': [1.0, 2.0, 3.0],
-        'prediction': [1.1, 2.1, 2.9],
-    })
-    
+    target_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'],
+            'target': [1.0, 2.0, 3.0],
+            'prediction': [1.1, 2.1, 2.9],
+        }
+    )
+
     input_data = {
         **metadata,
         'model_id': 'test_model_id',
@@ -673,23 +776,23 @@ async def test_calculate_simple_metrics_success_all_metrics(
     assert all(data_size == 3 for data_size in result['data_size'].values)
     assert all(interval_minutes == 5 for interval_minutes in result['interval_minutes'].values)
     model_metrics_activity.info.assert_called_once_with(
-        'Calculating simple metrics for model test_model_id: [\'rmse\', \'mse\', \'mae\', \'r2\']',
-        metadata['metadata']
+        "Calculating simple metrics for model test_model_id: ['rmse', 'mse', 'mae', 'r2']",
+        metadata['metadata'],
     )
     model_metrics_activity.debug.assert_called_once()
 
 
 @mark.asyncio
-async def test_calculate_simple_metrics_success_rmse_only(
-    model_metrics_activity
-):
+async def test_calculate_simple_metrics_success_rmse_only(model_metrics_activity):
     # Arrange
-    target_data = DataFrame({
-        'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
-        'target': [1.0, 2.0],
-        'prediction': [1.1, 2.1],
-    })
-    
+    target_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
+            'target': [1.0, 2.0],
+            'prediction': [1.1, 2.1],
+        }
+    )
+
     input_data = {
         **metadata,
         'model_id': 'test_model_id',
@@ -709,22 +812,21 @@ async def test_calculate_simple_metrics_success_rmse_only(
     assert result['data_size'].values[0] == 2
     assert result['interval_minutes'].values[0] == 5
     model_metrics_activity.info.assert_called_once_with(
-        'Calculating simple metrics for model test_model_id: [\'rmse\']',
-        metadata['metadata']
+        "Calculating simple metrics for model test_model_id: ['rmse']", metadata['metadata']
     )
 
 
 @mark.asyncio
-async def test_calculate_simple_metrics_success_mse_only(
-    model_metrics_activity
-):
+async def test_calculate_simple_metrics_success_mse_only(model_metrics_activity):
     # Arrange
-    target_data = DataFrame({
-        'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
-        'target': [1.0, 2.0],
-        'prediction': [1.1, 2.1],
-    })
-    
+    target_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
+            'target': [1.0, 2.0],
+            'prediction': [1.1, 2.1],
+        }
+    )
+
     input_data = {
         **metadata,
         'model_id': 'test_model_id',
@@ -744,22 +846,21 @@ async def test_calculate_simple_metrics_success_mse_only(
     assert result['data_size'].values[0] == 2
     assert result['interval_minutes'].values[0] == 5
     model_metrics_activity.info.assert_called_once_with(
-        'Calculating simple metrics for model test_model_id: [\'mse\']',
-        metadata['metadata']
+        "Calculating simple metrics for model test_model_id: ['mse']", metadata['metadata']
     )
 
 
 @mark.asyncio
-async def test_calculate_simple_metrics_success_mae_only(
-    model_metrics_activity
-):
+async def test_calculate_simple_metrics_success_mae_only(model_metrics_activity):
     # Arrange
-    target_data = DataFrame({
-        'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
-        'target': [1.0, 2.0],
-        'prediction': [1.1, 2.1],
-    })
-    
+    target_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
+            'target': [1.0, 2.0],
+            'prediction': [1.1, 2.1],
+        }
+    )
+
     input_data = {
         **metadata,
         'model_id': 'test_model_id',
@@ -779,22 +880,21 @@ async def test_calculate_simple_metrics_success_mae_only(
     assert result['data_size'].values[0] == 2
     assert result['interval_minutes'].values[0] == 5
     model_metrics_activity.info.assert_called_once_with(
-        'Calculating simple metrics for model test_model_id: [\'mae\']',
-        metadata['metadata']
+        "Calculating simple metrics for model test_model_id: ['mae']", metadata['metadata']
     )
 
 
 @mark.asyncio
-async def test_calculate_simple_metrics_success_r2_only(
-    model_metrics_activity
-):
+async def test_calculate_simple_metrics_success_r2_only(model_metrics_activity):
     # Arrange
-    target_data = DataFrame({
-        'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
-        'target': [1.0, 2.0],
-        'prediction': [1.1, 2.1],
-    })
-    
+    target_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
+            'target': [1.0, 2.0],
+            'prediction': [1.1, 2.1],
+        }
+    )
+
     input_data = {
         **metadata,
         'model_id': 'test_model_id',
@@ -814,23 +914,22 @@ async def test_calculate_simple_metrics_success_r2_only(
     assert result['data_size'].values[0] == 2
     assert result['interval_minutes'].values[0] == 5
     model_metrics_activity.info.assert_called_once_with(
-        'Calculating simple metrics for model test_model_id: [\'r2\']',
-        metadata['metadata']
+        "Calculating simple metrics for model test_model_id: ['r2']", metadata['metadata']
     )
 
 
 @mark.asyncio
-async def test_calculate_simple_metrics_r2_zero_ss_tot(
-    model_metrics_activity
-):
+async def test_calculate_simple_metrics_r2_zero_ss_tot(model_metrics_activity):
     # Arrange
     # All target values are the same, so ss_tot will be 0
-    target_data = DataFrame({
-        'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
-        'target': [1.0, 1.0],
-        'prediction': [1.1, 1.1],
-    })
-    
+    target_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
+            'target': [1.0, 1.0],
+            'prediction': [1.1, 1.1],
+        }
+    )
+
     input_data = {
         **metadata,
         'model_id': 'test_model_id',
@@ -851,22 +950,21 @@ async def test_calculate_simple_metrics_r2_zero_ss_tot(
     assert result['data_size'].values[0] == 2
     assert result['interval_minutes'].values[0] == 5
     model_metrics_activity.info.assert_called_once_with(
-        'Calculating simple metrics for model test_model_id: [\'r2\']',
-        metadata['metadata']
+        "Calculating simple metrics for model test_model_id: ['r2']", metadata['metadata']
     )
 
 
 @mark.asyncio
-async def test_calculate_simple_metrics_success_multiple_metrics_subset(
-    model_metrics_activity
-):
+async def test_calculate_simple_metrics_success_multiple_metrics_subset(model_metrics_activity):
     # Arrange
-    target_data = DataFrame({
-        'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'],
-        'target': [1.0, 2.0, 3.0],
-        'prediction': [1.1, 2.1, 2.9],
-    })
-    
+    target_data = DataFrame(
+        {
+            'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'],
+            'target': [1.0, 2.0, 3.0],
+            'prediction': [1.1, 2.1, 2.9],
+        }
+    )
+
     input_data = {
         **metadata,
         'model_id': 'test_model_id',
@@ -887,7 +985,5 @@ async def test_calculate_simple_metrics_success_multiple_metrics_subset(
     assert all(data_size == 3 for data_size in result['data_size'].values)
     assert all(interval_minutes == 5 for interval_minutes in result['interval_minutes'].values)
     model_metrics_activity.info.assert_called_once_with(
-        'Calculating simple metrics for model test_model_id: [\'rmse\', \'mae\']',
-        metadata['metadata']
+        "Calculating simple metrics for model test_model_id: ['rmse', 'mae']", metadata['metadata']
     )
-