vitor-aignosi
503d9aa485
Update dependencies and refactor input filter handling for consistency - Updated sientia-dataops-library dependency version from 1.10.3 to 1.10.4 in requirements.txt. - Refactored input filter handling in the Gates class to read policy and config keys in a case-insensitive manner. - Updated test cases to ensure consistency in filter key naming conventions across various scenarios.
324 lines
11 KiB
Python
324 lines
11 KiB
Python
import datetime
|
|
import os
|
|
from unittest.mock import ANY, AsyncMock, MagicMock, patch
|
|
|
|
from pytest import fixture, mark, raises
|
|
from sientia_do.notifications.models import NotificationLevel
|
|
from sientia_do.temporal.activities.postgres 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
|
|
|
|
|
|
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=AsyncMock(),
|
|
)
|
|
|
|
|
|
@patch('laborious.activities.storage.MinioRepository')
|
|
def test___init___not_hasattr(mock_minio_repository):
|
|
logger = MagicMock()
|
|
notification_handler = MagicMock()
|
|
metrics_controller = AsyncMock()
|
|
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 = AsyncMock()
|
|
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=AsyncMock(),
|
|
)
|
|
mock_minio_repository.assert_not_called()
|
|
assert storage.minio_repository is not None
|
|
|
|
|
|
def test_close(storage):
|
|
storage.minio_repository = MagicMock()
|
|
|
|
storage.close()
|
|
|
|
assert storage.minio_repository is None
|
|
|
|
|
|
def test___del__(storage):
|
|
storage.close = MagicMock()
|
|
|
|
storage.__del__()
|
|
|
|
storage.close.assert_called_once()
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_load_query_with_minio_offload_no_rows(storage):
|
|
storage.load_custom_query = AsyncMock(return_value=None)
|
|
storage_result = {'success': False}
|
|
with patch(
|
|
'laborious.activities.storage.MinioDataFramePayload.from_dataframe',
|
|
new_callable=AsyncMock,
|
|
return_value=storage_result,
|
|
) as mock_from_dataframe:
|
|
result = await 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_awaited_once()
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_load_query_with_minio_offload_inline(storage):
|
|
storage.load_custom_query = AsyncMock(return_value=[{'a': 1}])
|
|
storage_result = {'success': True, 'data': {'a': [1]}, 'object_key': None}
|
|
with patch(
|
|
'laborious.activities.storage.MinioDataFramePayload.from_dataframe',
|
|
new_callable=AsyncMock,
|
|
return_value=storage_result,
|
|
) as mock_from_dataframe:
|
|
result = await 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_awaited_once()
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_load_query_with_minio_offload_minio(storage):
|
|
storage.load_custom_query = AsyncMock(return_value=[{'a': 1}])
|
|
storage_result = {'success': True, 'data': None, 'object_key': 'object-key'}
|
|
with patch(
|
|
'laborious.activities.storage.MinioDataFramePayload.from_dataframe',
|
|
new_callable=AsyncMock,
|
|
return_value=storage_result,
|
|
) as mock_from_dataframe:
|
|
result = await 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_awaited_once()
|
|
|
|
|
|
@mark.asyncio
|
|
@patch.dict(os.environ, {'SIENTIA_MINIO_RETENTION_HOURS': '1'})
|
|
@patch('laborious.activities.storage.now')
|
|
async 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 = AsyncMock(
|
|
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 = AsyncMock()
|
|
storage.send_notification_async = AsyncMock()
|
|
|
|
data_mock = MagicMock()
|
|
data_mock.cleanup_prefix.return_value = 'training_datasets/m'
|
|
|
|
result = await 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'],
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_load_query_with_minio_offload_minio_not_initialized(storage):
|
|
storage.minio_repository = None
|
|
|
|
with raises(ValueError, match='Minio repository not initialized'):
|
|
await storage.load_query_with_minio_offload(
|
|
{**metadata, 'query': 'SELECT 1', 'model_name': 'm'}
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_export_payload_to_postgres(storage):
|
|
payload = AsyncMock()
|
|
payload.retrieve = AsyncMock(return_value=MagicMock())
|
|
storage.export_data_to_postgres = AsyncMock(return_value={'success': True})
|
|
|
|
result = await storage.export_payload_to_postgres(
|
|
{**metadata, 'data': payload, 'schema': 'public', 'table': 't'}
|
|
)
|
|
|
|
payload.retrieve.assert_awaited_once_with(storage.minio_repository, metadata['metadata'])
|
|
storage.export_data_to_postgres.assert_awaited_once()
|
|
assert result == {'success': True}
|
|
|
|
|
|
@mark.asyncio
|
|
async 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'):
|
|
await storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock})
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.activities.storage.now')
|
|
async 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 = AsyncMock(
|
|
return_value=['some/random/key-without-timestamp.parquet']
|
|
)
|
|
storage.minio_repository.delete_file = AsyncMock()
|
|
storage.send_notification_async = AsyncMock()
|
|
|
|
data_mock = MagicMock()
|
|
data_mock.cleanup_prefix.return_value = 'test'
|
|
result = await 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()
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.activities.storage.now')
|
|
async 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 = AsyncMock(return_value=[old_key])
|
|
storage.minio_repository.delete_file = AsyncMock(side_effect=Exception('delete error'))
|
|
storage.send_notification_async = AsyncMock()
|
|
|
|
data_mock = MagicMock()
|
|
data_mock.cleanup_prefix.return_value = 'training_datasets/m'
|
|
result = await 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'
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.activities.storage.now')
|
|
async 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 = AsyncMock(side_effect=Exception('list error'))
|
|
storage.send_notification_async = AsyncMock()
|
|
storage.error = MagicMock()
|
|
|
|
data_mock = MagicMock()
|
|
data_mock.cleanup_prefix.return_value = 'training_datasets/m'
|
|
result = await storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock})
|
|
|
|
assert result['deleted_count'] == 0
|
|
assert result['failed_count'] == 0
|
|
storage.send_notification_async.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()
|