SIENTIAPDE-1182

Refactor test data structure in gates.py and model_repository.py for improved clarity and consistency

- Updated test cases in test_gates.py to use dictionaries for prediction and response_time, associating values with timestamps.
- Modified test_predict_success in test_model_repository.py to create a DataFrame with named indices for better readability in assertions.
This commit is contained in:
vitor-aignosi
2025-09-04 12:16:14 -03:00
parent 2e420dadba
commit 8ab1b9b437
2 changed files with 47 additions and 13 deletions

View File

@@ -379,8 +379,14 @@ 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',
'data': {
'prediction': {
'2023-05-26 11:12:27': 1
},
'response_time': {
'2023-05-26 11:12:27': 0.1
}
},
'model_id': 'test_model',
'prediction_confidence': 0.9,
'prediction_store_policy': 'lts:1'
@@ -404,9 +410,18 @@ 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']},
'data': {
'prediction': {
'2023-05-26 11:12:27': 1,
'2023-05-26 11:12:28': 2,
'2023-05-26 11:12:29': 3,
},
'response_time': {
'2023-05-26 11:12:27': 0.1,
'2023-05-26 11:12:28': 0.2,
'2023-05-26 11:12:29': 0.3,
}
},
'model_id': 'test_model',
'prediction_confidence': 0.9,
'prediction_store_policy': 'erl:2'
@@ -431,9 +446,18 @@ 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']},
'data': {
'prediction': {
'2023-05-26 11:12:27': 1,
'2023-05-26 11:12:28': 2,
'2023-05-26 11:12:29': 3,
},
'response_time': {
'2023-05-26 11:12:27': 0.1,
'2023-05-26 11:12:28': 0.2,
'2023-05-26 11:12:29': 0.3,
}
},
'model_id': 'test_model',
'prediction_confidence': 0.9,
'prediction_store_policy': 'lts:2'

View File

@@ -59,7 +59,12 @@ def test_transform_error(mlflow_repository):
def test_predict_success(mlflow_repository):
data = 'data'
data = DataFrame({
'feat_1': {
'index_1': 2,
'index_2': 3
}
})
model_name = 'model'
mlflow_repository.model_serving.get_cached_predict.return_value = np.array(
[2, 3]
@@ -71,10 +76,15 @@ def test_predict_success(mlflow_repository):
model_name, data, 1)
assert output['success'] is True
assert output['content'] == {'prediction': {
0: 2,
1: 3
}, 'response_time': ANY}
assert output['content'] == {
'prediction': {
'index_1': 2,
'index_2': 3
}, 'response_time': {
'index_1': ANY,
'index_2': ANY
}
}
def test_predict_error(mlflow_repository):