SIENTIAPDE-1243: Initial commit of the model manager project, adding core files and configurations.
This commit introduces the initial project structure, including: - .env.example: Example environment configuration. - .github/workflows/quality-gate.yml: CI workflow for quality checks. - .gitignore: Specifies intentionally untracked files that Git should ignore. - Makefile: Automation of tasks like docker builds. - README.md: Project documentation. - Source code for model management, activities, utils, worker and workflows. - Test suite. - Dockerfile for the simulator. - sonar-project.properties: SonarQube configuration file. - values.yaml: Helm chart values for deployment.
This commit is contained in:
Bruno Domingues
502
tests/laborious/utils/repository/test_model_repository.py
Normal file
502
tests/laborious/utils/repository/test_model_repository.py
Normal file
@@ -0,0 +1,502 @@
|
||||
from unittest.mock import ANY, MagicMock, call, patch
|
||||
import numpy as np
|
||||
from pandas import DataFrame
|
||||
import pytest
|
||||
from datetime import datetime, timezone
|
||||
from pandas import Timestamp
|
||||
from laborious.utils.repository.model_repository import MLFlowRepository
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mlflow_repository():
|
||||
with patch('laborious.utils.repository.model_repository.ModelServing',
|
||||
autospec=True) as mock_model_serving:
|
||||
mock_instance = mock_model_serving.return_value
|
||||
mock_instance.get_transformed_data = MagicMock()
|
||||
|
||||
repo = MLFlowRepository(
|
||||
host='http://localhost:5000',
|
||||
username='admin',
|
||||
password='admin',
|
||||
logger=MagicMock()
|
||||
)
|
||||
return repo
|
||||
|
||||
|
||||
metadata = {
|
||||
"metadata": {
|
||||
"model_id": "test_model",
|
||||
"model_name": "test_model",
|
||||
"workflow_name": "test_workflow",
|
||||
"schema_name": "test_schedule",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class Any:
|
||||
pass
|
||||
|
||||
|
||||
invalid_cases = [
|
||||
(
|
||||
{
|
||||
'value': {
|
||||
'2024-01-01 12:00:00': 1,
|
||||
2024: 2
|
||||
}
|
||||
}
|
||||
),
|
||||
(
|
||||
{
|
||||
'value': {
|
||||
'2024-01-01': 1,
|
||||
'2024-01-02': 2
|
||||
}
|
||||
}
|
||||
),
|
||||
(
|
||||
{
|
||||
'value': {
|
||||
Any(): 1,
|
||||
Any(): 2
|
||||
}
|
||||
}
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("data", invalid_cases)
|
||||
def test_detect_and_parse_datetime_index_error_cases(mlflow_repository, data):
|
||||
input_data = DataFrame(
|
||||
data
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError) as e:
|
||||
mlflow_repository.detect_and_parse_datetime_index(
|
||||
input_data, metadata['metadata'])
|
||||
|
||||
assert str(e) == "Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format %Y-%m-%d %H:%M:%S"
|
||||
|
||||
|
||||
valid_cases = [
|
||||
(
|
||||
{
|
||||
'value': {
|
||||
'2024-01-01 12:00:00+0000': 1,
|
||||
'2024-01-02 12:00:00+0000': 2
|
||||
}
|
||||
}, ['2024-01-01 12:00:00+0000', '2024-01-02 12:00:00+0000']
|
||||
),
|
||||
(
|
||||
{
|
||||
'value': {
|
||||
datetime(2025, 1, 1, 12, 0, 0, tzinfo=timezone.utc): 1,
|
||||
datetime(2025, 1, 2, 12, 0, 0, tzinfo=timezone.utc): 2
|
||||
}
|
||||
}, ['2025-01-01 12:00:00+0000', '2025-01-02 12:00:00+0000']
|
||||
),
|
||||
(
|
||||
{
|
||||
'value': {
|
||||
Timestamp(2026, 1, 1, 12, 0, 0, tzinfo=timezone.utc): 1,
|
||||
Timestamp(2026, 1, 2, 12, 0, 0, tzinfo=timezone.utc): 2
|
||||
}
|
||||
}, ['2026-01-01 12:00:00+0000', '2026-01-02 12:00:00+0000']
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("data,expected", valid_cases)
|
||||
def test_detect_and_parse_datetime_index_valid_format(mlflow_repository, data, expected):
|
||||
input_data = DataFrame(data)
|
||||
|
||||
response = mlflow_repository.detect_and_parse_datetime_index(
|
||||
input_data, metadata['metadata'])
|
||||
|
||||
assert response.index.tolist() == expected
|
||||
|
||||
|
||||
def test_transform_success(mlflow_repository):
|
||||
data = MagicMock()
|
||||
model_name = 'model'
|
||||
|
||||
mlflow_repository.detect_and_parse_datetime_index = MagicMock()
|
||||
|
||||
output = mlflow_repository.transform(
|
||||
model_name, data, {}, metadata['metadata'])
|
||||
|
||||
mlflow_repository.model_serving.get_cached_transform.assert_called_once_with(
|
||||
model_name, data, 0, 'sklearn', False, 'model', 'predict')
|
||||
|
||||
mlflow_repository.detect_and_parse_datetime_index.assert_called_once_with(
|
||||
mlflow_repository.model_serving.get_cached_transform.return_value, metadata['metadata'])
|
||||
|
||||
assert output == {
|
||||
'success': True,
|
||||
'content': mlflow_repository.detect_and_parse_datetime_index.return_value.to_dict.return_value
|
||||
}
|
||||
|
||||
|
||||
def test_transform_error(mlflow_repository):
|
||||
data = MagicMock()
|
||||
model_name = 'model'
|
||||
|
||||
mlflow_repository.model_serving.get_cached_transform.side_effect = Exception(
|
||||
'error')
|
||||
|
||||
output = mlflow_repository.transform(
|
||||
model_name, data, {}, metadata['metadata'])
|
||||
|
||||
mlflow_repository.model_serving.get_cached_transform.assert_called_once_with(
|
||||
model_name, data, 0, 'sklearn', False, 'model', 'predict')
|
||||
|
||||
assert output == {
|
||||
'success': False,
|
||||
'content': {
|
||||
'message': 'error',
|
||||
'traceback': ANY
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def test_predict_success(mlflow_repository):
|
||||
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]
|
||||
)
|
||||
|
||||
output = mlflow_repository.predict(
|
||||
model_name, data, {}, metadata['metadata'])
|
||||
|
||||
mlflow_repository.model_serving.get_cached_predict.assert_called_once_with(
|
||||
model_name, data, 0, 'pyfunc', False, 'model')
|
||||
|
||||
assert output['success'] is True
|
||||
assert output['content'] == {
|
||||
'prediction': {
|
||||
'index_1': 2,
|
||||
'index_2': 3
|
||||
}, 'response_time': {
|
||||
'index_1': ANY,
|
||||
'index_2': ANY
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def test_predict_error(mlflow_repository):
|
||||
data = DataFrame({
|
||||
'feat_1': {
|
||||
'index_1': 2,
|
||||
'index_2': 3
|
||||
}
|
||||
})
|
||||
model_name = 'model'
|
||||
|
||||
mlflow_repository.model_serving.get_cached_predict = MagicMock(
|
||||
side_effect=Exception('error')
|
||||
)
|
||||
|
||||
output = mlflow_repository.predict(
|
||||
model_name, data, {}, metadata['metadata'])
|
||||
|
||||
mlflow_repository.model_serving.get_cached_predict.assert_called_once_with(
|
||||
model_name, data, 0, 'pyfunc', False, 'model')
|
||||
|
||||
assert output == {
|
||||
'success': False,
|
||||
'content': {
|
||||
'message': 'error',
|
||||
'traceback': ANY
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@patch('laborious.utils.repository.model_repository.mlflow')
|
||||
def test_get_experiment_by_run_id(mlflow, mlflow_repository):
|
||||
mlflow.get_run.return_value = MagicMock(
|
||||
info=MagicMock(
|
||||
experiment_id='0',
|
||||
)
|
||||
)
|
||||
mlflow.get_experiment.return_value = MagicMock()
|
||||
mlflow.get_experiment.return_value.name = 'test'
|
||||
|
||||
output = mlflow_repository.get_experiment_by_run_id('0')
|
||||
assert output == 'test'
|
||||
mlflow.get_run.assert_called_once_with('0')
|
||||
mlflow.get_experiment.assert_called_once_with('0')
|
||||
|
||||
|
||||
@patch('laborious.utils.repository.model_repository.mlflow')
|
||||
def test_get_next_run_name(mlflow, mlflow_repository):
|
||||
mlflow.search_runs.return_value = [1, 2, 3]
|
||||
output = mlflow_repository.get_next_run_name('run')
|
||||
assert output == 'run-4'
|
||||
mlflow.search_runs.assert_called_once_with(
|
||||
experiment_names=['run'],
|
||||
order_by=['start_time desc'],
|
||||
)
|
||||
|
||||
|
||||
@patch('laborious.utils.repository.model_repository.mlflow')
|
||||
def test_get_experiment_success(mlflow, mlflow_repository):
|
||||
mlflow.get_experiment_by_name.return_value = MagicMock(
|
||||
experiment_id='0')
|
||||
|
||||
output = mlflow_repository.get_experiment('test')
|
||||
|
||||
assert output == 0
|
||||
|
||||
|
||||
@patch('laborious.utils.repository.model_repository.mlflow')
|
||||
def test_get_experiment_error(mlflow, mlflow_repository):
|
||||
mlflow.get_experiment_by_name.return_value = None
|
||||
|
||||
try:
|
||||
mlflow_repository.get_experiment('test')
|
||||
except ValueError as e:
|
||||
assert str(e) == 'Experiment test not found'
|
||||
else:
|
||||
assert False
|
||||
|
||||
|
||||
@patch('laborious.utils.repository.model_repository.mlflow')
|
||||
def test_get_experiment_last_run(mlflow, mlflow_repository):
|
||||
mlflow.search_runs.return_value = DataFrame({
|
||||
'params.retrain': ['True', 'False', 'True', 'False'],
|
||||
'end_time': ['2021-01-01', '2021-01-02', '2021-01-03', '2021-01-04'],
|
||||
'run_id': ['0', '1', '2', '3'],
|
||||
})
|
||||
|
||||
output = mlflow_repository.get_experiment_last_run(0)
|
||||
|
||||
mlflow.search_runs.assert_called_once_with(
|
||||
experiment_ids=[0],
|
||||
filter_string="",
|
||||
output_format="pandas",
|
||||
)
|
||||
|
||||
assert output == '2'
|
||||
|
||||
|
||||
@patch('laborious.utils.repository.model_repository.mlflow')
|
||||
def test_get_experiment_last_run_error(mlflow, mlflow_repository):
|
||||
mlflow.search_runs.return_value = []
|
||||
|
||||
try:
|
||||
mlflow_repository.get_experiment_last_run(0)
|
||||
except ValueError as e:
|
||||
assert str(e) == 'Runs is not a pandas DataFrame'
|
||||
else:
|
||||
assert False
|
||||
|
||||
|
||||
@patch('laborious.utils.repository.model_repository.mlflow.sklearn')
|
||||
@patch('laborious.utils.repository.model_repository.mlflow.set_experiment')
|
||||
def test_create_model_experiment(set_experiment, sklearn, mlflow_repository):
|
||||
|
||||
mlflow_repository.model_serving.get_model_run_id = MagicMock(
|
||||
return_value='0')
|
||||
mlflow_repository.model_serving.get_model_uri = MagicMock(
|
||||
return_value='test')
|
||||
mlflow_repository.get_experiment_by_run_id = MagicMock()
|
||||
|
||||
data_model_mock = MagicMock()
|
||||
prediction_model_mock = MagicMock()
|
||||
|
||||
sklearn.load_model.side_effect = [data_model_mock, prediction_model_mock]
|
||||
|
||||
data_model_mock.fit.return_value = data_model_mock
|
||||
data_model_mock.predict.return_value = DataFrame({
|
||||
'x': [10, 20, 30],
|
||||
})
|
||||
data_model_mock.target_variable = 'y'
|
||||
|
||||
prediction_model_mock.fit.return_value = prediction_model_mock
|
||||
|
||||
data = DataFrame({
|
||||
'x': [1, 2, 3],
|
||||
'y': [4, 5, 6]
|
||||
})
|
||||
|
||||
output = mlflow_repository.create_model_experiment('test', data)
|
||||
|
||||
mlflow_repository.model_serving.get_model_run_id.assert_called_once_with(
|
||||
'test', stage='Production')
|
||||
mlflow_repository.model_serving.get_model_uri.assert_called_once_with(
|
||||
'0', prediction=False)
|
||||
|
||||
sklearn.load_model.assert_has_calls([
|
||||
call(mlflow_repository.model_serving.get_model_uri.return_value),
|
||||
call("models:/test/production"),
|
||||
])
|
||||
assert sklearn.load_model.call_count == 2
|
||||
|
||||
data_model_mock.fit.assert_called_once_with(data)
|
||||
data_model_mock.predict.assert_called_once_with(data)
|
||||
|
||||
fit_args = prediction_model_mock.fit.call_args[0][0]
|
||||
assert fit_args.equals(
|
||||
DataFrame({
|
||||
'x': [10, 20, 30],
|
||||
'y': [4, 5, 6],
|
||||
})
|
||||
)
|
||||
|
||||
mlflow_repository.get_experiment_by_run_id.assert_called_once_with('0')
|
||||
|
||||
set_experiment.assert_called_once_with(
|
||||
mlflow_repository.get_experiment_by_run_id.return_value
|
||||
)
|
||||
|
||||
assert output == (prediction_model_mock,
|
||||
data_model_mock,
|
||||
mlflow_repository.get_experiment_by_run_id.return_value)
|
||||
|
||||
|
||||
@patch('laborious.utils.repository.model_repository.mlflow.start_run')
|
||||
@patch('laborious.utils.repository.model_repository.mlflow.log_param')
|
||||
@patch('laborious.utils.repository.model_repository.mlflow.sklearn.log_model')
|
||||
@patch('laborious.utils.repository.model_repository.mlflow.log_artifact')
|
||||
def test_perform_model_retrain(log_artifact, log_model, log_param, start_run, mlflow_repository):
|
||||
|
||||
prediction_model_mock = MagicMock()
|
||||
data_model_mock = MagicMock()
|
||||
experiment = 'test'
|
||||
model_name = 'test'
|
||||
data = MagicMock()
|
||||
|
||||
mlflow_repository.get_next_run_name = MagicMock(
|
||||
return_value='test-1')
|
||||
run = MagicMock()
|
||||
start_run.__enter__.return_value = run
|
||||
|
||||
output = mlflow_repository.perform_model_retrain(
|
||||
prediction_model_mock, data_model_mock, experiment, model_name, data)
|
||||
|
||||
mlflow_repository.get_next_run_name.assert_called_once_with(experiment)
|
||||
start_run.assert_called_once_with(
|
||||
run_name='test-1', description='Retrain model test with new data')
|
||||
|
||||
log_model.assert_has_calls([
|
||||
call(data_model_mock, "data_model"),
|
||||
call(prediction_model_mock, "prediction_model"),
|
||||
])
|
||||
|
||||
data.to_csv.assert_called_once_with(
|
||||
"temp/raw_data_test.csv", index=True)
|
||||
|
||||
log_artifact.assert_called_once_with(
|
||||
"temp/raw_data_test.csv")
|
||||
|
||||
log_param.assert_has_calls([
|
||||
call("retrain", True),
|
||||
])
|
||||
|
||||
assert output == ("Model retrained successfully", experiment)
|
||||
|
||||
|
||||
def test_retrain_model(mlflow_repository):
|
||||
data = MagicMock()
|
||||
model_name = 'test'
|
||||
|
||||
mlflow_repository.create_model_experiment = MagicMock(
|
||||
return_value=('data_model', 'prediction_model', '0'))
|
||||
|
||||
mlflow_repository.perform_model_retrain = MagicMock(
|
||||
return_value='Model retrained successfully')
|
||||
|
||||
output = mlflow_repository.retrain_model(data, model_name)
|
||||
|
||||
mlflow_repository.create_model_experiment.assert_called_once_with(
|
||||
model_name, data)
|
||||
|
||||
mlflow_repository.perform_model_retrain.assert_called_once_with(
|
||||
'data_model', 'prediction_model', '0', model_name, data)
|
||||
|
||||
assert output == 'Model retrained successfully'
|
||||
|
||||
|
||||
@patch('laborious.utils.repository.model_repository.mlflow')
|
||||
def test_update_production_model_by_run_id(mlflow, mlflow_repository):
|
||||
client_mock = MagicMock()
|
||||
mlflow.tracking.MlflowClient.return_value = client_mock
|
||||
|
||||
client_mock.get_registered_model.return_value = MagicMock(
|
||||
latest_versions=[
|
||||
MagicMock(version='1'),
|
||||
MagicMock(version='2'),
|
||||
MagicMock(version='3'),
|
||||
]
|
||||
)
|
||||
output = mlflow_repository.update_production_model_by_run_id('0', 'test')
|
||||
|
||||
mlflow.register_model.assert_called_once_with(
|
||||
"runs:/0/prediction_model",
|
||||
'test',
|
||||
)
|
||||
|
||||
mlflow.tracking.MlflowClient.assert_called_once()
|
||||
client_mock.get_registered_model.assert_called_once_with('test')
|
||||
client_mock.transition_model_version_stage.assert_called_once_with(
|
||||
name='test',
|
||||
version='3',
|
||||
stage='Production',
|
||||
archive_existing_versions=True,
|
||||
)
|
||||
|
||||
assert output == {
|
||||
'model_name': 'test',
|
||||
'version': '3',
|
||||
'mlflow_run_id': '0',
|
||||
}
|
||||
|
||||
|
||||
@patch('laborious.utils.repository.model_repository.mlflow')
|
||||
def test_update_production_model_by_run_id_error(mlflow, mlflow_repository):
|
||||
mlflow.tracking.MlflowClient.return_value = MagicMock(
|
||||
get_registered_model=MagicMock(
|
||||
return_value=MagicMock(
|
||||
latest_versions={}
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
try:
|
||||
mlflow_repository.update_production_model_by_run_id('0', 'test')
|
||||
except Exception as e:
|
||||
assert str(e) == 'Model versions is not a list'
|
||||
else:
|
||||
assert False
|
||||
|
||||
|
||||
def test_update_production_model(mlflow_repository):
|
||||
connector = mlflow_repository
|
||||
|
||||
with patch.object(connector, 'get_experiment',
|
||||
return_value='0') as get_experiment:
|
||||
with patch.object(connector, 'get_experiment_last_run',
|
||||
return_value='2') as get_experiment_last_run:
|
||||
with patch.object(connector, 'update_production_model_by_run_id',
|
||||
return_value={'model_name': 'test', 'version': '3',
|
||||
'mlflow_run_id': '0'}) as update_production_model_by_run_id:
|
||||
|
||||
output = connector.update_production_model('0', 'test')
|
||||
|
||||
get_experiment.assert_called_once_with('0')
|
||||
get_experiment_last_run.assert_called_once_with('0')
|
||||
update_production_model_by_run_id.assert_called_once_with(
|
||||
'2', 'test')
|
||||
|
||||
assert output == {
|
||||
'model_name': 'test',
|
||||
'version': '3',
|
||||
'mlflow_run_id': '0',
|
||||
'mlflow_experiment_id': '0',
|
||||
}
|
||||
388
tests/laborious/utils/repository/test_opc_repository.py
Normal file
388
tests/laborious/utils/repository/test_opc_repository.py
Normal file
@@ -0,0 +1,388 @@
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, Mock, patch, MagicMock, ANY, call
|
||||
from asyncua.crypto.security_policies import SecurityPolicyBasic256
|
||||
from laborious.utils.repository.opc_repository import OpcRepository
|
||||
from sientia_do.notifications.models import NotificationLevel
|
||||
from datetime import datetime
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_logger():
|
||||
return Mock()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def opc_repository(mock_logger):
|
||||
return OpcRepository(
|
||||
id="test_repo",
|
||||
url="opc.tcp://localhost:4840",
|
||||
logger=mock_logger,
|
||||
notification_handler=Mock(),
|
||||
reconnection_interval=60,
|
||||
server_uri="urn:test:server",
|
||||
cert_path="/path/to/cert.pem",
|
||||
private_key_path="/path/to/key.pem",
|
||||
server_cert_path="/path/to/server_cert.pem"
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_client():
|
||||
with patch('laborious.utils.repository.opc_repository.Client') as mock:
|
||||
client_instance = AsyncMock()
|
||||
mock.return_value = client_instance
|
||||
yield client_instance
|
||||
|
||||
|
||||
metadata = {
|
||||
"metadata": {
|
||||
"model_id": "test_model",
|
||||
"model_name": "test_model",
|
||||
"workflow_name": "test_workflow",
|
||||
"schema_name": "test_schedule",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def test_init(opc_repository):
|
||||
assert opc_repository.id == "test_repo"
|
||||
assert opc_repository.url == "opc.tcp://localhost:4840"
|
||||
assert opc_repository.server_uri == "urn:test:server"
|
||||
assert opc_repository.cert_path == "/path/to/cert.pem"
|
||||
assert opc_repository.private_key_path == "/path/to/key.pem"
|
||||
assert opc_repository.server_cert_path == "/path/to/server_cert.pem"
|
||||
assert opc_repository.reconnection_interval == 60
|
||||
assert opc_repository.client is None
|
||||
assert opc_repository.last_reconnection_time is None
|
||||
assert opc_repository.error_count == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_security(opc_repository, mock_client):
|
||||
opc_repository.client = mock_client
|
||||
await opc_repository.set_security()
|
||||
|
||||
mock_client.application_uri = "urn:test:server"
|
||||
mock_client.set_security.assert_called_once_with(
|
||||
SecurityPolicyBasic256,
|
||||
certificate="/path/to/cert.pem",
|
||||
private_key="/path/to/key.pem",
|
||||
server_certificate="/path/to/server_cert.pem"
|
||||
)
|
||||
assert mock_client.secure_channel_timeout == 10000000
|
||||
assert mock_client.session_timeout == 10000000
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_security_missing_certificates(opc_repository):
|
||||
opc_repository.cert_path = None
|
||||
opc_repository.private_key_path = None
|
||||
|
||||
try:
|
||||
await opc_repository.set_security()
|
||||
except ValueError as e:
|
||||
assert str(
|
||||
e) == "Certificate and private key paths must be provided for secure connection."
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_connect_with_security(opc_repository, mock_client):
|
||||
opc_repository.try_connect = AsyncMock(return_value=(True, {}))
|
||||
result = await opc_repository.connect()
|
||||
|
||||
opc_repository.try_connect.assert_called_once()
|
||||
assert opc_repository.client == mock_client
|
||||
assert result == (True, {})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_connect_without_security(opc_repository, mock_client):
|
||||
opc_repository.cert_path = None
|
||||
opc_repository.try_connect = AsyncMock(return_value=(True, {}))
|
||||
opc_repository.set_security = AsyncMock()
|
||||
result = await opc_repository.connect()
|
||||
|
||||
opc_repository.try_connect.assert_called_once()
|
||||
opc_repository.set_security.assert_not_called()
|
||||
assert opc_repository.client == mock_client
|
||||
assert result == (True, {})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_try_connect_success(opc_repository):
|
||||
opc_repository.last_reconnection_time = None
|
||||
opc_repository.client = AsyncMock()
|
||||
result = await opc_repository.try_connect()
|
||||
|
||||
opc_repository.client.connect.assert_called_once()
|
||||
assert opc_repository.last_reconnection_time is not None
|
||||
assert result == (True, {})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_try_connect_fail(opc_repository):
|
||||
opc_repository.last_reconnection_time = None
|
||||
opc_repository.client = MagicMock()
|
||||
opc_repository.client.connect.side_effect = Exception("Test error")
|
||||
|
||||
is_connected, error_data = await opc_repository.try_connect()
|
||||
|
||||
opc_repository.client.connect.assert_called_once()
|
||||
assert is_connected is False
|
||||
assert error_data['notification_id'] == f"OPC_CONNECTION_ERROR_{opc_repository.id}"
|
||||
assert error_data['message'] == "Failed to connect to OPC server: Test error"
|
||||
assert error_data['block'] == "opc_repository"
|
||||
assert error_data['level'] == NotificationLevel.ERROR
|
||||
assert error_data['attachment_content'] is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_disconnect(opc_repository, mock_client):
|
||||
opc_repository.client = mock_client
|
||||
await opc_repository.disconnect()
|
||||
|
||||
mock_client.disconnect.assert_called_once()
|
||||
assert opc_repository.client is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_disconnect_no_client(opc_repository):
|
||||
opc_repository.client = None
|
||||
assert await opc_repository.disconnect() is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_disconnect_error(opc_repository, mock_client):
|
||||
opc_repository.client = mock_client
|
||||
mock_client.disconnect.side_effect = Exception("Test error")
|
||||
await opc_repository.disconnect()
|
||||
|
||||
opc_repository.logger.custom_error.assert_called_once_with(
|
||||
"Failed to disconnect from OPC server: Test error",
|
||||
ANY
|
||||
)
|
||||
assert opc_repository.client is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_validate_connection_none_client(opc_repository):
|
||||
opc_repository.client = None
|
||||
opc_repository.connect = AsyncMock(return_value=(True, {}))
|
||||
response = await opc_repository.validate_connection()
|
||||
assert response == (True, {})
|
||||
opc_repository.connect.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_validate_connection_error_count_disconnect_error(opc_repository):
|
||||
opc_repository.error_count = 6
|
||||
opc_repository.client = AsyncMock()
|
||||
opc_repository.disconnect = AsyncMock(
|
||||
side_effect=Exception("Test error")
|
||||
)
|
||||
opc_repository.connect = AsyncMock(return_value=(True, {}))
|
||||
|
||||
response = await opc_repository.validate_connection()
|
||||
assert response == opc_repository.connect.return_value
|
||||
opc_repository.disconnect.assert_called_once()
|
||||
opc_repository.connect.assert_called_once()
|
||||
opc_repository.logger.custom_error.assert_has_calls(
|
||||
[
|
||||
call("Failed to disconnect from OPC server: Test error", ANY),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_validate_connection_error_validate_connection_error(opc_repository):
|
||||
opc_repository.client = MagicMock(
|
||||
uaclient=Exception("Test error")
|
||||
)
|
||||
opc_repository.error_count = 0
|
||||
|
||||
response = await opc_repository.validate_connection()
|
||||
|
||||
assert response == (False, {
|
||||
"notification_id": f"OPC_CONNECTION_CHECK_ERROR_{opc_repository.id}",
|
||||
"message": "Failed to validate connection to OPC server: 'Exception' object has no attribute 'protocol'",
|
||||
"block": "opc_repository",
|
||||
"level": NotificationLevel.ERROR,
|
||||
"attachment_content": ANY
|
||||
})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch('laborious.utils.repository.opc_repository.datetime')
|
||||
async def test_validate_connection_lost_not_time_to_reconnect(_mock_datetime, opc_repository):
|
||||
_mock_datetime.now = MagicMock(
|
||||
return_value=datetime(2025, 1, 1, 0, 0, 0))
|
||||
opc_repository.error_count = 0
|
||||
opc_repository.client = MagicMock()
|
||||
opc_repository.client.uaclient.protocol = None
|
||||
opc_repository.last_reconnection_time = datetime(2025, 1, 1, 0, 0, 0)
|
||||
opc_repository.connect = MagicMock(return_value=(True, {}))
|
||||
|
||||
response = await opc_repository.validate_connection()
|
||||
opc_repository.connect.assert_not_called()
|
||||
assert response == (False, {
|
||||
"notification_id": f"OPC_CONNECTION_AWAITING_RECONNECTION_WINDOW_{opc_repository.id}",
|
||||
"message": f"OPC server {opc_repository.id} is not connected, waiting for next reconnection window...",
|
||||
"block": "opc_repository",
|
||||
"level": NotificationLevel.WARNING
|
||||
})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch('laborious.utils.repository.opc_repository.datetime')
|
||||
async def test_validate_connection_lost_time_to_reconnect(mock_datetime, opc_repository):
|
||||
mock_datetime.now = MagicMock(
|
||||
return_value=datetime(2025, 1, 1, 1, 0, 0))
|
||||
opc_repository.error_count = 0
|
||||
opc_repository.client = AsyncMock()
|
||||
opc_repository.client.uaclient.protocol = None
|
||||
opc_repository.last_reconnection_time = datetime(2025, 1, 1, 0, 0, 0)
|
||||
opc_repository.connect = AsyncMock(return_value=(True, {}))
|
||||
|
||||
response = await opc_repository.validate_connection()
|
||||
opc_repository.connect.assert_called_once()
|
||||
assert response == opc_repository.connect.return_value
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_validate_connection_success(opc_repository):
|
||||
opc_repository.client = MagicMock()
|
||||
opc_repository.error_count = 0
|
||||
opc_repository.client.uaclient.protocol = MagicMock()
|
||||
opc_repository.client.uaclient.protocol.state = "open"
|
||||
|
||||
output = await opc_repository.validate_connection()
|
||||
assert output == (True, {})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_data_validate_connection_do_nothing(opc_repository):
|
||||
opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
|
||||
opc_repository.client = AsyncMock(
|
||||
get_node=MagicMock()
|
||||
)
|
||||
mock_node = AsyncMock()
|
||||
opc_repository.client.get_node.return_value = mock_node
|
||||
|
||||
result = await opc_repository.write_data("ns=2;s=TestNode", 42.0,
|
||||
"float", opc_repository.logger, metadata['metadata'])
|
||||
|
||||
opc_repository.validate_connection.assert_called_once()
|
||||
opc_repository.client.get_node.assert_called_once_with("ns=2;s=TestNode")
|
||||
assert result == (True, {})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_data_validate_connection_failed(opc_repository):
|
||||
opc_repository.validate_connection = AsyncMock(return_value=(False, {}))
|
||||
opc_repository.client = AsyncMock()
|
||||
opc_repository.error_count = 0
|
||||
|
||||
result = await opc_repository.write_data("ns=2;s=TestNode", 42.0,
|
||||
"float", opc_repository.logger, metadata['metadata'])
|
||||
|
||||
opc_repository.validate_connection.assert_called_once()
|
||||
opc_repository.client.get_node.assert_not_called()
|
||||
assert result == (False, {})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_data_get_node_failed(opc_repository):
|
||||
opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
|
||||
opc_repository.client = AsyncMock()
|
||||
opc_repository.error_count = 0
|
||||
opc_repository.client.get_node = MagicMock(
|
||||
side_effect=Exception("Test error"))
|
||||
|
||||
is_success, error_data = await opc_repository.write_data("ns=2;s=TestNode", 42.0,
|
||||
"float", opc_repository.logger, metadata['metadata'])
|
||||
|
||||
opc_repository.validate_connection.assert_called_once()
|
||||
opc_repository.client.get_node.assert_called_once_with("ns=2;s=TestNode")
|
||||
assert is_success is False
|
||||
assert error_data['notification_id'] == f"OPC_WRITE_GET_NODE_ERROR_{opc_repository.id}"
|
||||
assert error_data['message'] == "Failed to get node from OPC server: Test error | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}"
|
||||
assert error_data['block'] == "opc_repository"
|
||||
assert error_data['level'] == NotificationLevel.ERROR
|
||||
assert error_data['attachment_content'] is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_data_invalid_data_type(opc_repository, mock_client):
|
||||
opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
|
||||
opc_repository.client = mock_client
|
||||
mock_node = AsyncMock()
|
||||
mock_client.get_node = MagicMock(return_value=mock_node)
|
||||
|
||||
is_success, error_data = await opc_repository.write_data("ns=2;s=TestNode", 42.0,
|
||||
"invalid_type", opc_repository.logger, metadata['metadata'])
|
||||
|
||||
opc_repository.validate_connection.assert_called_once()
|
||||
mock_client.get_node.assert_called_once_with("ns=2;s=TestNode")
|
||||
|
||||
assert is_success is False
|
||||
assert error_data['notification_id'] == f"OPC_WRITE_DATA_TYPE_ERROR_{opc_repository.id}"
|
||||
assert error_data['message'] == "Unsupported data type: invalid_type | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}"
|
||||
assert error_data['block'] == "opc_repository"
|
||||
assert error_data['level'] == NotificationLevel.ERROR
|
||||
assert error_data.get('attachment_content') is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch('laborious.utils.repository.opc_repository.metrics')
|
||||
async def test_write_data(mock_metrics, opc_repository, mock_client):
|
||||
opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
|
||||
opc_repository.client = mock_client
|
||||
mock_node = AsyncMock()
|
||||
mock_client.get_node = MagicMock(return_value=mock_node)
|
||||
|
||||
result = await opc_repository.write_data("ns=2;s=TestNode", 42.0,
|
||||
"float", opc_repository.logger, metadata['metadata'])
|
||||
|
||||
mock_client.get_node.assert_called_once_with("ns=2;s=TestNode")
|
||||
mock_node.write_value.assert_called_once()
|
||||
assert result == (True, {})
|
||||
|
||||
mock_metrics.PREDICTION_OPC_WRITING_COUNT.labels.assert_called_once_with(
|
||||
pod_id=opc_repository.pod_id,
|
||||
model_name=metadata['metadata']['model_name'],
|
||||
pipeline_name=metadata['metadata']['workflow_name'],
|
||||
opc_server_id=opc_repository.id
|
||||
)
|
||||
mock_metrics.PREDICTION_OPC_WRITING_COUNT.labels.return_value.inc.assert_called_once_with()
|
||||
|
||||
mock_metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR.labels.assert_called_once_with(
|
||||
pod_id=opc_repository.pod_id,
|
||||
model_name=metadata['metadata']['model_name'],
|
||||
pipeline_name=metadata['metadata']['workflow_name'],
|
||||
opc_server_id=opc_repository.id
|
||||
)
|
||||
mock_metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR.labels.return_value.observe.assert_called_once_with(
|
||||
ANY)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_data_write_value_failed(opc_repository, mock_client):
|
||||
opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
|
||||
opc_repository.client = mock_client
|
||||
mock_node = AsyncMock()
|
||||
opc_repository.error_count = 0
|
||||
mock_client.get_node = MagicMock(return_value=mock_node)
|
||||
mock_node.write_value.side_effect = Exception("Test error")
|
||||
|
||||
is_success, error_data = await opc_repository.write_data("ns=2;s=TestNode", 42.0,
|
||||
"float", opc_repository.logger, metadata['metadata'])
|
||||
|
||||
opc_repository.validate_connection.assert_called_once()
|
||||
mock_client.get_node.assert_called_once_with("ns=2;s=TestNode")
|
||||
mock_node.write_value.assert_called_once()
|
||||
assert is_success is False
|
||||
assert error_data['notification_id'] == f"OPC_WRITE_DATA_ERROR_{opc_repository.id}"
|
||||
assert error_data['message'] == "Failed to write data to OPC server: Test error | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}"
|
||||
assert error_data['block'] == "opc_repository"
|
||||
assert error_data['level'] == NotificationLevel.ERROR
|
||||
assert error_data['attachment_content'] is not None
|
||||
Reference in New Issue
Block a user