vitor-aignosi
e6018af23f
Update README, requirements, and E2E tests for improved configuration and functionality - Enhanced the README with updated model configuration examples, including the addition of an alias for production. - Removed the `requirements-light.txt` file and updated `requirements-local.txt` and `requirements.txt` to replace `asyncua` with `opcua`. - Refactored E2E test scenarios to utilize scenario input files for better maintainability and clarity. - Improved test coverage for MinIO offload functionality and added new helper functions for loading scenario inputs. - Updated `values.yaml` to reflect new global configurations and environment variables for the laborious worker.
325 lines
11 KiB
Python
325 lines
11 KiB
Python
import datetime
|
|
import os
|
|
from unittest.mock import ANY, MagicMock, patch
|
|
|
|
from pytest import fixture, raises
|
|
from sientia_do.notifications.models import NotificationLevel
|
|
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
|
|
from sientia_do.temporal.activities.postgres_sync import Postgres
|
|
|
|
from laborious.activities.storage import Storage
|
|
|
|
|
|
@fixture(autouse=True)
|
|
def _passthrough_from_dict():
|
|
with patch(
|
|
'laborious.activities.storage.MinioDataFramePayload.from_dict', side_effect=lambda x: x
|
|
):
|
|
yield
|
|
|
|
|
|
@fixture(autouse=True)
|
|
def _patch_monitoring_shutdown():
|
|
"""
|
|
Avoid running real async SientiaMonitoring.shutdown when Storage.close runs inside tests.
|
|
"""
|
|
with patch.object(SientiaMonitoring, 'shutdown') as mock_shutdown:
|
|
yield mock_shutdown
|
|
|
|
|
|
metadata = {
|
|
'metadata': {
|
|
'model_id': 'test_model_id',
|
|
'model_name': 'test_model',
|
|
'workflow_name': 'test_workflow',
|
|
'schedule_name': 'test_schedule',
|
|
}
|
|
}
|
|
|
|
|
|
@fixture
|
|
@patch('laborious.activities.storage.MinioRepository')
|
|
def storage(mock_minio_repository):
|
|
return Storage(
|
|
host='localhost',
|
|
port=5432,
|
|
user='postgres',
|
|
password='postgres',
|
|
dbname='postgres',
|
|
min_connections=1,
|
|
max_connections=10,
|
|
retention_hours=24,
|
|
minio_repository=mock_minio_repository.return_value,
|
|
logger=MagicMock(),
|
|
notification_handler=MagicMock(),
|
|
metrics_controller=MagicMock(),
|
|
)
|
|
|
|
|
|
@patch('laborious.activities.storage.MinioRepository')
|
|
def test___init___not_hasattr(mock_minio_repository):
|
|
logger = MagicMock()
|
|
notification_handler = MagicMock()
|
|
metrics_controller = MagicMock()
|
|
minio_repo = mock_minio_repository.return_value
|
|
storage = Storage(
|
|
host='localhost',
|
|
port=5432,
|
|
user='postgres',
|
|
password='postgres',
|
|
dbname='postgres',
|
|
min_connections=1,
|
|
max_connections=10,
|
|
retention_hours=24,
|
|
minio_repository=minio_repo,
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=metrics_controller,
|
|
)
|
|
assert isinstance(storage, Postgres)
|
|
|
|
assert storage.minio_repository is minio_repo
|
|
mock_minio_repository.assert_not_called()
|
|
|
|
|
|
@patch('laborious.activities.storage.MinioRepository')
|
|
def test___init___none_minio_repository(mock_minio_repository, storage):
|
|
storage.minio_repository = None
|
|
logger = MagicMock()
|
|
notification_handler = MagicMock()
|
|
metrics_controller = MagicMock()
|
|
storage.__init__(
|
|
host='localhost',
|
|
port=5432,
|
|
user='postgres',
|
|
password='postgres',
|
|
dbname='postgres',
|
|
min_connections=1,
|
|
max_connections=10,
|
|
retention_hours=24,
|
|
minio_repository=None,
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=metrics_controller,
|
|
)
|
|
|
|
assert storage.minio_repository is None
|
|
mock_minio_repository.assert_not_called()
|
|
|
|
|
|
@patch('laborious.activities.storage.MinioRepository')
|
|
def test___init___done_repository(mock_minio_repository, storage):
|
|
storage.__init__(
|
|
host='localhost',
|
|
port=5432,
|
|
user='postgres',
|
|
password='postgres',
|
|
dbname='postgres',
|
|
min_connections=1,
|
|
max_connections=10,
|
|
retention_hours=24,
|
|
minio_repository=mock_minio_repository.return_value,
|
|
logger=MagicMock(),
|
|
notification_handler=MagicMock(),
|
|
metrics_controller=MagicMock(),
|
|
)
|
|
mock_minio_repository.assert_not_called()
|
|
assert storage.minio_repository is not None
|
|
|
|
|
|
def test_close(storage, _patch_monitoring_shutdown):
|
|
storage.minio_repository = MagicMock()
|
|
|
|
storage.close()
|
|
|
|
assert storage.minio_repository is None
|
|
_patch_monitoring_shutdown.assert_called_once_with(storage)
|
|
|
|
|
|
def test_close_when_minio_repository_already_none(storage, _patch_monitoring_shutdown):
|
|
"""Closing without an initialized MinIO repository skips MinIO teardown."""
|
|
storage.minio_repository = None
|
|
|
|
storage.close()
|
|
|
|
assert storage.minio_repository is None
|
|
_patch_monitoring_shutdown.assert_called_once_with(storage)
|
|
|
|
|
|
def test_load_query_with_minio_offload_no_rows(storage):
|
|
storage.load_custom_query = MagicMock(return_value=None)
|
|
storage_result = {'success': False}
|
|
with patch(
|
|
'laborious.activities.storage.MinioDataFramePayload.from_dataframe',
|
|
new_callable=MagicMock,
|
|
return_value=storage_result,
|
|
) as mock_from_dataframe:
|
|
result = storage.load_query_with_minio_offload(
|
|
{**metadata, 'query': 'SELECT 1', 'model_name': 'm', 'key_prefix': 'predictions/s'}
|
|
)
|
|
assert result == storage_result
|
|
mock_from_dataframe.assert_called_once()
|
|
|
|
|
|
def test_load_query_with_minio_offload_inline(storage):
|
|
storage.load_custom_query = MagicMock(return_value=[{'a': 1}])
|
|
storage_result = {'success': True, 'data': {'a': [1]}, 'object_key': None}
|
|
with patch(
|
|
'laborious.activities.storage.MinioDataFramePayload.from_dataframe',
|
|
new_callable=MagicMock,
|
|
return_value=storage_result,
|
|
) as mock_from_dataframe:
|
|
result = storage.load_query_with_minio_offload(
|
|
{
|
|
**metadata,
|
|
'query': 'SELECT 1',
|
|
'model_name': 'my-model',
|
|
'key_prefix': 'predictions/s',
|
|
}
|
|
)
|
|
assert result == storage_result
|
|
mock_from_dataframe.assert_called_once()
|
|
|
|
|
|
def test_load_query_with_minio_offload_minio(storage):
|
|
storage.load_custom_query = MagicMock(return_value=[{'a': 1}])
|
|
storage_result = {'success': True, 'data': None, 'object_key': 'object-key'}
|
|
with patch(
|
|
'laborious.activities.storage.MinioDataFramePayload.from_dataframe',
|
|
new_callable=MagicMock,
|
|
return_value=storage_result,
|
|
) as mock_from_dataframe:
|
|
result = storage.load_query_with_minio_offload(
|
|
{**metadata, 'query': 'SELECT 1', 'model_name': 'm', 'key_prefix': 'predictions/s'}
|
|
)
|
|
|
|
assert result == storage_result
|
|
mock_from_dataframe.assert_called_once()
|
|
|
|
|
|
@patch.dict(os.environ, {'SIENTIA_MINIO_RETENTION_HOURS': '1'})
|
|
@patch('laborious.activities.storage.now')
|
|
def test_cleanup_minio_objects_expired(mock_now, storage):
|
|
mock_now.return_value = datetime.datetime(2025, 1, 10, 12, 0, 0)
|
|
storage.minio_repository.list_objects = MagicMock(
|
|
return_value=[
|
|
'sientia/streamlit-connectors/training_datasets/m/m-initial-2024-12-01_00-00-00.parquet',
|
|
'sientia/streamlit-connectors/training_datasets/m/m-initial-2025-01-10_12-00-00.parquet',
|
|
]
|
|
)
|
|
storage.minio_repository.delete_file = MagicMock()
|
|
storage.send_notification = MagicMock()
|
|
|
|
data_mock = MagicMock()
|
|
data_mock.cleanup_prefix.return_value = 'training_datasets/m'
|
|
|
|
result = storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock})
|
|
|
|
assert result['deleted_count'] == 1
|
|
assert result['failed_count'] == 0
|
|
deleted_key = (
|
|
'sientia/streamlit-connectors/training_datasets/m/m-initial-2024-12-01_00-00-00.parquet'
|
|
)
|
|
assert deleted_key in result['deleted']
|
|
assert result['deleted'][deleted_key]['success'] is True
|
|
storage.minio_repository.list_objects.assert_called_once_with(
|
|
prefix='training_datasets/m',
|
|
recursive=True,
|
|
metadata=metadata['metadata'],
|
|
)
|
|
storage.minio_repository.delete_file.assert_called_once_with(
|
|
object_name='sientia/streamlit-connectors/training_datasets/m/m-initial-2024-12-01_00-00-00.parquet',
|
|
metadata=metadata['metadata'],
|
|
)
|
|
|
|
|
|
def test_load_query_with_minio_offload_minio_not_initialized(storage):
|
|
storage.minio_repository = None
|
|
|
|
with raises(ValueError, match='Minio repository not initialized'):
|
|
storage.load_query_with_minio_offload({**metadata, 'query': 'SELECT 1', 'model_name': 'm'})
|
|
|
|
|
|
def test_export_payload_to_postgres(storage):
|
|
payload = MagicMock()
|
|
payload.retrieve = MagicMock(return_value=MagicMock())
|
|
storage.export_data_to_postgres = MagicMock(return_value={'success': True})
|
|
|
|
result = storage.export_payload_to_postgres(
|
|
{**metadata, 'data': payload, 'schema': 'public', 'table': 't'}
|
|
)
|
|
|
|
payload.retrieve.assert_called_once_with(storage.minio_repository, metadata['metadata'])
|
|
storage.export_data_to_postgres.assert_called_once()
|
|
assert result == {'success': True}
|
|
|
|
|
|
def test_cleanup_minio_objects_expired_minio_not_initialized(storage):
|
|
storage.minio_repository = None
|
|
|
|
data_mock = MagicMock()
|
|
data_mock.cleanup_prefix.return_value = 'test'
|
|
with raises(ValueError, match='Minio repository not initialized'):
|
|
storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock})
|
|
|
|
|
|
@patch('laborious.activities.storage.now')
|
|
def test_cleanup_minio_objects_expired_unparseable_key(mock_now, storage):
|
|
mock_now.return_value = datetime.datetime(2025, 1, 10, 12, 0, 0)
|
|
storage.minio_repository.list_objects = MagicMock(
|
|
return_value=['some/random/key-without-timestamp.parquet']
|
|
)
|
|
storage.minio_repository.delete_file = MagicMock()
|
|
storage.send_notification = MagicMock()
|
|
|
|
data_mock = MagicMock()
|
|
data_mock.cleanup_prefix.return_value = 'test'
|
|
result = storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock})
|
|
|
|
assert result['deleted_count'] == 0
|
|
assert result['failed_count'] == 0
|
|
storage.minio_repository.delete_file.assert_not_called()
|
|
|
|
|
|
@patch('laborious.activities.storage.now')
|
|
def test_cleanup_minio_objects_expired_delete_fails(mock_now, storage):
|
|
mock_now.return_value = datetime.datetime(2025, 1, 10, 12, 0, 0)
|
|
old_key = 'training_datasets/m/m-initial-2024-12-01_00-00-00.parquet'
|
|
storage.minio_repository.list_objects = MagicMock(return_value=[old_key])
|
|
storage.minio_repository.delete_file = MagicMock(side_effect=Exception('delete error'))
|
|
storage.send_notification = MagicMock()
|
|
|
|
data_mock = MagicMock()
|
|
data_mock.cleanup_prefix.return_value = 'training_datasets/m'
|
|
result = storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock})
|
|
|
|
assert result['deleted_count'] == 0
|
|
assert result['failed_count'] == 1
|
|
assert old_key in result['failed']
|
|
assert result['failed'][old_key]['success'] is False
|
|
assert result['failed'][old_key]['message'] == 'delete error'
|
|
|
|
|
|
@patch('laborious.activities.storage.now')
|
|
def test_cleanup_minio_objects_expired_list_objects_error(mock_now, storage):
|
|
mock_now.return_value = datetime.datetime(2025, 1, 10, 12, 0, 0)
|
|
storage.minio_repository.list_objects = MagicMock(side_effect=Exception('list error'))
|
|
storage.send_notification = MagicMock()
|
|
storage.error = MagicMock()
|
|
|
|
data_mock = MagicMock()
|
|
data_mock.cleanup_prefix.return_value = 'training_datasets/m'
|
|
result = storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock})
|
|
|
|
assert result['deleted_count'] == 0
|
|
assert result['failed_count'] == 0
|
|
storage.send_notification.assert_called_once_with(
|
|
metadata=metadata['metadata'],
|
|
notification_id='ERROR_CLEANUP_MINIO_OBJECTS_EXPIRED',
|
|
message='Error cleaning up MinIO objects: list error',
|
|
block='cleanup_minio_objects_expired',
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=ANY,
|
|
)
|
|
storage.error.assert_called_once()
|