sientia-dataops-model-manager/tests/activities/test_cleanup.py

"""Unit tests for the Cleanup activity, ensuring 100% code coverage."""

import asyncio
import os
import shutil
import tempfile
from datetime import datetime, timedelta
from importlib import reload
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

# Define mocks at the top level to be accessible by all tests


@pytest.fixture
def mock_logger():
    """Fixture for a mock logger."""
    return MagicMock()


@pytest.fixture
def mock_notification_handler():
    """Fixture for a mock notification handler."""
    return MagicMock()


@pytest.fixture
def mock_metrics_controller():
    """Fixture for a mock metrics controller with async methods."""
    controller = MagicMock()
    controller.shutdown = AsyncMock()
    controller.emit = AsyncMock()
    return controller


@pytest.fixture
def temp_dir():
    """Fixture to create and clean up a temporary directory."""
    path = tempfile.mkdtemp()
    yield path
    shutil.rmtree(path)


# --- Initialization Tests ---


@patch.dict(
    'model_manager.activities.cleanup.os.environ',
    {
        'CLEANUP_RETENTION_HOURS': '24',
        'CLEANUP_DRY_RUN': 'false',
    },
)
def test_cleanup_init_default_values(
    mock_logger,
    mock_notification_handler,
    mock_metrics_controller,
):
    """Test Cleanup initialization uses default environment values."""
    import model_manager.activities.cleanup

    reload(model_manager.activities.cleanup)
    from model_manager.activities.cleanup import Cleanup

    cleanup = Cleanup(
        logger=mock_logger,
        notification_handler=mock_notification_handler,
        metrics_controller=mock_metrics_controller,
    )

    assert cleanup.retention_hours == 24
    assert cleanup.dry_run is False


def test_cleanup_init_custom_env_values(
    mock_logger,
    mock_notification_handler,
    mock_metrics_controller,
):
    """Test Cleanup initialization with custom environment values."""
    with patch.dict(
        os.environ,
        {
            'CLEANUP_RETENTION_HOURS': '48',
            'CLEANUP_DRY_RUN': 'true',
        },
    ):
        import model_manager.activities.cleanup

        reload(model_manager.activities.cleanup)
        from model_manager.activities.cleanup import Cleanup

        cleanup = Cleanup(
            logger=mock_logger,
            notification_handler=mock_notification_handler,
            metrics_controller=mock_metrics_controller,
        )

        assert cleanup.retention_hours == 48
        assert cleanup.dry_run is True


@patch.dict(os.environ, {'CLEANUP_RETENTION_HOURS': 'invalid'})
def test_cleanup_init_invalid_env_value_raises_error():
    """Test Cleanup module raises ValueError for invalid environment variables on import."""
    import model_manager.activities.cleanup

    with pytest.raises(ValueError):
        reload(model_manager.activities.cleanup)


# --- Temp Directory Cleanup Tests ---


def test_cleanup_temp_directories_nonexistent_path(
    mock_logger,
    mock_notification_handler,
    mock_metrics_controller,
):
    """Test temp directory cleanup with a non-existent path."""
    from model_manager.activities.cleanup import Cleanup

    cleanup = Cleanup(
        logger=mock_logger,
        notification_handler=mock_notification_handler,
        metrics_controller=mock_metrics_controller,
    )
    cleanup._emit_metrics = AsyncMock()
    cleanup.warning = MagicMock()

    asyncio.run(
        cleanup.cleanup_temp_directories({'temp_path': '/nonexistent/path', 'metadata': {}})
    )

    cleanup.warning.assert_called_once()
    cleanup._emit_metrics.assert_called_once()


@patch.dict('model_manager.activities.cleanup.os.environ', {'CLEANUP_DRY_RUN': 'false'})
def test_cleanup_temp_directories_success_with_deletions(
    temp_dir,
    mock_logger,
    mock_notification_handler,
    mock_metrics_controller,
):
    """Test successful deletion of old temporary directories."""
    import model_manager.activities.cleanup

    reload(model_manager.activities.cleanup)
    from model_manager.activities.cleanup import Cleanup

    cleanup = Cleanup(
        logger=mock_logger,
        notification_handler=mock_notification_handler,
        metrics_controller=mock_metrics_controller,
    )
    cleanup._emit_metrics = AsyncMock()

    old_time = (datetime.now() - timedelta(hours=48)).strftime('%Y%m%d_%H%M%S_000000')
    old_dir = os.path.join(temp_dir, f'old_dir_{old_time}')
    os.makedirs(old_dir)

    recent_time = (datetime.now() - timedelta(hours=1)).strftime('%Y%m%d_%H%M%S_000000')
    recent_dir = os.path.join(temp_dir, f'recent_dir_{recent_time}')
    os.makedirs(recent_dir)

    asyncio.run(cleanup.cleanup_temp_directories({'temp_path': temp_dir, 'metadata': {}}))

    assert not os.path.exists(old_dir)
    assert os.path.exists(recent_dir)
    cleanup._emit_metrics.assert_called_once()


@patch.dict(os.environ, {'CLEANUP_DRY_RUN': 'true'})
def test_cleanup_temp_directories_dry_run(
    temp_dir,
    mock_logger,
    mock_notification_handler,
    mock_metrics_controller,
):
    """Test temp directory cleanup in dry_run mode does not delete."""
    import model_manager.activities.cleanup

    reload(model_manager.activities.cleanup)
    from model_manager.activities.cleanup import Cleanup

    cleanup = Cleanup(
        logger=mock_logger,
        notification_handler=mock_notification_handler,
        metrics_controller=mock_metrics_controller,
    )
    cleanup._emit_metrics = AsyncMock()

    old_time = (datetime.now() - timedelta(hours=48)).strftime('%Y%m%d_%H%M%S_000000')
    old_dir = os.path.join(temp_dir, f'old_dir_{old_time}')
    os.makedirs(old_dir)

    asyncio.run(cleanup.cleanup_temp_directories({'temp_path': temp_dir, 'metadata': {}}))

    assert os.path.exists(old_dir)
    cleanup._emit_metrics.assert_called_once()


@patch.dict('model_manager.activities.cleanup.os.environ', {'CLEANUP_DRY_RUN': 'false'})
def test_cleanup_temp_directories_delete_error(
    temp_dir,
    mock_logger,
    mock_notification_handler,
    mock_metrics_controller,
):
    """Test graceful handling of errors during directory deletion."""
    import model_manager.activities.cleanup

    reload(model_manager.activities.cleanup)
    from model_manager.activities.cleanup import Cleanup

    cleanup = Cleanup(
        logger=mock_logger,
        notification_handler=mock_notification_handler,
        metrics_controller=mock_metrics_controller,
    )
    cleanup._emit_metrics = AsyncMock()
    cleanup.error = MagicMock()

    old_time = (datetime.now() - timedelta(hours=48)).strftime('%Y%m%d_%H%M%S_000000')
    old_dir = os.path.join(temp_dir, f'old_dir_{old_time}')
    os.makedirs(old_dir)

    with patch('shutil.rmtree', side_effect=OSError('Permission Denied')):
        asyncio.run(cleanup.cleanup_temp_directories({'temp_path': temp_dir, 'metadata': {}}))

    cleanup.error.assert_called_once()
    cleanup._emit_metrics.assert_called_once()


# --- Metrics and Utility Tests ---


def test_emit_metrics(
    mock_logger,
    mock_notification_handler,
    mock_metrics_controller,
):
    """Test that _emit_metrics calls the public emit_metric method."""
    from model_manager.activities.cleanup import Cleanup

    cleanup = Cleanup(
        logger=mock_logger,
        notification_handler=mock_notification_handler,
        metrics_controller=mock_metrics_controller,
    )
    cleanup.emit_metric = AsyncMock()

    asyncio.run(
        cleanup._emit_metrics(
            metadata={'pod_id': 'p1', 'workflow_name': 'wf1'},
            metrics_status='success',
            activity_name='test_activity',
            emit_workflow_metric=True,
        )
    )

    assert cleanup.emit_metric.call_count == 2


def test_cleanup_temp_directories_with_files_and_unmatched_dirs(
    temp_dir,
    mock_logger,
    mock_notification_handler,
    mock_metrics_controller,
):
    """Test that files and directories with non-matching names are skipped."""
    import model_manager.activities.cleanup

    reload(model_manager.activities.cleanup)
    from model_manager.activities.cleanup import Cleanup

    cleanup = Cleanup(
        logger=mock_logger,
        notification_handler=mock_notification_handler,
        metrics_controller=mock_metrics_controller,
    )
    cleanup._emit_metrics = AsyncMock()
    cleanup.debug = MagicMock()

    # Create a file and a directory with a non-matching name
    with open(os.path.join(temp_dir, 'a_file.txt'), 'w') as f:
        f.write('hello')
    os.makedirs(os.path.join(temp_dir, 'a_directory_with_no_timestamp'))

    asyncio.run(cleanup.cleanup_temp_directories({'temp_path': temp_dir, 'metadata': {}}))

    # Ensure the debug message for skipping was called for the unmatched directory
    cleanup.debug.assert_called_with(
        'Skipping directory without timestamp pattern: a_directory_with_no_timestamp', {}
    )
    cleanup._emit_metrics.assert_called_once()


def test_cleanup_temp_directories_invalid_timestamp_format(
    temp_dir,
    mock_logger,
    mock_notification_handler,
    mock_metrics_controller,
):
    """Test that a directory with an invalid timestamp format is handled correctly."""
    import model_manager.activities.cleanup

    reload(model_manager.activities.cleanup)
    from model_manager.activities.cleanup import Cleanup

    cleanup = Cleanup(
        logger=mock_logger,
        notification_handler=mock_notification_handler,
        metrics_controller=mock_metrics_controller,
    )
    cleanup._emit_metrics = AsyncMock()
    cleanup.error = MagicMock()

    # Create a directory with a malformed timestamp that matches the regex but fails parsing
    malformed_dir_name = 'dir_20239999_999999_999999'
    os.makedirs(os.path.join(temp_dir, malformed_dir_name))

    asyncio.run(cleanup.cleanup_temp_directories({'temp_path': temp_dir, 'metadata': {}}))

    cleanup.error.assert_called_once()
    cleanup._emit_metrics.assert_called_once()


def test_cleanup_temp_directories_generic_exception(
    temp_dir,
    mock_logger,
    mock_notification_handler,
    mock_metrics_controller,
):
    """Test that a generic exception during directory cleanup is handled."""
    import model_manager.activities.cleanup

    reload(model_manager.activities.cleanup)
    from model_manager.activities.cleanup import Cleanup

    cleanup = Cleanup(
        logger=mock_logger,
        notification_handler=mock_notification_handler,
        metrics_controller=mock_metrics_controller,
    )
    cleanup._emit_metrics = AsyncMock()
    cleanup.send_notification = MagicMock()

    with patch('os.listdir', side_effect=Exception('Unexpected OS Error')):
        with pytest.raises(Exception, match='Unexpected OS Error'):
            asyncio.run(cleanup.cleanup_temp_directories({'temp_path': temp_dir, 'metadata': {}}))

    cleanup.send_notification.assert_called_once()
    cleanup._emit_metrics.assert_called_once()


def test_emit_metrics_activity_only(
    mock_logger,
    mock_notification_handler,
    mock_metrics_controller,
):
    """Test that _emit_metrics can emit only the activity metric."""
    from model_manager.activities.cleanup import Cleanup

    cleanup = Cleanup(
        logger=mock_logger,
        notification_handler=mock_notification_handler,
        metrics_controller=mock_metrics_controller,
    )
    cleanup.emit_metric = AsyncMock()

    asyncio.run(
        cleanup._emit_metrics(
            metadata={'pod_id': 'p1', 'workflow_name': 'wf1'},
            metrics_status='success',
            activity_name='test_activity',
            emit_workflow_metric=False,
        )
    )

    cleanup.emit_metric.assert_called_once()