SIENTIAPDE-1712
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.
This commit is contained in:
vitor-aignosi
@@ -1,8 +1,14 @@
|
||||
from datetime import datetime
|
||||
from io import BytesIO
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from pytest import mark
|
||||
import pytest
|
||||
from pandas import DataFrame
|
||||
|
||||
from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload
|
||||
from laborious.utils.models.minio_dataframe_payload import (
|
||||
MinioDataFramePayload,
|
||||
_build_object_key,
|
||||
)
|
||||
|
||||
|
||||
def test_parse_object_timestamp_hyphenated_model():
|
||||
@@ -21,34 +27,190 @@ def test_parse_object_timestamp_invalid():
|
||||
assert MinioDataFramePayload.parse_object_timestamp('bad.parquet') is None
|
||||
|
||||
|
||||
def test_is_offloaded_dict_true_false():
|
||||
assert MinioDataFramePayload.is_offloaded_dict({'object_key': 'k', 'data': None}) is True
|
||||
assert MinioDataFramePayload.is_offloaded_dict({'object_key': 'k', 'data': {}}) is False
|
||||
assert MinioDataFramePayload.is_offloaded_dict({'data': {}}) is False
|
||||
def test_estimate_size_bytes_returns_positive_for_nonempty_frame():
|
||||
df = DataFrame({'a': [1, 2]})
|
||||
size = MinioDataFramePayload.estimate_size_bytes(df)
|
||||
assert isinstance(size, int)
|
||||
assert size > 0
|
||||
|
||||
|
||||
def test_cleanup_prefix_from_payload_dict():
|
||||
p = {
|
||||
'object_key': 'sientia/streamlit-connectors/training_datasets/m/m-initial-2024-01-01_00-00-00.parquet',
|
||||
'bucket': 'b',
|
||||
'data': None,
|
||||
}
|
||||
assert MinioDataFramePayload.cleanup_prefix_from_payload_dict(p) == 'training_datasets/m'
|
||||
def test_cleanup_prefix_when_offloaded_returns_object_prefix():
|
||||
payload = MinioDataFramePayload(
|
||||
last_timestamp='t',
|
||||
data=None,
|
||||
object_key='training_datasets/m/m-initial-2024-01-01_00-00-00.parquet',
|
||||
object_prefix='training_datasets/m',
|
||||
)
|
||||
assert MinioDataFramePayload.cleanup_prefix(payload) == 'training_datasets/m'
|
||||
|
||||
|
||||
def test_cleanup_prefix_from_explicit_object_prefix():
|
||||
p = {'object_key': 'x.parquet', 'object_prefix': 'my/prefix', 'data': None}
|
||||
assert MinioDataFramePayload.cleanup_prefix_from_payload_dict(p) == 'my/prefix'
|
||||
def test_cleanup_prefix_when_inline_returns_none():
|
||||
payload = MinioDataFramePayload(last_timestamp='t', data={'x': [1]}, object_key=None)
|
||||
assert MinioDataFramePayload.cleanup_prefix(payload) is None
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_resolve_dict_if_offloaded_noop():
|
||||
d = {'success': True, 'data': {'a': [1]}}
|
||||
out = await MinioDataFramePayload.resolve_dict_if_offloaded(d, None, {})
|
||||
assert out is d
|
||||
def test_has_data_true_when_object_key_set():
|
||||
payload = MinioDataFramePayload(last_timestamp='t', data=None, object_key='k')
|
||||
assert payload.has_data() is True
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_dataframe_from_wire_list():
|
||||
df = await MinioDataFramePayload.dataframe_from_wire([{'a': 1}], None, {})
|
||||
assert list(df.columns) == ['a']
|
||||
@pytest.mark.asyncio
|
||||
async def test_retrieve_inline_dict_as_dataframe():
|
||||
payload = MinioDataFramePayload(last_timestamp='t', data={'a': [1, 2]})
|
||||
minio = AsyncMock()
|
||||
out = await payload.retrieve(minio, {'metadata': {}})
|
||||
assert list(out.columns) == ['a']
|
||||
minio.download_file.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retrieve_downloads_parquet_when_offloaded():
|
||||
source = DataFrame({'a': [1, 2]})
|
||||
buf = BytesIO()
|
||||
source.to_parquet(buf, engine='pyarrow', index=True)
|
||||
file_bytes = buf.getvalue()
|
||||
|
||||
payload = MinioDataFramePayload(
|
||||
last_timestamp='t',
|
||||
data=None,
|
||||
object_key='training_datasets/m/f.parquet',
|
||||
object_prefix='training_datasets/m',
|
||||
)
|
||||
minio = AsyncMock()
|
||||
minio.download_file = AsyncMock(return_value=file_bytes)
|
||||
|
||||
out = await payload.retrieve(minio, {'metadata': {}})
|
||||
|
||||
minio.download_file.assert_awaited_once_with(
|
||||
object_name='training_datasets/m/f.parquet',
|
||||
metadata={'metadata': {}},
|
||||
)
|
||||
assert list(out.columns) == ['a']
|
||||
|
||||
|
||||
def test_build_object_key():
|
||||
key, prefix = _build_object_key('my-model', 'initial', '2024-01-01_00-00-00')
|
||||
assert key == 'training_datasets/my-model/my-model-initial-2024-01-01_00-00-00.parquet'
|
||||
assert prefix == 'training_datasets/my-model'
|
||||
|
||||
|
||||
def test_build_object_key_strips_slashes():
|
||||
key, prefix = _build_object_key(' /my-model/ ', 'transform', '2024-06-15_10-30-45')
|
||||
assert prefix == 'training_datasets/my-model'
|
||||
assert key.startswith('training_datasets/my-model/')
|
||||
|
||||
|
||||
def test_estimate_size_bytes_fallback():
|
||||
df = DataFrame({'a': [1, 2]})
|
||||
original_to_dict = df.to_dict
|
||||
df.to_dict = lambda *a, **kw: (_ for _ in ()).throw(RuntimeError('to_dict failed'))
|
||||
size = MinioDataFramePayload.estimate_size_bytes(df)
|
||||
df.to_dict = original_to_dict
|
||||
assert isinstance(size, int)
|
||||
assert size > 0
|
||||
|
||||
|
||||
def test_parse_object_timestamp_bad_datetime():
|
||||
key = 'p/m-initial-9999-99-99_99-99-99.parquet'
|
||||
assert MinioDataFramePayload.parse_object_timestamp(key) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retrieve_empty_when_no_data():
|
||||
payload = MinioDataFramePayload(last_timestamp='t', data=None, object_key=None)
|
||||
minio = AsyncMock()
|
||||
out = await payload.retrieve(minio, {})
|
||||
assert out.empty
|
||||
minio.download_file.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch('laborious.utils.models.minio_dataframe_payload.now')
|
||||
async def test_from_dataframe_none(mock_now):
|
||||
mock_now.return_value = datetime(2024, 1, 1, 0, 0, 0)
|
||||
minio = AsyncMock()
|
||||
result = await MinioDataFramePayload.from_dataframe(
|
||||
dataframe=None,
|
||||
minio_repo=minio,
|
||||
model_name='m',
|
||||
operation='initial',
|
||||
status={'success': False, 'message': 'no data'},
|
||||
)
|
||||
assert result.data is None
|
||||
assert result.status == {'success': False, 'message': 'no data'}
|
||||
assert result.object_key is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch('laborious.utils.models.minio_dataframe_payload.now')
|
||||
async def test_from_dataframe_empty(mock_now):
|
||||
mock_now.return_value = datetime(2024, 1, 1, 0, 0, 0)
|
||||
minio = AsyncMock()
|
||||
mock_df = MagicMock()
|
||||
mock_df.__bool__ = MagicMock(return_value=True)
|
||||
mock_df.empty = True
|
||||
result = await MinioDataFramePayload.from_dataframe(
|
||||
dataframe=mock_df,
|
||||
minio_repo=minio,
|
||||
model_name='m',
|
||||
operation='initial',
|
||||
)
|
||||
assert result.data is None
|
||||
assert result.object_key is None
|
||||
|
||||
|
||||
def _mock_dataframe(data_dict, timestamp_values=None):
|
||||
"""Build a MagicMock that behaves enough like a DataFrame for from_dataframe."""
|
||||
mock_df = MagicMock()
|
||||
mock_df.__bool__ = MagicMock(return_value=True)
|
||||
mock_df.empty = False
|
||||
if timestamp_values is None:
|
||||
timestamp_values = data_dict.get('timestamp', ['2024-01-01'])
|
||||
ts_col = MagicMock()
|
||||
ts_col.values.tolist.return_value = timestamp_values
|
||||
mock_df.__getitem__ = MagicMock(return_value=ts_col)
|
||||
mock_df.to_dict.return_value = data_dict
|
||||
buf = BytesIO()
|
||||
DataFrame(data_dict).to_parquet(buf, engine='pyarrow', index=True)
|
||||
mock_df.to_parquet = MagicMock(side_effect=lambda b, **kw: b.write(buf.getvalue()))
|
||||
return mock_df
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch('laborious.utils.models.minio_dataframe_payload.OFFLOAD_THRESHOLD_BYTES', 10**9)
|
||||
async def test_from_dataframe_inline():
|
||||
minio = AsyncMock()
|
||||
df = _mock_dataframe({'timestamp': ['2024-01-01'], 'value': [42]})
|
||||
result = await MinioDataFramePayload.from_dataframe(
|
||||
dataframe=df,
|
||||
minio_repo=minio,
|
||||
model_name='m',
|
||||
operation='initial',
|
||||
)
|
||||
assert result.data is not None
|
||||
assert result.object_key is None
|
||||
assert result.last_timestamp == '2024-01-01'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch('laborious.utils.models.minio_dataframe_payload.now')
|
||||
@patch('laborious.utils.models.minio_dataframe_payload.OFFLOAD_THRESHOLD_BYTES', 0)
|
||||
async def test_from_dataframe_offloaded(mock_now):
|
||||
mock_now.return_value = datetime(2024, 1, 1, 0, 0, 0)
|
||||
minio = AsyncMock()
|
||||
minio.upload_file = AsyncMock(return_value={'minio_object_name': 'full/key.parquet'})
|
||||
minio.bucket = 'test-bucket'
|
||||
|
||||
df = _mock_dataframe({'timestamp': ['2024-01-01'], 'value': [42]})
|
||||
result = await MinioDataFramePayload.from_dataframe(
|
||||
dataframe=df,
|
||||
minio_repo=minio,
|
||||
model_name='m',
|
||||
operation='initial',
|
||||
workflow_metadata={'wf': 'data'},
|
||||
)
|
||||
assert result.data is None
|
||||
assert result.object_key == 'full/key.parquet'
|
||||
assert result.bucket == 'test-bucket'
|
||||
assert result.uri == 's3://test-bucket/full/key.parquet'
|
||||
minio.upload_file.assert_awaited_once()
|
||||
|
||||
Reference in New Issue
Block a user