Snapshot of fix/QTZPOC-13 source tree

Code-only import without upstream history.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

tests/laborious/utils/models/test_minio_dataframe_payload.py

@@ -0,0 +1,266 @@
+from datetime import datetime
+from io import BytesIO
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+from pandas import DataFrame
+
+from laborious.utils.models.minio_dataframe_payload import (
+    MinioDataFramePayload,
+    _build_object_key,
+)
+
+
+def test_parse_object_timestamp_hyphenated_model():
+    key = 'predictions/sched/my-long-model-initial-2024-06-15_10-30-45.parquet'
+    ts = MinioDataFramePayload.parse_object_timestamp(key)
+    assert ts == datetime(2024, 6, 15, 10, 30, 45)
+
+
+def test_parse_object_timestamp_transform():
+    key = 'p/m-transform-2024-01-02_03-04-05.parquet'
+    ts = MinioDataFramePayload.parse_object_timestamp(key)
+    assert ts == datetime(2024, 1, 2, 3, 4, 5)
+
+
+def test_parse_object_timestamp_invalid():
+    assert MinioDataFramePayload.parse_object_timestamp('bad.parquet') is None
+
+
+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_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_when_inline_returns_none():
+    payload = MinioDataFramePayload(last_timestamp='t', data={'x': [1]}, object_key=None)
+    assert MinioDataFramePayload.cleanup_prefix(payload) is None
+
+
+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
+
+
+@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 == 'prediction_datasets/my-model/my-model-initial-2024-01-01_00-00-00.parquet'
+    assert prefix == 'prediction_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 == 'prediction_datasets/my-model'
+    assert key.startswith('prediction_datasets/my-model/')
+
+
+def test_estimate_size_bytes_fallback():
+    df = DataFrame({'a': [1, 2]})
+    with patch.object(df, 'to_dict', side_effect=RuntimeError('to_dict failed')):
+        size = MinioDataFramePayload.estimate_size_bytes(df)
+    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()
+
+
+def test_from_dict_inline():
+    raw = {
+        'last_timestamp': '2024-01-01T00:00:00+00:00',
+        'status': None,
+        'data': {'col1': {0: 'val1'}},
+        'bucket': None,
+        'object_key': None,
+        'object_prefix': None,
+        'uri': None,
+    }
+    payload = MinioDataFramePayload.from_dict(raw)
+    assert isinstance(payload, MinioDataFramePayload)
+    assert payload.last_timestamp == '2024-01-01T00:00:00+00:00'
+    assert payload.data == {'col1': {0: 'val1'}}
+    assert payload.object_key is None
+
+
+def test_from_dict_offloaded():
+    raw = {
+        'last_timestamp': '2024-06-15T10:30:45+00:00',
+        'status': {'success': True},
+        'data': None,
+        'bucket': 'my-bucket',
+        'object_key': 'training_datasets/model/model-initial-2024-06-15_10-30-45.parquet',
+        'object_prefix': 'training_datasets/model',
+        'uri': 's3://my-bucket/training_datasets/model/model-initial-2024-06-15_10-30-45.parquet',
+    }
+    payload = MinioDataFramePayload.from_dict(raw)
+    assert isinstance(payload, MinioDataFramePayload)
+    assert payload.data is None
+    assert payload.bucket == 'my-bucket'
+    assert payload.object_key == raw['object_key']
+    assert payload.object_prefix == 'training_datasets/model'
+    assert payload.uri == raw['uri']
+    assert payload.status == {'success': True}
+
+
+def test_from_dict_minimal_keys():
+    raw = {'last_timestamp': '2024-01-01'}
+    payload = MinioDataFramePayload.from_dict(raw)
+    assert payload.last_timestamp == '2024-01-01'
+    assert payload.data is None
+    assert payload.bucket is None
+    assert payload.object_key is None
+
+
+def test_from_dict_passthrough_existing_instance():
+    original = MinioDataFramePayload(last_timestamp='2024-01-01', data={'a': 1}, bucket='b')
+    result = MinioDataFramePayload.from_dict(original)
+    assert result is original