SIENTIAPDE-1182
Implement prediction store policy handling in Gates activity - Added a new method `get_prediction_store_policy` to validate and parse the prediction store policy. - Updated `format_prediction` method to utilize the new policy handling, allowing for sorting of predictions based on the specified policy. - Enhanced test coverage for the new policy handling, including various scenarios for valid and invalid policies. - Removed the obsolete `coverage.sh` script.
This commit is contained in:
@@ -322,15 +322,68 @@ async def test_mlflow_content_gate_with_filter(gates_activity):
|
||||
gates_activity.send_notification.assert_called()
|
||||
|
||||
|
||||
def test_get_prediction_store_policy_invalid_policy(gates_activity):
|
||||
# Arrange
|
||||
prediction_store_policy = 'INVALID_POLICY'
|
||||
|
||||
# Act
|
||||
policy_type, policy_value = gates_activity.get_prediction_store_policy(
|
||||
prediction_store_policy, metadata)
|
||||
|
||||
# Assert
|
||||
assert policy_type == 'lts'
|
||||
assert policy_value == 1
|
||||
|
||||
|
||||
def test_get_prediction_store_policy_invalid_policy_value(gates_activity):
|
||||
# Arrange
|
||||
prediction_store_policy = 'abc:INVALID_VALUE'
|
||||
|
||||
# Act
|
||||
policy_type, policy_value = gates_activity.get_prediction_store_policy(
|
||||
prediction_store_policy, metadata)
|
||||
|
||||
# Assert
|
||||
assert policy_type == 'lts'
|
||||
assert policy_value == 1
|
||||
|
||||
|
||||
def test_get_prediction_store_policy_valid_policy_type(gates_activity):
|
||||
# Arrange
|
||||
prediction_store_policy = 'abc:1'
|
||||
|
||||
# Act
|
||||
policy_type, policy_value = gates_activity.get_prediction_store_policy(
|
||||
prediction_store_policy, metadata)
|
||||
|
||||
# Assert
|
||||
assert policy_type == 'lts'
|
||||
assert policy_value == 1
|
||||
|
||||
|
||||
def test_get_prediction_store_policy_valid_policy(gates_activity):
|
||||
# Arrange
|
||||
prediction_store_policy = 'erl:1'
|
||||
|
||||
# Act
|
||||
policy_type, policy_value = gates_activity.get_prediction_store_policy(
|
||||
prediction_store_policy, metadata)
|
||||
|
||||
# Assert
|
||||
assert policy_type == 'erl'
|
||||
assert policy_value == 1
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_format_prediction(gates_activity):
|
||||
async def test_format_prediction_no_timestamp(gates_activity):
|
||||
# Arrange
|
||||
input_data = {
|
||||
**metadata,
|
||||
'data': {'prediction': [1], 'response_time': [0.1]},
|
||||
'timestamp': '2023-05-26 11:12:27',
|
||||
'model_id': 'test_model',
|
||||
'prediction_confidence': 0.9
|
||||
'prediction_confidence': 0.9,
|
||||
'prediction_store_policy': 'lts:1'
|
||||
}
|
||||
|
||||
# Act
|
||||
@@ -346,6 +399,83 @@ async def test_format_prediction(gates_activity):
|
||||
assert result['comments'] == {0: ""}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_format_prediction_with_timestamp_erl(gates_activity):
|
||||
# Arrange
|
||||
input_data = {
|
||||
**metadata,
|
||||
'data': {'prediction': [1, 2, 3],
|
||||
'response_time': [0.1, 0.2, 0.3],
|
||||
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29']},
|
||||
'model_id': 'test_model',
|
||||
'prediction_confidence': 0.9,
|
||||
'prediction_store_policy': 'erl:2'
|
||||
}
|
||||
|
||||
# Act
|
||||
result = await gates_activity.format_prediction(input_data)
|
||||
|
||||
# Assert
|
||||
assert result['prediction'] == {0: 2, 1: 1}
|
||||
assert result['response_time'] == {0: 0.2, 1: 0.1}
|
||||
assert result['timestamp'] == {
|
||||
0: '2023-05-26 11:12:28', 1: '2023-05-26 11:12:27'}
|
||||
assert result['model_id'] == {0: 'test_model', 1: 'test_model'}
|
||||
assert result['prediction_confidence'] == {0: 0.9, 1: 0.9}
|
||||
assert result['prediction_status'] == {0: 'Good', 1: 'Good'}
|
||||
assert result['comments'] == {0: "", 1: ""}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_format_prediction_with_timestamp_lts(gates_activity):
|
||||
# Arrange
|
||||
input_data = {
|
||||
**metadata,
|
||||
'data': {'prediction': [1, 2, 3],
|
||||
'response_time': [0.1, 0.2, 0.3],
|
||||
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29']},
|
||||
'model_id': 'test_model',
|
||||
'prediction_confidence': 0.9,
|
||||
'prediction_store_policy': 'lts:2'
|
||||
}
|
||||
|
||||
# Act
|
||||
result = await gates_activity.format_prediction(input_data)
|
||||
|
||||
# Assert
|
||||
assert result['prediction'] == {0: 3, 1: 2}
|
||||
assert result['response_time'] == {0: 0.3, 1: 0.2}
|
||||
assert result['timestamp'] == {
|
||||
0: '2023-05-26 11:12:29', 1: '2023-05-26 11:12:28'}
|
||||
assert result['model_id'] == {0: 'test_model', 1: 'test_model'}
|
||||
assert result['prediction_confidence'] == {0: 0.9, 1: 0.9}
|
||||
assert result['prediction_status'] == {0: 'Good', 1: 'Good'}
|
||||
assert result['comments'] == {0: "", 1: ""}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_format_prediction_with_timestamp_invalid_policy(gates_activity):
|
||||
# Arrange
|
||||
input_data = {
|
||||
**metadata,
|
||||
'data': {'prediction': [1, 2, 3],
|
||||
'response_time': [0.1, 0.2, 0.3],
|
||||
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29']},
|
||||
'model_id': 'test_model',
|
||||
'prediction_confidence': 0.9,
|
||||
'prediction_store_policy': 'lts:2'
|
||||
}
|
||||
gates_activity.get_prediction_store_policy = MagicMock(
|
||||
return_value=('invalid', 1))
|
||||
|
||||
try:
|
||||
result = await gates_activity.format_prediction(input_data)
|
||||
except ValueError as e:
|
||||
assert str(e) == "Invalid policy type: invalid"
|
||||
else:
|
||||
assert False, "Expected ValueError"
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_format_default_prediction(gates_activity):
|
||||
# Arrange
|
||||
|
||||
@@ -82,7 +82,8 @@ async def test_request_transform(mock_max, mock_dataframe, mlflow):
|
||||
}
|
||||
|
||||
# Mock the transform response
|
||||
expected_response = {'prediction': [0.5, 0.6]}
|
||||
expected_response = {'prediction': [0.5, 0.6], 'timestamp': [
|
||||
'2024-01-01', '2024-01-02']}
|
||||
mlflow.model_monitoring_repository.transform.return_value = expected_response
|
||||
|
||||
mock_dataframe.return_value.sort_values.return_value = mock_dataframe.return_value
|
||||
@@ -98,7 +99,7 @@ async def test_request_transform(mock_max, mock_dataframe, mlflow):
|
||||
)
|
||||
mock_dataframe = mock_dataframe.return_value.pivot.return_value
|
||||
mock_dataframe.fillna.assert_called_once_with(np.nan, inplace=True)
|
||||
mock_dataframe.reset_index.assert_called_once()
|
||||
# mock_dataframe.reset_index.assert_called_once()
|
||||
mock_dataframe.columns.name = None
|
||||
|
||||
# Verify the response
|
||||
|
||||
@@ -35,7 +35,8 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
|
||||
"schema": "test_schema",
|
||||
"table_name": "test_table",
|
||||
"opc_servers": ["test_server"],
|
||||
"opc_output_config": {"test": "config"}
|
||||
"opc_output_config": {"test": "config"},
|
||||
"prediction_store_policy": "erl:1"
|
||||
}
|
||||
|
||||
await format_and_export_prediction.run(input_data)
|
||||
@@ -48,6 +49,7 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
|
||||
'timestamp': input_data['timestamp'],
|
||||
'model_id': input_data['model_id'],
|
||||
'prediction_confidence': input_data['prediction_confidence'],
|
||||
'prediction_store_policy': input_data['prediction_store_policy'],
|
||||
**metadata
|
||||
},
|
||||
retry_policy=ANY,
|
||||
|
||||
Reference in New Issue
Block a user