SIENTIAPDE-1646 Sync full repo from local release/SIENTIAPDE-1646 (272e02d)
This commit is contained in:
vitor-aignosi
@@ -1,23 +1,36 @@
|
||||
from pandas import DataFrame
|
||||
|
||||
from laborious.utils.filters.conditional_filters import (
|
||||
filter_empty_data,
|
||||
filter_specific_variables_null_values,
|
||||
filter_empty_data
|
||||
)
|
||||
|
||||
|
||||
def test_filter_specific_variables_null_values():
|
||||
assert filter_specific_variables_null_values(
|
||||
DataFrame(
|
||||
{'variable': ['variable1', 'variable2'], 'value': [1, 2]}),
|
||||
config={'variables': ['variable2']}) is False
|
||||
assert (
|
||||
filter_specific_variables_null_values(
|
||||
DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, 2]}),
|
||||
config={'variables': ['variable2']},
|
||||
)
|
||||
is False
|
||||
)
|
||||
|
||||
|
||||
def test_filter_specific_variables_null_values_with_empty_data():
|
||||
assert (
|
||||
filter_specific_variables_null_values(DataFrame(), config={'variables': ['variable2']})
|
||||
is False
|
||||
)
|
||||
|
||||
|
||||
def test_filter_specific_variables_null_values_with_null_values():
|
||||
assert filter_specific_variables_null_values(
|
||||
DataFrame(
|
||||
{'variable': ['variable1', 'variable2'], 'value': [1, None]}),
|
||||
config={'variables': ['variable2']}) is True
|
||||
assert (
|
||||
filter_specific_variables_null_values(
|
||||
DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, None]}),
|
||||
config={'variables': ['variable2']},
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
def test_filter_empty_data():
|
||||
@@ -25,6 +38,7 @@ def test_filter_empty_data():
|
||||
|
||||
|
||||
def test_filter_empty_data_with_data():
|
||||
assert filter_empty_data(
|
||||
DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, 2]}),
|
||||
{}) is False
|
||||
assert (
|
||||
filter_empty_data(DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, 2]}), {})
|
||||
is False
|
||||
)
|
||||
|
||||
@@ -1,22 +1,23 @@
|
||||
from pandas import DataFrame
|
||||
|
||||
from laborious.utils.filters.mlflow_filters import api_error_filter, nan_values_filter
|
||||
|
||||
|
||||
def test_api_error_filter_invalid_response():
|
||||
assert api_error_filter(None, {}) == True # NOSONAR
|
||||
assert api_error_filter(None, {}) is True # NOSONAR
|
||||
|
||||
|
||||
def test_api_error_filter_valid_response_fail():
|
||||
assert api_error_filter({'success': False}, {}) == True
|
||||
assert api_error_filter({'success': False}, {}) is True
|
||||
|
||||
|
||||
def test_api_error_filter_valid_response_success():
|
||||
assert api_error_filter({'success': True}, {}) == False
|
||||
assert api_error_filter({'success': True}, {}) is False
|
||||
|
||||
|
||||
def test_nan_values_filter_all_nan_values():
|
||||
assert nan_values_filter(DataFrame({'variable': [None, None]}), {}) == True
|
||||
assert nan_values_filter(DataFrame({'variable': [None, None]}), {}) is True
|
||||
|
||||
|
||||
def test_nan_values_filter_no_nan_values():
|
||||
assert nan_values_filter(DataFrame({'variable': [1, 2]}), {}) == False
|
||||
assert nan_values_filter(DataFrame({'variable': [1, 2]}), {}) is False
|
||||
|
||||
278
tests/laborious/utils/models/test_minio_dataframe_payload.py
Normal file
278
tests/laborious/utils/models/test_minio_dataframe_payload.py
Normal file
@@ -0,0 +1,278 @@
|
||||
from datetime import datetime
|
||||
from io import BytesIO
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from pandas import DataFrame
|
||||
|
||||
from laborious.utils.models.minio_dataframe_payload import (
|
||||
MinioDataFramePayload,
|
||||
_build_object_key,
|
||||
)
|
||||
|
||||
|
||||
def test_parse_object_timestamp_hyphenated_model():
|
||||
key = 'predictions/sched/my-long-model-initial-2024-06-15_10-30-45.parquet'
|
||||
ts = MinioDataFramePayload.parse_object_timestamp(key)
|
||||
assert ts == datetime(2024, 6, 15, 10, 30, 45)
|
||||
|
||||
|
||||
def test_parse_object_timestamp_transform():
|
||||
key = 'p/m-transform-2024-01-02_03-04-05.parquet'
|
||||
ts = MinioDataFramePayload.parse_object_timestamp(key)
|
||||
assert ts == datetime(2024, 1, 2, 3, 4, 5)
|
||||
|
||||
|
||||
def test_parse_object_timestamp_invalid():
|
||||
assert MinioDataFramePayload.parse_object_timestamp('bad.parquet') is None
|
||||
|
||||
|
||||
def test_estimate_size_bytes_returns_positive_for_nonempty_frame():
|
||||
df = DataFrame({'a': [1, 2]})
|
||||
size = MinioDataFramePayload.estimate_size_bytes(df)
|
||||
assert isinstance(size, int)
|
||||
assert size > 0
|
||||
|
||||
|
||||
def test_cleanup_prefix_when_offloaded_returns_object_prefix():
|
||||
payload = MinioDataFramePayload(
|
||||
last_timestamp='t',
|
||||
data=None,
|
||||
object_key='training_datasets/m/m-initial-2024-01-01_00-00-00.parquet',
|
||||
object_prefix='training_datasets/m',
|
||||
)
|
||||
assert MinioDataFramePayload.cleanup_prefix(payload) == 'training_datasets/m'
|
||||
|
||||
|
||||
def test_cleanup_prefix_when_inline_returns_none():
|
||||
payload = MinioDataFramePayload(last_timestamp='t', data={'x': [1]}, object_key=None)
|
||||
assert MinioDataFramePayload.cleanup_prefix(payload) is None
|
||||
|
||||
|
||||
def test_has_data_true_when_object_key_set():
|
||||
payload = MinioDataFramePayload(last_timestamp='t', data=None, object_key='k')
|
||||
assert payload.has_data() is True
|
||||
|
||||
|
||||
def test_retrieve_inline_dict_as_dataframe():
|
||||
payload = MinioDataFramePayload(last_timestamp='t', data={'a': [1, 2]})
|
||||
minio = MagicMock()
|
||||
out = payload.retrieve(minio, {'metadata': {}})
|
||||
assert list(out.columns) == ['a']
|
||||
minio.download_file.assert_not_called()
|
||||
|
||||
|
||||
def test_retrieve_downloads_parquet_when_offloaded():
|
||||
source = DataFrame({'a': [1, 2]})
|
||||
buf = BytesIO()
|
||||
source.to_parquet(buf, engine='pyarrow', index=True)
|
||||
file_bytes = buf.getvalue()
|
||||
|
||||
payload = MinioDataFramePayload(
|
||||
last_timestamp='t',
|
||||
data=None,
|
||||
object_key='training_datasets/m/f.parquet',
|
||||
object_prefix='training_datasets/m',
|
||||
)
|
||||
minio = MagicMock()
|
||||
minio.download_file = MagicMock(return_value=file_bytes)
|
||||
|
||||
out = payload.retrieve(minio, {'metadata': {}})
|
||||
|
||||
minio.download_file.assert_called_once_with(
|
||||
object_name='training_datasets/m/f.parquet',
|
||||
metadata={'metadata': {}},
|
||||
)
|
||||
assert list(out.columns) == ['a']
|
||||
|
||||
|
||||
def test_build_object_key():
|
||||
key, prefix = _build_object_key('my-model', 'initial', '2024-01-01_00-00-00')
|
||||
assert key == 'prediction_datasets/my-model/my-model-initial-2024-01-01_00-00-00.parquet'
|
||||
assert prefix == 'prediction_datasets/my-model'
|
||||
|
||||
|
||||
def test_build_object_key_strips_slashes():
|
||||
key, prefix = _build_object_key(' /my-model/ ', 'transform', '2024-06-15_10-30-45')
|
||||
assert prefix == 'prediction_datasets/my-model'
|
||||
assert key.startswith('prediction_datasets/my-model/')
|
||||
|
||||
|
||||
def test_estimate_size_bytes_fallback():
|
||||
df = DataFrame({'a': [1, 2]})
|
||||
with patch.object(df, 'to_dict', side_effect=RuntimeError('to_dict failed')):
|
||||
size = MinioDataFramePayload.estimate_size_bytes(df)
|
||||
assert isinstance(size, int)
|
||||
assert size > 0
|
||||
|
||||
|
||||
def test_parse_object_timestamp_bad_datetime():
|
||||
key = 'p/m-initial-9999-99-99_99-99-99.parquet'
|
||||
assert MinioDataFramePayload.parse_object_timestamp(key) is None
|
||||
|
||||
|
||||
def test_retrieve_empty_when_no_data():
|
||||
payload = MinioDataFramePayload(last_timestamp='t', data=None, object_key=None)
|
||||
minio = MagicMock()
|
||||
out = payload.retrieve(minio, {})
|
||||
assert out.empty
|
||||
minio.download_file.assert_not_called()
|
||||
|
||||
|
||||
@patch('laborious.utils.models.minio_dataframe_payload.now')
|
||||
def test_from_dataframe_none(mock_now):
|
||||
mock_now.return_value = datetime(2024, 1, 1, 0, 0, 0)
|
||||
minio = MagicMock()
|
||||
result = MinioDataFramePayload.from_dataframe(
|
||||
dataframe=None,
|
||||
minio_repo=minio,
|
||||
model_name='m',
|
||||
operation='initial',
|
||||
status={'success': False, 'message': 'no data'},
|
||||
)
|
||||
assert result.data is None
|
||||
assert result.status == {'success': False, 'message': 'no data'}
|
||||
assert result.object_key is None
|
||||
|
||||
|
||||
@patch('laborious.utils.models.minio_dataframe_payload.now')
|
||||
def test_from_dataframe_empty(mock_now):
|
||||
mock_now.return_value = datetime(2024, 1, 1, 0, 0, 0)
|
||||
minio = MagicMock()
|
||||
mock_df = MagicMock()
|
||||
mock_df.__bool__ = MagicMock(return_value=True)
|
||||
mock_df.empty = True
|
||||
result = MinioDataFramePayload.from_dataframe(
|
||||
dataframe=mock_df,
|
||||
minio_repo=minio,
|
||||
model_name='m',
|
||||
operation='initial',
|
||||
)
|
||||
assert result.data is None
|
||||
assert result.object_key is None
|
||||
|
||||
|
||||
def _mock_dataframe(data_dict, timestamp_values=None):
|
||||
"""Build a MagicMock that behaves enough like a DataFrame for from_dataframe."""
|
||||
mock_df = MagicMock()
|
||||
mock_df.__bool__ = MagicMock(return_value=True)
|
||||
mock_df.empty = False
|
||||
if timestamp_values is None:
|
||||
timestamp_values = data_dict.get('timestamp', ['2024-01-01'])
|
||||
ts_col = MagicMock()
|
||||
ts_col.values.tolist.return_value = timestamp_values
|
||||
mock_df.__getitem__ = MagicMock(return_value=ts_col)
|
||||
mock_df.to_dict.return_value = data_dict
|
||||
buf = BytesIO()
|
||||
DataFrame(data_dict).to_parquet(buf, engine='pyarrow', index=True)
|
||||
mock_df.to_parquet = MagicMock(side_effect=lambda b, **kw: b.write(buf.getvalue()))
|
||||
return mock_df
|
||||
|
||||
|
||||
@patch('laborious.utils.models.minio_dataframe_payload.OFFLOAD_THRESHOLD_BYTES', 10**9)
|
||||
def test_from_dataframe_inline():
|
||||
minio = MagicMock()
|
||||
df = _mock_dataframe({'timestamp': ['2024-01-01'], 'value': [42]})
|
||||
result = MinioDataFramePayload.from_dataframe(
|
||||
dataframe=df,
|
||||
minio_repo=minio,
|
||||
model_name='m',
|
||||
operation='initial',
|
||||
)
|
||||
assert result.data is not None
|
||||
assert result.object_key is None
|
||||
assert result.last_timestamp == '2024-01-01'
|
||||
|
||||
|
||||
@patch('laborious.utils.models.minio_dataframe_payload.OFFLOAD_THRESHOLD_BYTES', 10**9)
|
||||
def test_from_dataframe_inline_uses_provided_last_timestamp():
|
||||
minio = MagicMock()
|
||||
df = _mock_dataframe({'timestamp': ['2024-01-01'], 'value': [42]})
|
||||
result = MinioDataFramePayload.from_dataframe(
|
||||
dataframe=df,
|
||||
minio_repo=minio,
|
||||
model_name='m',
|
||||
operation='initial',
|
||||
last_timestamp='2024-01-02',
|
||||
)
|
||||
assert result.last_timestamp == '2024-01-02'
|
||||
|
||||
|
||||
@patch('laborious.utils.models.minio_dataframe_payload.now')
|
||||
@patch('laborious.utils.models.minio_dataframe_payload.OFFLOAD_THRESHOLD_BYTES', 0)
|
||||
def test_from_dataframe_offloaded(mock_now):
|
||||
mock_now.return_value = datetime(2024, 1, 1, 0, 0, 0)
|
||||
minio = MagicMock()
|
||||
minio.upload_file = MagicMock(return_value={'minio_object_name': 'full/key.parquet'})
|
||||
minio.bucket = 'test-bucket'
|
||||
|
||||
df = _mock_dataframe({'timestamp': ['2024-01-01'], 'value': [42]})
|
||||
result = MinioDataFramePayload.from_dataframe(
|
||||
dataframe=df,
|
||||
minio_repo=minio,
|
||||
model_name='m',
|
||||
operation='initial',
|
||||
workflow_metadata={'wf': 'data'},
|
||||
)
|
||||
assert result.data is None
|
||||
assert result.object_key == 'full/key.parquet'
|
||||
assert result.bucket == 'test-bucket'
|
||||
assert result.uri == 's3://test-bucket/full/key.parquet'
|
||||
minio.upload_file.assert_called_once()
|
||||
|
||||
|
||||
def test_from_dict_inline():
|
||||
raw = {
|
||||
'last_timestamp': '2024-01-01T00:00:00+00:00',
|
||||
'status': None,
|
||||
'data': {'col1': {0: 'val1'}},
|
||||
'bucket': None,
|
||||
'object_key': None,
|
||||
'object_prefix': None,
|
||||
'uri': None,
|
||||
}
|
||||
payload = MinioDataFramePayload.from_dict(raw)
|
||||
assert isinstance(payload, MinioDataFramePayload)
|
||||
assert payload.last_timestamp == '2024-01-01T00:00:00+00:00'
|
||||
assert payload.data == {'col1': {0: 'val1'}}
|
||||
assert payload.object_key is None
|
||||
|
||||
|
||||
def test_from_dict_offloaded():
|
||||
raw = {
|
||||
'last_timestamp': '2024-06-15T10:30:45+00:00',
|
||||
'status': {'success': True},
|
||||
'data': None,
|
||||
'bucket': 'my-bucket',
|
||||
'object_key': 'training_datasets/model/model-initial-2024-06-15_10-30-45.parquet',
|
||||
'object_prefix': 'training_datasets/model',
|
||||
'uri': 's3://my-bucket/training_datasets/model/model-initial-2024-06-15_10-30-45.parquet',
|
||||
}
|
||||
payload = MinioDataFramePayload.from_dict(raw)
|
||||
assert isinstance(payload, MinioDataFramePayload)
|
||||
assert payload.data is None
|
||||
assert payload.bucket == 'my-bucket'
|
||||
assert payload.object_key == raw['object_key']
|
||||
assert payload.object_prefix == 'training_datasets/model'
|
||||
assert payload.uri == raw['uri']
|
||||
assert payload.status == {'success': True}
|
||||
|
||||
|
||||
def test_from_dict_minimal_keys():
|
||||
raw = {'last_timestamp': '2024-01-01'}
|
||||
payload = MinioDataFramePayload.from_dict(raw)
|
||||
assert payload.last_timestamp == '2024-01-01'
|
||||
assert payload.data is None
|
||||
assert payload.bucket is None
|
||||
assert payload.object_key is None
|
||||
|
||||
|
||||
def test_from_dict_passthrough_existing_instance():
|
||||
original = MinioDataFramePayload(last_timestamp='2024-01-01', data={'a': 1}, bucket='b')
|
||||
result = MinioDataFramePayload.from_dict(original)
|
||||
assert result is original
|
||||
|
||||
|
||||
def test_debug_with_logger_calls_custom_debug():
|
||||
logger = MagicMock()
|
||||
MinioDataFramePayload._debug(logger, 'msg', {'a': 1})
|
||||
logger.custom_debug.assert_called_once_with('msg', {'a': 1})
|
||||
@@ -1,502 +0,0 @@
|
||||
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',
|
||||
}
|
||||
@@ -1,9 +1,20 @@
|
||||
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
|
||||
import concurrent.futures
|
||||
import json
|
||||
from datetime import datetime
|
||||
from unittest.mock import ANY, MagicMock, Mock, patch
|
||||
|
||||
import pytest
|
||||
from asyncua.crypto import security_policies
|
||||
from asyncua.ua.uaerrors import BadNodeIdUnknown, BadSessionIdInvalid
|
||||
from sientia_do.notifications.models import NotificationLevel
|
||||
|
||||
from laborious.utils.repository.opc_repository import (
|
||||
OpcClientAlreadyExistsError,
|
||||
OpcClientNotInitializedError,
|
||||
OpcRepository,
|
||||
OpcSessionAlreadyConnectedError,
|
||||
is_reconnectable_opcua_bad,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
@@ -13,376 +24,561 @@ def mock_logger():
|
||||
|
||||
@pytest.fixture
|
||||
def opc_repository(mock_logger):
|
||||
return OpcRepository(
|
||||
id="test_repo",
|
||||
url="opc.tcp://localhost:4840",
|
||||
repository = OpcRepository(
|
||||
opc_id='test_repo',
|
||||
server_name='test_server',
|
||||
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"
|
||||
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',
|
||||
metrics_controller=MagicMock(),
|
||||
)
|
||||
repository.disconnection_interval = 0.1
|
||||
repository.send_notification = MagicMock()
|
||||
repository.send_notification = MagicMock()
|
||||
repository.emit_metric_sync = MagicMock()
|
||||
repository.info = MagicMock()
|
||||
repository.error = MagicMock()
|
||||
repository.warning = MagicMock()
|
||||
repository.debug = MagicMock()
|
||||
repository._session_ready.set()
|
||||
return repository
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_client():
|
||||
with patch('laborious.utils.repository.opc_repository.Client') as mock:
|
||||
client_instance = AsyncMock()
|
||||
client_instance = MagicMock()
|
||||
aio = MagicMock()
|
||||
client_instance.aio_obj = aio
|
||||
aio.uaclient = MagicMock()
|
||||
aio.uaclient.protocol = MagicMock(state='closed')
|
||||
aio.session_timeout = 600_000
|
||||
aio.secure_channel_timeout = 600_000
|
||||
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",
|
||||
'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.id == 'test_repo'
|
||||
assert opc_repository.server_name == 'test_server'
|
||||
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):
|
||||
def test_set_security(opc_repository, mock_client):
|
||||
opc_repository.client = mock_client
|
||||
await opc_repository.set_security()
|
||||
opc_repository.set_security()
|
||||
|
||||
mock_client.application_uri = "urn:test:server"
|
||||
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"
|
||||
security_policies.SecurityPolicyBasic256,
|
||||
'/path/to/cert.pem',
|
||||
'/path/to/key.pem',
|
||||
None,
|
||||
'/path/to/server_cert.pem',
|
||||
)
|
||||
assert mock_client.secure_channel_timeout == 10000000
|
||||
assert mock_client.session_timeout == 10000000
|
||||
assert mock_client.aio_obj.secure_channel_timeout == 600_000
|
||||
assert mock_client.aio_obj.session_timeout == 600_000
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_set_security_missing_certificates(opc_repository):
|
||||
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()
|
||||
opc_repository.set_security()
|
||||
except ValueError as e:
|
||||
assert str(
|
||||
e) == "Certificate and private key paths must be provided for secure connection."
|
||||
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()
|
||||
def test_set_security_missing_client(opc_repository):
|
||||
opc_repository.client = None
|
||||
try:
|
||||
opc_repository.set_security()
|
||||
except ValueError as e:
|
||||
assert str(e) == 'Client must be initialized before setting security'
|
||||
|
||||
opc_repository.try_connect.assert_called_once()
|
||||
assert opc_repository.client == mock_client
|
||||
|
||||
def test_connect_with_security(opc_repository, mock_client):
|
||||
opc_repository._create_client = MagicMock()
|
||||
opc_repository._open_session = MagicMock(return_value=(True, {}))
|
||||
result = opc_repository.connect()
|
||||
|
||||
opc_repository._create_client.assert_called_once()
|
||||
opc_repository._open_session.assert_called_once()
|
||||
assert result == (True, {})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_connect_without_security(opc_repository, mock_client):
|
||||
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._create_client = MagicMock()
|
||||
opc_repository._open_session = MagicMock(return_value=(True, {}))
|
||||
opc_repository.set_security = MagicMock()
|
||||
result = opc_repository.connect()
|
||||
|
||||
opc_repository.try_connect.assert_called_once()
|
||||
opc_repository._create_client.assert_called_once()
|
||||
opc_repository._open_session.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()
|
||||
def test_connect_raises_when_session_already_open(opc_repository, mock_client):
|
||||
opc_repository.client = mock_client
|
||||
proto = MagicMock()
|
||||
proto.state = 'open'
|
||||
mock_client.aio_obj.uaclient.protocol = proto
|
||||
|
||||
with pytest.raises(OpcSessionAlreadyConnectedError, match='disconnect'):
|
||||
opc_repository.connect()
|
||||
|
||||
|
||||
def test_create_client_raises_when_client_exists(opc_repository, mock_client):
|
||||
opc_repository.client = mock_client
|
||||
|
||||
with pytest.raises(OpcClientAlreadyExistsError, match='already exists'):
|
||||
opc_repository._create_client()
|
||||
|
||||
|
||||
def test_open_session_success(opc_repository):
|
||||
closed_proto = MagicMock()
|
||||
closed_proto.state = 'closed'
|
||||
opc_repository.client = MagicMock()
|
||||
aio = MagicMock()
|
||||
opc_repository.client.aio_obj = aio
|
||||
aio.uaclient = MagicMock(protocol=closed_proto)
|
||||
aio.session_timeout = 600_000
|
||||
aio.secure_channel_timeout = 600_000
|
||||
|
||||
open_proto = MagicMock()
|
||||
open_proto.state = 'open'
|
||||
open_proto.authentication_token = 'tok'
|
||||
|
||||
def connect_side_effect():
|
||||
aio.uaclient.protocol = open_proto
|
||||
|
||||
opc_repository.client.connect = MagicMock(side_effect=connect_side_effect)
|
||||
|
||||
result = opc_repository._open_session()
|
||||
|
||||
opc_repository.client.connect.assert_called_once()
|
||||
assert opc_repository.last_reconnection_time is not None
|
||||
assert result == (True, {})
|
||||
assert opc_repository._session_ready.is_set()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_try_connect_fail(opc_repository):
|
||||
opc_repository.last_reconnection_time = None
|
||||
def test_open_session_raises_when_already_connected(opc_repository, mock_client):
|
||||
opc_repository.client = mock_client
|
||||
proto = MagicMock()
|
||||
proto.state = 'open'
|
||||
mock_client.aio_obj.uaclient.protocol = proto
|
||||
|
||||
with pytest.raises(OpcSessionAlreadyConnectedError, match='disconnect'):
|
||||
opc_repository._open_session()
|
||||
|
||||
|
||||
def test_open_session_fail(opc_repository):
|
||||
opc_repository._disconnect_locked = MagicMock()
|
||||
opc_repository.client = MagicMock()
|
||||
opc_repository.client.connect.side_effect = Exception("Test error")
|
||||
aio = MagicMock()
|
||||
opc_repository.client.aio_obj = aio
|
||||
aio.uaclient = MagicMock(protocol=MagicMock(state='closed'))
|
||||
opc_repository.client.connect = MagicMock(side_effect=Exception('Test error'))
|
||||
|
||||
is_connected, error_data = await opc_repository.try_connect()
|
||||
is_connected, error_data = opc_repository._open_session()
|
||||
|
||||
opc_repository._disconnect_locked.assert_called_once()
|
||||
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['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):
|
||||
def test_open_session_raises_when_no_client(opc_repository):
|
||||
opc_repository.client = None
|
||||
|
||||
with pytest.raises(OpcClientNotInitializedError, match='not initialized'):
|
||||
opc_repository._open_session()
|
||||
|
||||
|
||||
def test_disconnection_fallback_success(opc_repository, mock_client):
|
||||
opc_repository.client = mock_client
|
||||
await opc_repository.disconnect()
|
||||
mock_client.disconnect.return_value = True
|
||||
result = opc_repository._disconnection_fallback()
|
||||
|
||||
mock_client.disconnect.assert_called_once()
|
||||
assert opc_repository.client is None
|
||||
assert result == []
|
||||
|
||||
|
||||
@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):
|
||||
def test_disconnection_fallback_fail(opc_repository, mock_client):
|
||||
opc_repository.client = mock_client
|
||||
mock_client.disconnect.side_effect = Exception("Test error")
|
||||
await opc_repository.disconnect()
|
||||
mock_client.disconnect.side_effect = Exception('Test error')
|
||||
result = opc_repository._disconnection_fallback()
|
||||
assert result == [
|
||||
{'attempt': 1, 'error': 'Test error', 'traceback': ANY},
|
||||
{'attempt': 2, 'error': 'Test error', 'traceback': ANY},
|
||||
{'attempt': 3, 'error': 'Test error', 'traceback': ANY},
|
||||
{'attempt': 4, 'error': 'Test error', 'traceback': ANY},
|
||||
{'attempt': 5, 'error': 'Test error', 'traceback': ANY},
|
||||
]
|
||||
assert mock_client.disconnect.call_count == 5
|
||||
|
||||
opc_repository.logger.custom_error.assert_called_once_with(
|
||||
"Failed to disconnect from OPC server: Test error",
|
||||
ANY
|
||||
|
||||
def test_disconnect(opc_repository, mock_client):
|
||||
opc_repository.client = mock_client
|
||||
opc_repository._disconnection_fallback = MagicMock(return_value=[])
|
||||
opc_repository.disconnect()
|
||||
|
||||
opc_repository._disconnection_fallback.assert_called_once()
|
||||
assert opc_repository.client is None
|
||||
assert opc_repository._allow_reconnect is False
|
||||
|
||||
|
||||
def test_disconnect_no_client(opc_repository):
|
||||
opc_repository.client = None
|
||||
assert opc_repository.disconnect() is None
|
||||
|
||||
|
||||
def test_disconnect_error(opc_repository, mock_client):
|
||||
opc_repository.client = mock_client
|
||||
opc_repository._disconnection_fallback = MagicMock(
|
||||
return_value=[{'attempt': 1, 'error': 'Test error', 'traceback': 'text'}]
|
||||
)
|
||||
opc_repository.disconnect()
|
||||
|
||||
opc_repository._disconnection_fallback.assert_called_once()
|
||||
opc_repository.send_notification.assert_called_once_with(
|
||||
metadata=opc_repository.metadata,
|
||||
notification_id=f'OPC_DISCONNECTION_ERROR_{opc_repository.id}',
|
||||
message='Failed to disconnect from OPC server in 5 attempts.',
|
||||
block='opc_repository',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=json.dumps(
|
||||
[{'attempt': 1, 'error': 'Test error', 'traceback': 'text'}], indent=4
|
||||
),
|
||||
)
|
||||
assert opc_repository.client is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_validate_connection_none_client(opc_repository):
|
||||
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()
|
||||
response = opc_repository.validate_connection()
|
||||
assert response == (False, opc_repository._not_connected_error())
|
||||
|
||||
|
||||
@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
|
||||
def test_validate_connection_session_not_open(opc_repository):
|
||||
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, {}))
|
||||
opc_repository.client.aio_obj.uaclient.protocol = None
|
||||
|
||||
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
|
||||
})
|
||||
response = opc_repository.validate_connection()
|
||||
|
||||
assert response == (False, opc_repository._not_connected_error())
|
||||
opc_repository.error.assert_called_once()
|
||||
|
||||
|
||||
@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):
|
||||
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"
|
||||
proto = MagicMock()
|
||||
proto.state = 'open'
|
||||
opc_repository.client.aio_obj.uaclient.protocol = proto
|
||||
|
||||
output = await opc_repository.validate_connection()
|
||||
output = 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()
|
||||
def test_write_data_validate_connection_do_nothing(opc_repository):
|
||||
opc_repository.validate_connection = MagicMock(return_value=(True, {}))
|
||||
opc_repository.client = MagicMock(get_node=MagicMock())
|
||||
mock_node = MagicMock()
|
||||
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'])
|
||||
result = opc_repository.write_data('ns=2;s=TestNode', 42.0, 'float', 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, {})
|
||||
opc_repository.client.get_node.assert_called_once_with('ns=2;s=TestNode')
|
||||
assert result == (True, {'response_time': ANY})
|
||||
|
||||
|
||||
@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
|
||||
def test_write_data_validate_connection_failed(opc_repository):
|
||||
opc_repository.validate_connection = MagicMock(return_value=(False, {}))
|
||||
opc_repository.client = MagicMock()
|
||||
opc_repository._start_reconnect = MagicMock()
|
||||
|
||||
result = await opc_repository.write_data("ns=2;s=TestNode", 42.0,
|
||||
"float", opc_repository.logger, metadata['metadata'])
|
||||
is_success, error_data = opc_repository.write_data(
|
||||
'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
|
||||
)
|
||||
|
||||
opc_repository.validate_connection.assert_called_once()
|
||||
opc_repository.client.get_node.assert_not_called()
|
||||
assert result == (False, {})
|
||||
opc_repository._start_reconnect.assert_called_once()
|
||||
assert opc_repository._start_reconnect.call_args.args[0] == 'ProtocolClosed'
|
||||
assert is_success is False
|
||||
assert error_data['opc_error_kind'] == 'connection_lost'
|
||||
assert error_data['opc_status'] == 'ProtocolClosed'
|
||||
|
||||
|
||||
@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"))
|
||||
def test_write_data_get_node_failed(opc_repository):
|
||||
opc_repository.validate_connection = MagicMock(return_value=(True, {}))
|
||||
opc_repository.client = MagicMock()
|
||||
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'])
|
||||
is_success, error_data = opc_repository.write_data(
|
||||
'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
|
||||
)
|
||||
|
||||
opc_repository.validate_connection.assert_called_once()
|
||||
opc_repository.client.get_node.assert_called_once_with("ns=2;s=TestNode")
|
||||
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['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, {}))
|
||||
def test_write_data_invalid_data_type(opc_repository, mock_client):
|
||||
opc_repository.validate_connection = MagicMock(return_value=(True, {}))
|
||||
opc_repository.client = mock_client
|
||||
mock_node = AsyncMock()
|
||||
mock_node = MagicMock()
|
||||
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'])
|
||||
is_success, error_data = opc_repository.write_data(
|
||||
'ns=2;s=TestNode', 42.0, 'invalid_type', metadata['metadata']
|
||||
)
|
||||
|
||||
opc_repository.validate_connection.assert_called_once()
|
||||
mock_client.get_node.assert_called_once_with("ns=2;s=TestNode")
|
||||
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['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, {}))
|
||||
def test_write_data(opc_repository, mock_client):
|
||||
opc_repository.validate_connection = MagicMock(return_value=(True, {}))
|
||||
opc_repository.client = mock_client
|
||||
mock_node = AsyncMock()
|
||||
mock_node = MagicMock()
|
||||
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'])
|
||||
result = opc_repository.write_data('ns=2;s=TestNode', 42.0, 'float', metadata['metadata'])
|
||||
|
||||
mock_client.get_node.assert_called_once_with("ns=2;s=TestNode")
|
||||
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)
|
||||
assert result == (True, {'response_time': ANY})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_data_write_value_failed(opc_repository, mock_client):
|
||||
opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
|
||||
def test_write_data_write_value_failed(opc_repository, mock_client):
|
||||
opc_repository.validate_connection = MagicMock(return_value=(True, {}))
|
||||
opc_repository.client = mock_client
|
||||
mock_node = AsyncMock()
|
||||
opc_repository.error_count = 0
|
||||
mock_node = MagicMock()
|
||||
mock_client.get_node = MagicMock(return_value=mock_node)
|
||||
mock_node.write_value.side_effect = Exception("Test error")
|
||||
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'])
|
||||
is_success, error_data = opc_repository.write_data(
|
||||
'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
|
||||
)
|
||||
|
||||
opc_repository.validate_connection.assert_called_once()
|
||||
mock_client.get_node.assert_called_once_with("ns=2;s=TestNode")
|
||||
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['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
|
||||
|
||||
|
||||
def test_is_reconnectable_opcua_bad():
|
||||
assert is_reconnectable_opcua_bad(BadSessionIdInvalid()) is True
|
||||
assert is_reconnectable_opcua_bad(BadNodeIdUnknown()) is False
|
||||
assert is_reconnectable_opcua_bad(Exception('other')) is False
|
||||
|
||||
|
||||
def test_write_data_bad_session_id_invalid_schedules_reconnect(opc_repository, mock_client):
|
||||
opc_repository.validate_connection = MagicMock(return_value=(True, {}))
|
||||
opc_repository.client = mock_client
|
||||
opc_repository._start_reconnect = MagicMock()
|
||||
mock_node = MagicMock()
|
||||
mock_client.get_node = MagicMock(return_value=mock_node)
|
||||
mock_node.write_value.side_effect = BadSessionIdInvalid()
|
||||
|
||||
is_success, error_data = opc_repository.write_data(
|
||||
'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
|
||||
)
|
||||
|
||||
mock_node.write_value.assert_called_once()
|
||||
opc_repository._start_reconnect.assert_called_once()
|
||||
assert is_success is False
|
||||
assert error_data['opc_error_kind'] == 'session_bad'
|
||||
assert error_data['opc_status'] == 'BadSessionIdInvalid'
|
||||
|
||||
|
||||
def test_write_data_reconnect_in_progress_immediate(opc_repository):
|
||||
opc_repository._session_ready.clear()
|
||||
opc_repository._reconnect_thread = MagicMock()
|
||||
opc_repository._reconnect_thread.is_alive.return_value = True
|
||||
opc_repository.validate_connection = MagicMock()
|
||||
|
||||
is_success, error_data = opc_repository.write_data(
|
||||
'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
|
||||
)
|
||||
|
||||
opc_repository.validate_connection.assert_not_called()
|
||||
assert is_success is False
|
||||
assert error_data['opc_error_kind'] == 'reconnect_in_progress'
|
||||
|
||||
|
||||
def test_start_reconnect_skips_within_interval(opc_repository):
|
||||
opc_repository.last_reconnection_time = datetime.now()
|
||||
opc_repository.reconnection_interval = 3600
|
||||
|
||||
opc_repository._start_reconnect('BadSessionIdInvalid', 'tok')
|
||||
|
||||
assert opc_repository._reconnect_thread is None
|
||||
|
||||
|
||||
def test_write_data_protocol_closed_schedules_reconnect(opc_repository):
|
||||
opc_repository.client = MagicMock()
|
||||
opc_repository.client.aio_obj.uaclient.protocol = MagicMock(state='closed')
|
||||
opc_repository._start_reconnect = MagicMock()
|
||||
|
||||
is_success, error_data = opc_repository.write_data(
|
||||
'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
|
||||
)
|
||||
|
||||
opc_repository._start_reconnect.assert_called_once()
|
||||
assert opc_repository._start_reconnect.call_args.args[0] == 'ProtocolClosed'
|
||||
assert is_success is False
|
||||
assert error_data['opc_error_kind'] == 'connection_lost'
|
||||
assert error_data['opc_status'] == 'ProtocolClosed'
|
||||
|
||||
|
||||
def test_write_data_protocol_closed_skips_reconnect_within_interval(opc_repository):
|
||||
opc_repository.client = MagicMock()
|
||||
opc_repository.client.aio_obj.uaclient.protocol = MagicMock(state='closed')
|
||||
opc_repository.last_reconnection_time = datetime.now()
|
||||
opc_repository.reconnection_interval = 3600
|
||||
|
||||
is_success, error_data = opc_repository.write_data(
|
||||
'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
|
||||
)
|
||||
|
||||
assert opc_repository._reconnect_thread is None
|
||||
assert is_success is False
|
||||
assert error_data['opc_error_kind'] == 'connection_lost'
|
||||
|
||||
|
||||
def test_write_data_after_failed_reconnect_schedules_again(opc_repository):
|
||||
opc_repository._session_ready.clear()
|
||||
opc_repository.reconnection_interval = 0
|
||||
opc_repository.last_reconnection_time = None
|
||||
opc_repository._reconnect_locked = MagicMock(
|
||||
return_value=(False, {'message': 'connect failed'})
|
||||
)
|
||||
|
||||
opc_repository.write_data('ns=2;s=TestNode', 42.0, 'float', metadata['metadata'])
|
||||
if opc_repository._reconnect_thread is not None:
|
||||
opc_repository._reconnect_thread.join(timeout=2)
|
||||
assert opc_repository._reconnect_locked.call_count == 1
|
||||
assert not opc_repository._reconnect_thread_in_progress()
|
||||
|
||||
opc_repository.write_data('ns=2;s=TestNode', 42.0, 'float', metadata['metadata'])
|
||||
if opc_repository._reconnect_thread is not None:
|
||||
opc_repository._reconnect_thread.join(timeout=2)
|
||||
assert opc_repository._reconnect_locked.call_count == 2
|
||||
|
||||
|
||||
def test_write_data_after_disconnect_does_not_schedule_reconnect(opc_repository, mock_client):
|
||||
opc_repository.client = mock_client
|
||||
proto = MagicMock()
|
||||
proto.state = 'closed'
|
||||
mock_client.aio_obj.uaclient.protocol = proto
|
||||
opc_repository._disconnection_fallback = MagicMock(return_value=[])
|
||||
opc_repository.disconnect()
|
||||
|
||||
is_success, error_data = opc_repository.write_data(
|
||||
'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
|
||||
)
|
||||
|
||||
assert opc_repository._reconnect_thread is None
|
||||
assert is_success is False
|
||||
assert error_data['opc_error_kind'] == 'connection_lost'
|
||||
|
||||
|
||||
def test_parallel_bad_writes_single_reconnect_task(opc_repository, mock_client):
|
||||
opc_repository.validate_connection = MagicMock(return_value=(True, {}))
|
||||
opc_repository.client = mock_client
|
||||
opc_repository.reconnection_interval = 0
|
||||
opc_repository.last_reconnection_time = None
|
||||
mock_node = MagicMock()
|
||||
mock_client.get_node = MagicMock(return_value=mock_node)
|
||||
mock_node.write_value.side_effect = BadSessionIdInvalid()
|
||||
opc_repository._start_reconnect = MagicMock()
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
|
||||
futures = [
|
||||
executor.submit(
|
||||
opc_repository.write_data,
|
||||
node,
|
||||
value,
|
||||
'float',
|
||||
metadata['metadata'],
|
||||
)
|
||||
for node, value in (('ns=2;s=TestNode', 1.0), ('ns=2;s=TestNode2', 2.0))
|
||||
]
|
||||
results = [future.result() for future in futures]
|
||||
|
||||
assert 1 <= opc_repository._start_reconnect.call_count <= 2
|
||||
assert mock_node.write_value.call_count == 2
|
||||
error_kinds = [r[1].get('opc_error_kind') for r in results]
|
||||
assert error_kinds.count('session_bad') >= 1
|
||||
assert all(k in ('session_bad', 'reconnect_in_progress') for k in error_kinds)
|
||||
|
||||
|
||||
@patch('laborious.utils.repository.opc_repository.datetime')
|
||||
def test_reconnect_locked_sets_last_reconnection_time(mock_datetime, opc_repository):
|
||||
mock_datetime.now = MagicMock(return_value=datetime(2025, 1, 1, 12, 0, 0))
|
||||
opc_repository._disconnect_locked = MagicMock()
|
||||
opc_repository._connect_locked = MagicMock(return_value=(True, {}))
|
||||
|
||||
result = opc_repository._reconnect_locked()
|
||||
|
||||
opc_repository._disconnect_locked.assert_called_once()
|
||||
opc_repository._connect_locked.assert_called_once()
|
||||
assert result == (True, {})
|
||||
assert opc_repository.last_reconnection_time == datetime(2025, 1, 1, 12, 0, 0)
|
||||
|
||||
@@ -1,14 +1,16 @@
|
||||
from os import environ
|
||||
from laborious.utils.connectors_config import (build_mlflow_config,
|
||||
build_opc_config,
|
||||
build_postgres_config,
|
||||
build_mongodb_config)
|
||||
|
||||
from laborious.utils.connectors_config import (
|
||||
build_minio_config,
|
||||
build_mlflow_config,
|
||||
build_opc_config,
|
||||
build_plugin_store_config,
|
||||
)
|
||||
|
||||
|
||||
def test_build_mlflow_config_with_env_vars():
|
||||
# Arrange
|
||||
environ['MLFLOW_HOST'] = 'http://test-host'
|
||||
environ['MLFLOW_PORT'] = '8080'
|
||||
environ['MLFLOW_URL'] = 'http://test-host:8080'
|
||||
environ['MLFLOW_USERNAME'] = 'test-user'
|
||||
environ['MLFLOW_PASSWORD'] = 'test-pass'
|
||||
|
||||
@@ -16,17 +18,25 @@ def test_build_mlflow_config_with_env_vars():
|
||||
config = build_mlflow_config()
|
||||
|
||||
# Assert
|
||||
assert config['host'] == 'http://test-host'
|
||||
assert config['port'] == 8080
|
||||
assert config['url'] == 'http://test-host:8080'
|
||||
assert config['username'] == 'test-user'
|
||||
assert config['password'] == 'test-pass'
|
||||
|
||||
|
||||
def test_build_mlflow_config_host_already_has_port():
|
||||
environ['MLFLOW_URL'] = 'http://tracker.example.com:443'
|
||||
environ['MLFLOW_USERNAME'] = 'u'
|
||||
environ['MLFLOW_PASSWORD'] = 'p'
|
||||
|
||||
config = build_mlflow_config()
|
||||
|
||||
assert config['url'] == 'http://tracker.example.com:443'
|
||||
|
||||
|
||||
def test_build_mlflow_config_with_defaults():
|
||||
# Arrange
|
||||
# Clear any existing env vars
|
||||
environ.pop('MLFLOW_HOST', None)
|
||||
environ.pop('MLFLOW_PORT', None)
|
||||
environ.pop('MLFLOW_URL', None)
|
||||
environ.pop('MLFLOW_USERNAME', None)
|
||||
environ.pop('MLFLOW_PASSWORD', None)
|
||||
|
||||
@@ -34,12 +44,30 @@ def test_build_mlflow_config_with_defaults():
|
||||
config = build_mlflow_config()
|
||||
|
||||
# Assert
|
||||
assert config['host'] == 'http://localhost'
|
||||
assert config['port'] == 5080
|
||||
assert config['url'] == 'http://localhost:5080'
|
||||
assert config['username'] == 'aignosi'
|
||||
assert config['password'] == 'aignosi'
|
||||
|
||||
|
||||
def test_build_plugin_store_config_defaults():
|
||||
environ.pop('STORE_BASE_URL', None)
|
||||
environ.pop('STORE_OWNER', None)
|
||||
environ.pop('STORE_REPO', None)
|
||||
environ.pop('STORE_BRANCH', None)
|
||||
environ.pop('STORE_USERNAME', None)
|
||||
environ.pop('STORE_PASSWORD', None)
|
||||
environ.pop('STORE_CACHE_TTL_SECONDS', None)
|
||||
environ.pop('PYPI_SERVER', None)
|
||||
environ.pop('PYPI_USERNAME', None)
|
||||
environ.pop('PYPI_PASSWORD', None)
|
||||
|
||||
cfg = build_plugin_store_config()
|
||||
assert cfg['base_url'] == 'http://localhost:3000'
|
||||
assert cfg['owner'] == 'sientia'
|
||||
assert cfg['repo'] == 'model-library-store'
|
||||
assert cfg['pypi_index_url'] == 'http://localhost:5000'
|
||||
|
||||
|
||||
def test_build_opc_config_with_env_vars():
|
||||
# Arrange
|
||||
environ['OPC_CONFIG'] = '{"opc": {"name": "test-opc", "url": "opc.tcp://test:4840"}}'
|
||||
@@ -88,74 +116,37 @@ def test_build_opc_config_with_defaults():
|
||||
assert config['1']['reconnection_interval'] == 120
|
||||
|
||||
|
||||
def test_build_postgres_config_with_env_vars():
|
||||
# Arrange
|
||||
environ['POSTGRES_HOST'] = 'test-host'
|
||||
environ['POSTGRES_PORT'] = '5433'
|
||||
environ['POSTGRES_USER'] = 'test-user'
|
||||
environ['POSTGRES_PASSWORD'] = 'test-pass'
|
||||
environ['POSTGRES_DBNAME'] = 'test-db'
|
||||
environ['POSTGRES_MIN_CONNECTIONS'] = '10'
|
||||
environ['POSTGRES_MAX_CONNECTIONS'] = '30'
|
||||
|
||||
# Act
|
||||
config = build_postgres_config()
|
||||
|
||||
# Assert
|
||||
assert config['host'] == 'test-host'
|
||||
assert config['port'] == 5433
|
||||
assert config['user'] == 'test-user'
|
||||
assert config['password'] == 'test-pass'
|
||||
assert config['dbname'] == 'test-db'
|
||||
assert config['min_connections'] == 10
|
||||
assert config['max_connections'] == 30
|
||||
|
||||
|
||||
def test_build_postgres_config_with_defaults():
|
||||
# Arrange
|
||||
environ.pop('POSTGRES_HOST', None)
|
||||
environ.pop('POSTGRES_PORT', None)
|
||||
environ.pop('POSTGRES_USER', None)
|
||||
environ.pop('POSTGRES_PASSWORD', None)
|
||||
environ.pop('POSTGRES_DBNAME', None)
|
||||
environ.pop('POSTGRES_MIN_CONNECTIONS', None)
|
||||
environ.pop('POSTGRES_MAX_CONNECTIONS', None)
|
||||
|
||||
# Act
|
||||
config = build_postgres_config()
|
||||
|
||||
# Assert
|
||||
assert config['host'] == 'localhost'
|
||||
assert config['port'] == 5432
|
||||
assert config['user'] == 'sientia'
|
||||
assert config['password'] == 'sientia'
|
||||
assert config['dbname'] == 'sientia'
|
||||
assert config['min_connections'] == 5
|
||||
assert config['max_connections'] == 20
|
||||
|
||||
|
||||
def test_build_mongo_db_config_with_env_vars():
|
||||
environ['MONGODB_USERNAME'] = 'sientia1'
|
||||
environ['MONGODB_PASSWORD'] = 'sientia1'
|
||||
environ['MONGODB_URL'] = 'localhost:27018'
|
||||
environ['MONGODB_DATABASE_NAME'] = 'test_db'
|
||||
environ['MONGODB_TTL_INDEX_HOURS'] = '1'
|
||||
|
||||
assert build_mongodb_config() == {
|
||||
'connection_string': 'mongodb://sientia1:sientia1@localhost:27018',
|
||||
'database_name': 'test_db',
|
||||
'ttl_index_seconds': 3600
|
||||
def test_build_minio_config_with_env_vars():
|
||||
environ['MINIO_ENDPOINT_URL'] = 'http://test-host'
|
||||
environ['MINIO_ACCESS_KEY'] = 'test-key'
|
||||
environ['MINIO_SECRET_KEY'] = 'test-secret'
|
||||
environ['MINIO_REGION_NAME'] = 'test-region'
|
||||
environ['MINIO_DEFAULT_BUCKET'] = 'test-bucket'
|
||||
# Isolate from IDE/CI env (e.g. VS Code may export MINIO_SECURE=true).
|
||||
environ['MINIO_SECURE'] = 'false'
|
||||
assert build_minio_config() == {
|
||||
'endpoint_url': 'http://test-host',
|
||||
'access_key': 'test-key',
|
||||
'secret_key': 'test-secret',
|
||||
'default_bucket': 'test-bucket',
|
||||
'retention_hours': 24,
|
||||
'secure': False,
|
||||
}
|
||||
|
||||
|
||||
def test_build_mongo_db_config_with_defaults():
|
||||
environ.pop('MONGODB_USERNAME', None)
|
||||
environ.pop('MONGODB_PASSWORD', None)
|
||||
environ.pop('MONGODB_DATABASE_NAME', None)
|
||||
environ.pop('MONGODB_URL', None)
|
||||
environ.pop('MONGODB_TTL_INDEX_HOURS', None)
|
||||
assert build_mongodb_config() == {
|
||||
'connection_string': 'mongodb://root:wKZDbMNU1c@localhost:27018',
|
||||
'database_name': 'sientia',
|
||||
'ttl_index_seconds': 3600
|
||||
def test_build_minio_config_with_defaults():
|
||||
environ.pop('MINIO_ENDPOINT_URL', None)
|
||||
environ.pop('MINIO_ACCESS_KEY', None)
|
||||
environ.pop('MINIO_SECRET_KEY', None)
|
||||
environ.pop('MINIO_REGION_NAME', None)
|
||||
environ.pop('MINIO_DEFAULT_BUCKET', None)
|
||||
environ.pop('MINIO_SECURE', None)
|
||||
environ.pop('MINIO_RETENTION_HOURS', None)
|
||||
assert build_minio_config() == {
|
||||
'endpoint_url': 'http://localhost:9000',
|
||||
'access_key': 'minioadmin',
|
||||
'secret_key': 'minioadmin',
|
||||
'default_bucket': 'laborious',
|
||||
'retention_hours': 24,
|
||||
'secure': False,
|
||||
}
|
||||
|
||||
12
tests/laborious/utils/test_dataframe_debug.py
Normal file
12
tests/laborious/utils/test_dataframe_debug.py
Normal file
@@ -0,0 +1,12 @@
|
||||
from pandas import DataFrame
|
||||
|
||||
from laborious.utils.dataframe_debug import build_dataframe_debug_message
|
||||
|
||||
|
||||
def test_build_dataframe_debug_message_skips_large_dataframe():
|
||||
df = DataFrame({'a': [1, 2, 3]})
|
||||
|
||||
msg = build_dataframe_debug_message('payload', df, max_rows=1)
|
||||
|
||||
assert 'skipped because dataframe has 3 rows' in msg
|
||||
assert '(max: 1)' in msg
|
||||
Reference in New Issue
Block a user