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()