"""Unit tests for local worker factory."""

from unittest.mock import MagicMock, patch


def test_prepare_worker_train_queue_uses_train_limits():
    from model_manager.worker.prepare_worker import prepare_worker
    from model_manager.workflows.train_model import TrainModel

    fake_worker = MagicMock()
    fake_client = MagicMock()
    fake_logger = MagicMock()

    with patch(
        'model_manager.worker.prepare_worker.Worker', return_value=fake_worker
    ) as worker_class:
        with patch.dict(
            'os.environ',
            {
                'TRAINMODEL_ACTIVITY_EXECUTOR_MAX_WORKERS': '3',
                'TRAINMODEL_MAX_CONCURRENT_ACTIVITIES': '6',
                'TRAINMODEL_MAX_CONCURRENT_WORKFLOW_TASKS': '10',
            },
            clear=False,
        ):
            worker = prepare_worker(
                main_workflow=TrainModel,
                other_workflows=[],
                activities=[],
                temporal_client=fake_client,
                logger=fake_logger,
            )

    assert worker is fake_worker
    worker_class.assert_called_once()
    kwargs = worker_class.call_args.kwargs
    assert kwargs['task_queue'] == 'train_model-queue'
    assert kwargs['max_concurrent_activities'] == 6
    assert kwargs['max_concurrent_workflow_tasks'] == 10
    assert kwargs['activity_executor']._max_workers == 3
    kwargs['activity_executor'].shutdown(wait=True, cancel_futures=True)


def test_prepare_worker_cleanup_queue_uses_cleanup_limits():
    from model_manager.worker.prepare_worker import prepare_worker
    from model_manager.workflows.cleanup_files import CleanupFiles

    fake_worker = MagicMock()
    fake_client = MagicMock()
    fake_logger = MagicMock()

    with patch(
        'model_manager.worker.prepare_worker.Worker', return_value=fake_worker
    ) as worker_class:
        with patch.dict(
            'os.environ',
            {
                'CLEANUPFILES_ACTIVITY_EXECUTOR_MAX_WORKERS': '5',
                'CLEANUPFILES_MAX_CONCURRENT_ACTIVITIES': '7',
            },
            clear=False,
        ):
            worker = prepare_worker(
                main_workflow=CleanupFiles,
                other_workflows=[],
                activities=[],
                temporal_client=fake_client,
                logger=fake_logger,
            )

    assert worker is fake_worker
    kwargs = worker_class.call_args.kwargs
    assert kwargs['task_queue'] == 'cleanup_files-queue'
    assert kwargs['max_concurrent_activities'] == 7
    assert kwargs['activity_executor']._max_workers == 5
    kwargs['activity_executor'].shutdown(wait=True, cancel_futures=True)