vitor-aignosi
ac795c7c53
Update .gitignore and refactor metrics.py for improved logging and consistency - Added coverage.xml to .gitignore to prevent tracking of coverage reports. - Refactored metric labels in metrics.py for consistency in string formatting and improved readability. - Enhanced logging messages in various activities to ensure uniformity in message formatting.
204 lines
6.0 KiB
Python
204 lines
6.0 KiB
Python
import datetime
|
|
from unittest.mock import ANY, AsyncMock, MagicMock, patch
|
|
|
|
from pytest import fixture, mark
|
|
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,
|
|
minio_config={
|
|
'endpoint_url': 'localhost:9000',
|
|
'access_key': 'minio',
|
|
'secret_key': 'minio123',
|
|
'region_name': 'us-east-1',
|
|
'default_bucket': 'test',
|
|
},
|
|
logger=MagicMock(),
|
|
notification_handler=MagicMock(),
|
|
)
|
|
|
|
|
|
@patch('laborious.activities.storage.MinioRepository')
|
|
def test___init___not_hasattr(mock_minio_repository):
|
|
logger = MagicMock()
|
|
notification_handler = MagicMock()
|
|
storage = Storage(
|
|
host='localhost',
|
|
port=5432,
|
|
user='postgres',
|
|
password='postgres',
|
|
dbname='postgres',
|
|
min_connections=1,
|
|
max_connections=10,
|
|
minio_config={
|
|
'endpoint_url': 'localhost:9000',
|
|
'access_key': 'minio',
|
|
'secret_key': 'minio123',
|
|
'region_name': 'us-east-1',
|
|
'default_bucket': 'test',
|
|
},
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
)
|
|
assert isinstance(storage, Postgres)
|
|
|
|
mock_minio_repository.assert_called_once_with(
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
minio_endpoint_url='localhost:9000',
|
|
minio_access_key='minio',
|
|
minio_secret_key='minio123',
|
|
minio_region_name='us-east-1',
|
|
minio_default_bucket='test',
|
|
)
|
|
|
|
|
|
@patch('laborious.activities.storage.MinioRepository')
|
|
def test___init___none_minio_repository(mock_minio_repository, storage):
|
|
storage.minio_repository = None
|
|
logger = MagicMock()
|
|
notification_handler = MagicMock()
|
|
|
|
storage.__init__(
|
|
host='localhost',
|
|
port=5432,
|
|
user='postgres',
|
|
password='postgres',
|
|
dbname='postgres',
|
|
min_connections=1,
|
|
max_connections=10,
|
|
minio_config={
|
|
'endpoint_url': 'localhost:9000',
|
|
'access_key': 'minio',
|
|
'secret_key': 'minio123',
|
|
'region_name': 'us-east-1',
|
|
'default_bucket': 'test',
|
|
},
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
)
|
|
|
|
mock_minio_repository.assert_called_once_with(
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
minio_endpoint_url='localhost:9000',
|
|
minio_access_key='minio',
|
|
minio_secret_key='minio123',
|
|
minio_region_name='us-east-1',
|
|
minio_default_bucket='test',
|
|
)
|
|
|
|
|
|
@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,
|
|
minio_config={
|
|
'endpoint_url': 'localhost:9000',
|
|
'access_key': 'minio',
|
|
'secret_key': 'minio123',
|
|
'region_name': 'us-east-1',
|
|
'default_bucket': 'test',
|
|
},
|
|
logger=MagicMock(),
|
|
notification_handler=MagicMock(),
|
|
)
|
|
mock_minio_repository.assert_not_called()
|
|
assert storage.minio_repository is not None
|
|
|
|
|
|
@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.minio_bucket = 'test'
|
|
|
|
result = await storage.query_to_minio({'object_prefix': 'test', **metadata})
|
|
|
|
dataframe.assert_called_once_with(data)
|
|
|
|
storage.minio_repository.store_dataframe_as_parquet.assert_called_once_with(
|
|
dataframe=dataframe.return_value,
|
|
uri='s3://test/test_2024-01-01_00-00-00.parquet',
|
|
object_name='test_2024-01-01_00-00-00.parquet',
|
|
metadata=metadata['metadata'],
|
|
)
|
|
|
|
assert result['success'] is True
|
|
assert result['object_key'] == 'test_2024-01-01_00-00-00.parquet'
|
|
assert result['uri'] == 's3://test/test_2024-01-01_00-00-00.parquet'
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_query_to_minio_error(storage):
|
|
storage.send_notification = MagicMock()
|
|
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.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()
|