SIENTIAPDE-1273
Enhance MLFlowRepository and Activities classes with new methods and metrics - Added `check_artifact_exists` method to MLFlowRepository for verifying artifact presence in the MLflow Model Registry. - Implemented `get_prediction_data` method in MLFlowRepository to retrieve prediction data from models. - Updated Activities class to integrate ModelMetrics for improved metrics handling. - Enhanced tests for artifact existence checks and prediction data retrieval, ensuring robust coverage for new functionalities. - Updated various workflows to include `transform_table_name` in input data for better data handling.
This commit is contained in:
641
tests/laborious/activities/test_model_metrics.py
Normal file
641
tests/laborious/activities/test_model_metrics.py
Normal file
@@ -0,0 +1,641 @@
|
||||
from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
|
||||
|
||||
from pandas import DataFrame
|
||||
from pytest import fixture, mark
|
||||
from sientia_do.notifications.models import NotificationLevel
|
||||
|
||||
from laborious.activities.model_metrics import ModelMetrics
|
||||
|
||||
|
||||
@fixture
|
||||
def model_metrics_activity():
|
||||
model_metrics = ModelMetrics(
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
metrics_controller=AsyncMock(),
|
||||
)
|
||||
model_metrics.error = MagicMock()
|
||||
model_metrics.debug = MagicMock()
|
||||
model_metrics.info = MagicMock()
|
||||
model_metrics.warning = MagicMock()
|
||||
model_metrics.critical = MagicMock()
|
||||
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.observe_lag = AsyncMock()
|
||||
model_metrics.pod_id = 'test_pod'
|
||||
return model_metrics
|
||||
|
||||
|
||||
metadata = {
|
||||
'metadata': {
|
||||
'model_id': 'test_model',
|
||||
'model_name': 'test_model',
|
||||
'workflow_name': 'test_workflow',
|
||||
'schema_name': 'test_schedule',
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_calculate_drift_invalid_chunk_period(model_metrics_activity):
|
||||
# Arrange
|
||||
input_data = {
|
||||
**metadata,
|
||||
'model_name': 'test_model',
|
||||
'model_id': 'test_model_id',
|
||||
'reference_data': None,
|
||||
'target_data': {
|
||||
'timestamp': ['2023-05-26 11:12:27'],
|
||||
'variable': ['feature1'],
|
||||
'value': [1.0],
|
||||
},
|
||||
'target_name': 'target',
|
||||
'drift_metrics': ['ks_test'],
|
||||
'chunk_period': 'invalid',
|
||||
}
|
||||
|
||||
# Act & Assert
|
||||
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\""
|
||||
model_metrics_activity.error.assert_called_once_with(
|
||||
'Invalid chunk period: invalid', metadata['metadata']
|
||||
)
|
||||
else:
|
||||
raise AssertionError('Expected ValueError')
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.activities.model_metrics.DataFrame')
|
||||
@patch('laborious.activities.model_metrics.to_datetime')
|
||||
async def test_calculate_drift_with_reference_data(
|
||||
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_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.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]}
|
||||
|
||||
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],
|
||||
})
|
||||
|
||||
mock_target_df = MagicMock()
|
||||
mock_target_df.pivot.return_value = mock_target_df
|
||||
mock_target_df.index = ['2023-05-26 11:12:27']
|
||||
mock_target_df.reset_index.return_value = mock_target_df
|
||||
mock_target_df.dropna.return_value = mock_target_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',
|
||||
'model_id': 'test_model_id',
|
||||
'reference_data': reference_data.to_dict(),
|
||||
'target_data': {
|
||||
'timestamp': ['2023-05-26 11:12:27'],
|
||||
'variable': ['feature1'],
|
||||
'value': [1.0],
|
||||
},
|
||||
'target_name': 'target',
|
||||
'drift_metrics': ['ks_test'],
|
||||
'chunk_period': 'min',
|
||||
}
|
||||
|
||||
# Act
|
||||
result = await model_metrics_activity.calculate_drift(input_data)
|
||||
|
||||
# Assert
|
||||
assert isinstance(result, dict)
|
||||
assert result == mock_drift_df.to_dict.return_value
|
||||
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.__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()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.activities.model_metrics.DataFrame')
|
||||
@patch('laborious.activities.model_metrics.to_datetime')
|
||||
async def test_calculate_drift_without_reference_data(
|
||||
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_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.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]}
|
||||
|
||||
model_metrics_activity.get_drift_metrics = AsyncMock(return_value=mock_drift_df)
|
||||
|
||||
target_data_dict = {
|
||||
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'],
|
||||
'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.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',
|
||||
'model_id': 'test_model_id',
|
||||
'reference_data': None,
|
||||
'target_data': target_data_dict,
|
||||
'target_name': 'target',
|
||||
'drift_metrics': ['ks_test'],
|
||||
'chunk_period': 's',
|
||||
}
|
||||
|
||||
# Act
|
||||
result = await model_metrics_activity.calculate_drift(input_data)
|
||||
|
||||
# Assert
|
||||
assert isinstance(result, dict)
|
||||
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(
|
||||
metadata=metadata['metadata'],
|
||||
notification_id='MODEL_METRICS_REFERENCE_DATA_WARNING',
|
||||
message='Using 30% first rows of target data as reference data',
|
||||
block='model_metrics',
|
||||
level=NotificationLevel.WARNING,
|
||||
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.__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()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.activities.model_metrics.DataFrame')
|
||||
@patch('laborious.activities.model_metrics.to_datetime')
|
||||
async def test_calculate_drift_empty_drift_df(
|
||||
mock_to_datetime, mock_dataframe, model_metrics_activity
|
||||
):
|
||||
# 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],
|
||||
})
|
||||
|
||||
mock_target_df = MagicMock()
|
||||
mock_target_df.pivot.return_value = mock_target_df
|
||||
mock_target_df.index = ['2023-05-26 11:12:27']
|
||||
mock_target_df.reset_index.return_value = mock_target_df
|
||||
mock_target_df.dropna.return_value = mock_target_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',
|
||||
'model_id': 'test_model_id',
|
||||
'reference_data': reference_data.to_dict(),
|
||||
'target_data': {
|
||||
'timestamp': ['2023-05-26 11:12:27'],
|
||||
'variable': ['feature1'],
|
||||
'value': [1.0],
|
||||
},
|
||||
'target_name': 'target',
|
||||
'drift_metrics': ['ks_test'],
|
||||
'chunk_period': 'min',
|
||||
}
|
||||
|
||||
# Act
|
||||
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'])
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.activities.model_metrics.DataFrame')
|
||||
@patch('laborious.activities.model_metrics.to_datetime')
|
||||
async def test_calculate_drift_empty_after_timestamp_filter(
|
||||
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_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],
|
||||
})
|
||||
|
||||
mock_target_df = MagicMock()
|
||||
mock_target_df.pivot.return_value = mock_target_df
|
||||
mock_target_df.index = ['2023-05-26 11:12:27']
|
||||
mock_target_df.reset_index.return_value = mock_target_df
|
||||
mock_target_df.dropna.return_value = mock_target_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',
|
||||
'model_id': 'test_model_id',
|
||||
'reference_data': reference_data.to_dict(),
|
||||
'target_data': {
|
||||
'timestamp': ['2023-05-26 11:12:27'],
|
||||
'variable': ['feature1'],
|
||||
'value': [1.0],
|
||||
},
|
||||
'target_name': 'target',
|
||||
'drift_metrics': ['ks_test'],
|
||||
'chunk_period': 'min',
|
||||
}
|
||||
|
||||
# Act
|
||||
result = await model_metrics_activity.calculate_drift(input_data)
|
||||
|
||||
# Assert
|
||||
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']
|
||||
)
|
||||
# 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.__getitem__.assert_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.activities.model_metrics.DataFrame')
|
||||
@patch('laborious.activities.model_metrics.to_datetime')
|
||||
async def test_calculate_drift_success_min(
|
||||
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_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.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]}
|
||||
|
||||
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],
|
||||
})
|
||||
|
||||
mock_target_df = MagicMock()
|
||||
mock_target_df.pivot.return_value = mock_target_df
|
||||
mock_target_df.index = ['2023-05-26 11:12:27']
|
||||
mock_target_df.reset_index.return_value = mock_target_df
|
||||
mock_target_df.dropna.return_value = mock_target_df
|
||||
mock_target_df.__getitem__.return_value.isin.return_value = [True]
|
||||
mock_target_df.drop.return_value.columns = ['feature1']
|
||||
mock_dataframe.return_value = mock_target_df
|
||||
|
||||
input_data = {
|
||||
**metadata,
|
||||
'model_name': 'test_model',
|
||||
'model_id': 'test_model_id',
|
||||
'reference_data': reference_data.to_dict(),
|
||||
'target_data': {
|
||||
'timestamp': ['2023-05-26 11:12:27'],
|
||||
'variable': ['feature1'],
|
||||
'value': [1.0],
|
||||
},
|
||||
'target_name': 'target',
|
||||
'drift_metrics': ['ks_test'],
|
||||
'chunk_period': 'min',
|
||||
}
|
||||
|
||||
# Act
|
||||
result = await model_metrics_activity.calculate_drift(input_data)
|
||||
|
||||
# Assert
|
||||
assert isinstance(result, dict)
|
||||
assert result == mock_drift_df.to_dict.return_value
|
||||
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.__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()
|
||||
|
||||
|
||||
@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
|
||||
):
|
||||
# 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_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.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]}
|
||||
|
||||
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],
|
||||
})
|
||||
|
||||
mock_target_df = MagicMock()
|
||||
mock_target_df.pivot.return_value = mock_target_df
|
||||
mock_target_df.index = ['2023-05-26 11:12:27']
|
||||
mock_target_df.reset_index.return_value = mock_target_df
|
||||
mock_target_df.dropna.return_value = mock_target_df
|
||||
mock_target_df.__getitem__.return_value.isin.return_value = [True]
|
||||
mock_target_df.drop.return_value.columns = ['feature1']
|
||||
mock_dataframe.return_value = mock_target_df
|
||||
|
||||
input_data = {
|
||||
**metadata,
|
||||
'model_name': 'test_model',
|
||||
'model_id': 'test_model_id',
|
||||
'reference_data': reference_data.to_dict(),
|
||||
'target_data': {
|
||||
'timestamp': ['2023-05-26 11:12:27'],
|
||||
'variable': ['feature1'],
|
||||
'value': [1.0],
|
||||
},
|
||||
'target_name': 'target',
|
||||
'drift_metrics': ['ks_test'],
|
||||
'chunk_period': 's',
|
||||
}
|
||||
|
||||
# Act
|
||||
result = await model_metrics_activity.calculate_drift(input_data)
|
||||
|
||||
# Assert
|
||||
assert isinstance(result, dict)
|
||||
assert result == mock_drift_df.to_dict.return_value
|
||||
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.__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()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.activities.model_metrics.DataFrame')
|
||||
@patch('laborious.activities.model_metrics.to_datetime')
|
||||
async def test_calculate_drift_get_drift_metrics_error(
|
||||
mock_to_datetime, mock_dataframe, model_metrics_activity
|
||||
):
|
||||
# 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],
|
||||
})
|
||||
|
||||
mock_target_df = MagicMock()
|
||||
mock_target_df.pivot.return_value = mock_target_df
|
||||
mock_target_df.index = ['2023-05-26 11:12:27']
|
||||
mock_target_df.reset_index.return_value = mock_target_df
|
||||
mock_target_df.dropna.return_value = mock_target_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',
|
||||
'model_id': 'test_model_id',
|
||||
'reference_data': reference_data.to_dict(),
|
||||
'target_data': {
|
||||
'timestamp': ['2023-05-26 11:12:27'],
|
||||
'variable': ['feature1'],
|
||||
'value': [1.0],
|
||||
},
|
||||
'target_name': 'target',
|
||||
'drift_metrics': ['ks_test'],
|
||||
'chunk_period': 'min',
|
||||
}
|
||||
|
||||
# Act
|
||||
result = await model_metrics_activity.calculate_drift(input_data)
|
||||
|
||||
# Assert
|
||||
assert result == {}
|
||||
model_metrics_activity.error.assert_called_once_with(
|
||||
'Error getting drift metrics: Get drift metrics error',
|
||||
metadata['metadata']
|
||||
)
|
||||
model_metrics_activity.send_notification_async.assert_called_once_with(
|
||||
metadata=metadata['metadata'],
|
||||
notification_id='MODEL_METRICS_GET_DRIFT_METRICS_ERROR',
|
||||
message='Error getting drift metrics: Get drift metrics error',
|
||||
block='model_metrics',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=ANY,
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.activities.model_metrics.to_datetime')
|
||||
@patch('laborious.activities.model_metrics.time.time')
|
||||
@patch('laborious.activities.model_metrics.ModelAnalysis')
|
||||
@patch('laborious.activities.model_metrics.metrics')
|
||||
async def test_get_drift_metrics_success(
|
||||
mock_metrics, mock_model_analysis, mock_time, mock_to_datetime, model_metrics_activity
|
||||
):
|
||||
# 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_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_columns = reference_data.drop(
|
||||
columns=['target', 'timestamp'], errors='ignore'
|
||||
).columns
|
||||
|
||||
# Act
|
||||
result = await model_metrics_activity.get_drift_metrics(
|
||||
reference_data=reference_data,
|
||||
target_data=target_data,
|
||||
target_name='target',
|
||||
reference_columns=reference_columns,
|
||||
drift_metrics=['ks_test'],
|
||||
chunk_period='min',
|
||||
metadata=metadata['metadata'],
|
||||
)
|
||||
|
||||
# Assert
|
||||
assert isinstance(result, DataFrame)
|
||||
model_metrics_activity.debug.assert_called()
|
||||
model_metrics_activity.observe_lag.assert_called()
|
||||
model_metrics_activity.emit_metric.assert_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.activities.model_metrics.to_datetime')
|
||||
@patch('laborious.activities.model_metrics.time.time')
|
||||
@patch('laborious.activities.model_metrics.ModelAnalysis')
|
||||
@patch('laborious.activities.model_metrics.metrics')
|
||||
async def test_get_drift_metrics_univariate_error(
|
||||
mock_metrics, mock_model_analysis, mock_time, mock_to_datetime, model_metrics_activity
|
||||
):
|
||||
# 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],
|
||||
})
|
||||
|
||||
reference_columns = reference_data.drop(
|
||||
columns=['target', 'timestamp'], errors='ignore'
|
||||
).columns
|
||||
|
||||
# Act & Assert
|
||||
try:
|
||||
await model_metrics_activity.get_drift_metrics(
|
||||
reference_data=reference_data,
|
||||
target_data=target_data,
|
||||
target_name='target',
|
||||
reference_columns=reference_columns,
|
||||
drift_metrics=['ks_test'],
|
||||
chunk_period='min',
|
||||
metadata=metadata['metadata'],
|
||||
)
|
||||
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']
|
||||
)
|
||||
model_metrics_activity.emit_metric.assert_called_with(
|
||||
metric_object=mock_metrics.MODEL_ANALYZE_ERROR_COUNT,
|
||||
tags=ANY
|
||||
)
|
||||
else:
|
||||
raise AssertionError('Expected Exception')
|
||||
|
||||
Reference in New Issue
Block a user