vitor-aignosi
cefef0b1e9
Enhance MLFlow activity to sort data by created_at and remove duplicates; update tests to reflect changes
132 lines
4.6 KiB
Python
132 lines
4.6 KiB
Python
from unittest.mock import MagicMock, patch
|
|
|
|
import numpy as np
|
|
from pytest import fixture, mark
|
|
from laborious.activities.mlflow import MLFlow
|
|
|
|
|
|
@patch("laborious.activities.mlflow.MLFlowRepository")
|
|
def test___init__(mock_mlflow_repository):
|
|
mlflow = MLFlow(
|
|
mlflow_host="http://localhost",
|
|
mlflow_port=5000,
|
|
mlflow_username="admin",
|
|
mlflow_password="admin",
|
|
logger=MagicMock(),
|
|
notification_handler=MagicMock()
|
|
)
|
|
|
|
assert mlflow.mlflow_host == "http://localhost"
|
|
assert mlflow.mlflow_port == 5000
|
|
assert mlflow.mlflow_username == "admin"
|
|
assert mlflow.mlflow_password == "admin"
|
|
|
|
mock_mlflow_repository.assert_called_once_with(
|
|
"http://localhost:5000", "admin", "admin"
|
|
)
|
|
|
|
|
|
@fixture
|
|
@patch("laborious.activities.mlflow.MLFlowRepository")
|
|
def mlflow(mock_mlflow_repository):
|
|
return MLFlow(
|
|
mlflow_host="http://localhost:5000",
|
|
mlflow_port=5000,
|
|
mlflow_username="admin",
|
|
mlflow_password="admin",
|
|
logger=MagicMock(),
|
|
notification_handler=MagicMock()
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
@patch("laborious.activities.mlflow.DataFrame")
|
|
@patch("laborious.activities.mlflow.max")
|
|
async def test_request_transform(mock_max, mock_dataframe, mlflow):
|
|
mock_max.return_value = '2024-01-02'
|
|
# Mock input data
|
|
input_data = {
|
|
'data': [
|
|
{'timestamp': '2024-01-01', 'variable': 'var1',
|
|
'value': 1.0, 'created_at': '2024-01-01 12:00:00'},
|
|
{'timestamp': '2024-01-01', 'variable': 'var2',
|
|
'value': 2.0, 'created_at': '2024-01-01 12:00:00'},
|
|
{'timestamp': '2024-01-02', 'variable': 'var1',
|
|
'value': 3.0, 'created_at': '2024-01-02 12:00:00'},
|
|
{'timestamp': '2024-01-02', 'variable': 'var2',
|
|
'value': 4.0, 'created_at': '2024-01-02 12:00:00'},
|
|
{'timestamp': '2024-01-02', 'variable': 'var1',
|
|
'value': 1.0, 'created_at': '2024-01-01 12:00:00'},
|
|
{'timestamp': '2024-01-02', 'variable': 'var2',
|
|
'value': 1.0, 'created_at': '2024-01-01 12:00:00'}
|
|
],
|
|
'model_name': 'test_model',
|
|
'model_retention': 30
|
|
}
|
|
|
|
# Mock the transform response
|
|
expected_response = {'prediction': [0.5, 0.6]}
|
|
mlflow.model_monitoring_repository.transform.return_value = expected_response
|
|
|
|
mock_dataframe.return_value.sort_values.return_value = mock_dataframe.return_value
|
|
mock_dataframe.return_value.drop_duplicates.return_value = mock_dataframe.return_value
|
|
|
|
# Call the method
|
|
response_data = await mlflow.request_transform(input_data)
|
|
|
|
# Verify the data was correctly transformed
|
|
mock_dataframe.assert_called_once_with(input_data['data'])
|
|
mock_dataframe.return_value.pivot.assert_called_once_with(
|
|
index='timestamp', columns='variable', values='value'
|
|
)
|
|
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.columns.name = None
|
|
|
|
# Verify the response
|
|
assert response_data == expected_response
|
|
|
|
# Verify the repository was called with correct arguments
|
|
mlflow.model_monitoring_repository.transform.assert_called_once_with(
|
|
'test_model', mock_dataframe, 30
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
@patch("laborious.activities.mlflow.DataFrame")
|
|
@patch("laborious.activities.mlflow.max")
|
|
async def test_request_predict(mock_max, mock_dataframe, mlflow):
|
|
mock_max.return_value = '2024-01-02'
|
|
# Mock input data
|
|
input_data = {
|
|
'data': [
|
|
{'timestamp': '2024-01-01', 'variable': 'var1', 'value': 1.0},
|
|
{'timestamp': '2024-01-01', 'variable': 'var2', 'value': 2.0},
|
|
{'timestamp': '2024-01-02', 'variable': 'var1', 'value': 3.0},
|
|
{'timestamp': '2024-01-02', 'variable': 'var2', 'value': 4.0}
|
|
],
|
|
'model_name': 'test_model',
|
|
'model_retention': 30
|
|
}
|
|
|
|
# Mock the predict response
|
|
expected_response = {'prediction': [0.5, 0.6]}
|
|
mlflow.model_monitoring_repository.predict.return_value = expected_response
|
|
|
|
# Call the method
|
|
response_data = await mlflow.request_predict(input_data)
|
|
|
|
mock_dataframe.assert_called_once_with(input_data['data'])
|
|
mock_dataframe.return_value.replace.assert_called_once_with(
|
|
np.nan, None, inplace=True
|
|
)
|
|
|
|
# Verify the response
|
|
assert response_data == expected_response
|
|
|
|
# Verify the repository was called with correct arguments
|
|
mlflow.model_monitoring_repository.predict.assert_called_once_with(
|
|
'test_model', mock_dataframe.return_value, 30
|
|
)
|