vitor-aignosi
5d0d049082
Remove code validation script and refactor imports in activities and workflows - Deleted the `validate.sh` script, which was responsible for running code quality checks. - Cleaned up import statements in `activities.py`, `gates.py`, `mlflow.py`, and `storage.py` by removing unused imports and organizing them. - Refactored initialization methods in `MinioManager` and `MLFlow` classes for improved readability. - Updated various workflows to ensure compatibility with the new structure and removed unnecessary comments. - Enhanced test cases to accommodate changes in the activities and workflows, ensuring proper mocking of dependencies.
384 lines
13 KiB
Python
384 lines
13 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
|
|
|
|
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
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_query_to_minio_minio_repository_not_initialized(storage):
|
|
storage.minio_repository = None
|
|
|
|
with raises(ValueError) as e:
|
|
await storage.query_to_minio({})
|
|
|
|
assert str(e.value) == 'Minio repository not initialized'
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_query_to_minio_not_data(storage):
|
|
storage.load_custom_query = AsyncMock(return_value=None)
|
|
result = await storage.query_to_minio({})
|
|
|
|
storage.load_custom_query.assert_called_once_with({})
|
|
assert result['success'] is False
|
|
assert result['message'] == 'No data returned from query'
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.activities.storage.pd.DataFrame')
|
|
@patch('laborious.activities.storage.now')
|
|
async def test_query_to_minio_success(now, dataframe, storage):
|
|
data = [{'a': 1}, {'a': 2}, {'a': 3}]
|
|
storage.load_custom_query = AsyncMock(return_value=data)
|
|
now.return_value = datetime.datetime(2024, 1, 1, 0, 0, 0)
|
|
storage.minio_repository.upload_file = AsyncMock(
|
|
return_value={
|
|
'minio_object_name': 'sientia/streamlit-connectors/training_datasets/test_model/test_2024-01-01_00-00-00.parquet'
|
|
}
|
|
)
|
|
storage.minio_repository.bucket = 'test'
|
|
|
|
result = await storage.query_to_minio({'object_prefix': 'test', **metadata})
|
|
|
|
dataframe.assert_called_once_with(data)
|
|
|
|
storage.minio_repository.upload_file.assert_called_once_with(
|
|
file_bytes=ANY,
|
|
relative_key='training_datasets/test_model/test_2024-01-01_00-00-00.parquet',
|
|
metadata=metadata['metadata'],
|
|
)
|
|
|
|
assert result['success'] is True
|
|
assert (
|
|
result['object_key']
|
|
== 'sientia/streamlit-connectors/training_datasets/test_model/test_2024-01-01_00-00-00.parquet'
|
|
)
|
|
assert (
|
|
result['uri']
|
|
== 's3://test/sientia/streamlit-connectors/training_datasets/test_model/test_2024-01-01_00-00-00.parquet'
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_query_to_minio_error(storage):
|
|
storage.send_notification = MagicMock()
|
|
storage.send_notification_async = AsyncMock()
|
|
|
|
storage.load_custom_query = AsyncMock(side_effect=Exception('test'))
|
|
result = await storage.query_to_minio({**metadata, 'object_prefix': 'test'})
|
|
assert result['success'] is False
|
|
assert result['message'] == 'test'
|
|
storage.send_notification_async.assert_called_once_with(
|
|
metadata=metadata['metadata'],
|
|
notification_id='ERROR_STORING_QUERY_TO_MINIO',
|
|
message='Error storing query to MinIO: test',
|
|
block='query_to_minio',
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=ANY,
|
|
)
|
|
|
|
|
|
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()
|
|
|
|
result = await storage.cleanup_minio_objects_expired(
|
|
{**metadata, 'prefix': 'training_datasets/m'}
|
|
)
|
|
|
|
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
|
|
|
|
with raises(ValueError, match='Minio repository not initialized'):
|
|
await storage.cleanup_minio_objects_expired({**metadata, 'prefix': 'test'})
|
|
|
|
|
|
@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()
|
|
|
|
result = await storage.cleanup_minio_objects_expired({**metadata, 'prefix': 'test'})
|
|
|
|
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()
|
|
|
|
result = await storage.cleanup_minio_objects_expired(
|
|
{**metadata, 'prefix': 'training_datasets/m'}
|
|
)
|
|
|
|
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()
|
|
|
|
result = await storage.cleanup_minio_objects_expired(
|
|
{**metadata, 'prefix': 'training_datasets/m'}
|
|
)
|
|
|
|
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()
|