import errno
import os
import socket
import tempfile
from itertools import cycle
from unittest.mock import MagicMock, Mock, patch

import pytest
from billiard.pool import ApplyResult
from kombu.asynchronous import Hub

import t.skip
from celery.app.defaults import DEFAULTS
from celery.concurrency.asynpool import iterate_file_descriptors_safely
from celery.utils.collections import AttributeDict
from celery.utils.functional import noop
from celery.utils.objects import Bunch

try:
    from celery.concurrency import asynpool
    from celery.concurrency import prefork as mp
except ImportError:

    class _mp:
        RUN = 0x1

        class TaskPool:
            _pool = Mock()

            def __init__(self, *args, **kwargs):
                pass

            def start(self):
                pass

            def stop(self):
                pass

            def apply_async(self, *args, **kwargs):
                pass
    mp = _mp()
    asynpool = None


class MockResult:

    def __init__(self, value, pid):
        self.value = value
        self.pid = pid

    def worker_pids(self):
        return [self.pid]

    def get(self):
        return self.value


@patch('celery.platforms.set_mp_process_title')
class test_process_initializer:

    @staticmethod
    def Loader(*args, **kwargs):
        loader = Mock(*args, **kwargs)
        loader.conf = {}
        loader.override_backends = {}
        return loader

    @patch('celery.platforms.signals')
    def test_process_initializer(self, _signals, set_mp_process_title, restore_logging):
        from celery import signals
        from celery._state import _tls
        from celery.concurrency.prefork import WORKER_SIGIGNORE, WORKER_SIGRESET, process_initializer
        on_worker_process_init = Mock()
        signals.worker_process_init.connect(on_worker_process_init)

        with self.Celery(loader=self.Loader) as app:
            app.conf = AttributeDict(DEFAULTS)
            process_initializer(app, 'awesome.worker.com')
            _signals.ignore.assert_any_call(*WORKER_SIGIGNORE)
            _signals.reset.assert_any_call(*WORKER_SIGRESET)
            assert app.loader.init_worker.call_count
            on_worker_process_init.assert_called()
            assert _tls.current_app is app
            set_mp_process_title.assert_called_with(
                'celeryd', hostname='awesome.worker.com',
            )

            with patch('celery.app.trace.setup_worker_optimizations') as S:
                os.environ['FORKED_BY_MULTIPROCESSING'] = '1'
                try:
                    process_initializer(app, 'luke.worker.com')
                    S.assert_called_with(app, 'luke.worker.com')
                finally:
                    os.environ.pop('FORKED_BY_MULTIPROCESSING', None)

            os.environ['CELERY_LOG_FILE'] = 'worker%I.log'
            app.log.setup = Mock(name='log_setup')
            try:
                process_initializer(app, 'luke.worker.com')
            finally:
                os.environ.pop('CELERY_LOG_FILE', None)

    @patch('celery.platforms.set_pdeathsig')
    def test_pdeath_sig(self, _set_pdeathsig, set_mp_process_title, restore_logging):
        from celery import signals
        on_worker_process_init = Mock()
        signals.worker_process_init.connect(on_worker_process_init)
        from celery.concurrency.prefork import process_initializer

        with self.Celery(loader=self.Loader) as app:
            app.conf = AttributeDict(DEFAULTS)
            process_initializer(app, 'awesome.worker.com')
        _set_pdeathsig.assert_called_once_with('SIGKILL')


class test_process_destructor:

    @patch('celery.concurrency.prefork.signals')
    def test_process_destructor(self, signals):
        mp.process_destructor(13, -3)
        signals.worker_process_shutdown.send.assert_called_with(
            sender=None, pid=13, exitcode=-3,
        )


class MockPool:
    started = False
    closed = False
    joined = False
    terminated = False
    _state = None

    def __init__(self, *args, **kwargs):
        self.started = True
        self._timeout_handler = Mock()
        self._result_handler = Mock()
        self.maintain_pool = Mock()
        self._state = mp.RUN
        self._processes = kwargs.get('processes')
        self._proc_alive_timeout = kwargs.get('proc_alive_timeout')
        self._pool = [Bunch(pid=i, inqW_fd=1, outqR_fd=2)
                      for i in range(self._processes)]
        self._current_proc = cycle(range(self._processes))

    def close(self):
        self.closed = True
        self._state = 'CLOSE'

    def join(self):
        self.joined = True

    def terminate(self):
        self.terminated = True

    def terminate_job(self, *args, **kwargs):
        pass

    def restart(self, *args, **kwargs):
        pass

    def handle_result_event(self, *args, **kwargs):
        pass

    def flush(self):
        pass

    def grow(self, n=1):
        self._processes += n

    def shrink(self, n=1):
        self._processes -= n

    def apply_async(self, *args, **kwargs):
        pass

    def register_with_event_loop(self, loop):
        pass


class ExeMockPool(MockPool):

    def apply_async(self, target, args=(), kwargs={}, callback=noop):
        from threading import Timer
        res = target(*args, **kwargs)
        Timer(0.1, callback, (res,)).start()
        return MockResult(res, next(self._current_proc))


class TaskPool(mp.TaskPool):
    Pool = BlockingPool = MockPool


class ExeMockTaskPool(mp.TaskPool):
    Pool = BlockingPool = ExeMockPool


@t.skip.if_win32
class test_AsynPool:

    def setup_method(self):
        pytest.importorskip('multiprocessing')

    def test_gen_not_started(self):

        def gen():
            yield 1
            assert not asynpool.gen_not_started(g)
            yield 2
        g = gen()
        assert asynpool.gen_not_started(g)
        next(g)
        assert not asynpool.gen_not_started(g)
        list(g)
        assert not asynpool.gen_not_started(g)

        def gen2():
            yield 1
            raise RuntimeError('generator error')
        g = gen2()
        assert asynpool.gen_not_started(g)
        next(g)
        assert not asynpool.gen_not_started(g)
        with pytest.raises(RuntimeError):
            next(g)
        assert not asynpool.gen_not_started(g)

    @patch('select.select', create=True)
    def test_select(self, __select):
        ebadf = socket.error()
        ebadf.errno = errno.EBADF
        with patch('select.poll', create=True) as poller:
            poll = poller.return_value = Mock(name='poll.poll')
            poll.return_value = {3}, set(), 0
            assert asynpool._select({3}, poll=poll) == ({3}, set(), 0)

            poll.return_value = {3}, set(), 0
            assert asynpool._select({3}, None, {3}, poll=poll) == (
                {3}, set(), 0,
            )

            eintr = socket.error()
            eintr.errno = errno.EINTR
            poll.side_effect = eintr

            readers = {3}
            assert asynpool._select(readers, poll=poll) == (set(), set(), 1)
            assert 3 in readers

        with patch('select.poll', create=True) as poller:
            poll = poller.return_value = Mock(name='poll.poll')
            poll.side_effect = ebadf
            with patch('select.select') as selcheck:
                selcheck.side_effect = ebadf
                readers = {3}
                assert asynpool._select(readers, poll=poll) == (
                    set(), set(), 1,
                )
                assert 3 not in readers

        with patch('select.poll', create=True) as poller:
            poll = poller.return_value = Mock(name='poll.poll')
            poll.side_effect = MemoryError()
            with pytest.raises(MemoryError):
                asynpool._select({1}, poll=poll)

        with patch('select.poll', create=True) as poller:
            poll = poller.return_value = Mock(name='poll.poll')
            with patch('select.select') as selcheck:

                def se(*args):
                    selcheck.side_effect = MemoryError()
                    raise ebadf
                poll.side_effect = se
                with pytest.raises(MemoryError):
                    asynpool._select({3}, poll=poll)

        with patch('select.poll', create=True) as poller:
            poll = poller.return_value = Mock(name='poll.poll')
            with patch('select.select') as selcheck:

                def se2(*args):
                    selcheck.side_effect = socket.error()
                    selcheck.side_effect.errno = 1321
                    raise ebadf
                poll.side_effect = se2
                with pytest.raises(socket.error):
                    asynpool._select({3}, poll=poll)

        with patch('select.poll', create=True) as poller:
            poll = poller.return_value = Mock(name='poll.poll')

            poll.side_effect = socket.error()
            poll.side_effect.errno = 34134
            with pytest.raises(socket.error):
                asynpool._select({3}, poll=poll)

    def test_select_unpatched(self):
        with tempfile.TemporaryFile('w') as f:
            _, writeable, _ = asynpool._select(writers={f, }, err={f, })
            assert f.fileno() in writeable

        with tempfile.TemporaryFile('r') as f:
            readable, _, _ = asynpool._select(readers={f, }, err={f, })
            assert f.fileno() in readable

    def test_promise(self):
        fun = Mock()
        x = asynpool.promise(fun, (1,), {'foo': 1})
        x()
        assert x.ready
        fun.assert_called_with(1, foo=1)

    def test_Worker(self):
        w = asynpool.Worker(Mock(), Mock())
        w.on_loop_start(1234)
        w.outq.put.assert_called_with((asynpool.WORKER_UP, (1234,)))

    def test_iterate_file_descriptors_safely_source_data_list(self):
        # Given: a list of integers that could be file descriptors
        fd_iter = [1, 2, 3, 4, 5]

        # Given: a mock hub method that does nothing to call
        def _fake_hub(*args, **kwargs):
            raise OSError

        # When Calling the helper to iterate_file_descriptors_safely
        iterate_file_descriptors_safely(
            fd_iter, fd_iter, _fake_hub,
            "arg1", "arg2", kw1="kw1", kw2="kw2",
        )

        # Then: all items were removed from the managed data source
        assert fd_iter == [], "Expected all items removed from managed list"

    def test_iterate_file_descriptors_safely_source_data_set(self):
        # Given: a list of integers that could be file descriptors
        fd_iter = {1, 2, 3, 4, 5}

        # Given: a mock hub method that does nothing to call
        def _fake_hub(*args, **kwargs):
            raise OSError

        # When Calling the helper to iterate_file_descriptors_safely
        iterate_file_descriptors_safely(
            fd_iter, fd_iter, _fake_hub,
            "arg1", "arg2", kw1="kw1", kw2="kw2",
        )

        # Then: all items were removed from the managed data source
        assert fd_iter == set(), "Expected all items removed from managed set"

    def test_iterate_file_descriptors_safely_source_data_dict(self):
        # Given: a list of integers that could be file descriptors
        fd_iter = {1: 1, 2: 2, 3: 3, 4: 4, 5: 5}

        # Given: a mock hub method that does nothing to call
        def _fake_hub(*args, **kwargs):
            raise OSError

        # When Calling the helper to iterate_file_descriptors_safely
        iterate_file_descriptors_safely(
            fd_iter, fd_iter, _fake_hub,
            "arg1", "arg2", kw1="kw1", kw2="kw2",
        )

        # Then: all items were removed from the managed data source
        assert fd_iter == {}, "Expected all items removed from managed dict"

    def _get_hub(self):
        hub = Hub()
        hub.readers = {}
        hub.writers = {}
        hub.timer = Mock(name='hub.timer')
        hub.timer._queue = [Mock()]
        hub.fire_timers = Mock(name='hub.fire_timers')
        hub.fire_timers.return_value = 1.7
        hub.poller = Mock(name='hub.poller')
        hub.close = Mock(name='hub.close()')
        return hub

    @t.skip.if_pypy
    def test_schedule_writes_hub_remove_writer_ready_fd_not_in_all_inqueues(self):
        pool = asynpool.AsynPool(threads=False)
        hub = self._get_hub()

        writer = Mock(name='writer')
        reader = Mock(name='reader')

        # add 2 fake fds with the same id
        hub.add_reader(6, reader, 6)
        hub.add_writer(6, writer, 6)
        pool._all_inqueues.clear()
        pool._create_write_handlers(hub)

        # check schedule_writes write fds remove not remove the reader one from the hub.
        hub.consolidate_callback(ready_fds=[6])
        assert 6 in hub.readers
        assert 6 not in hub.writers

    @t.skip.if_pypy
    def test_schedule_writes_hub_remove_writers_from_active_writers_when_get_index_error(self):
        pool = asynpool.AsynPool(threads=False)
        hub = self._get_hub()

        writer = Mock(name='writer')
        reader = Mock(name='reader')

        # add 3 fake fds with the same id to reader and writer
        hub.add_reader(6, reader, 6)
        hub.add_reader(8, reader, 8)
        hub.add_reader(9, reader, 9)
        hub.add_writer(6, writer, 6)
        hub.add_writer(8, writer, 8)
        hub.add_writer(9, writer, 9)

        # add fake fd to pool _all_inqueues to make sure we try to read from outbound_buffer
        # set active_writes to 6 to make sure we remove all write fds except 6
        pool._active_writes = {6}
        pool._all_inqueues = {2, 6, 8, 9}

        pool._create_write_handlers(hub)

        # clear outbound_buffer to get IndexError when trying to pop any message
        # in this case all active_writers fds will be removed from the hub
        pool.outbound_buffer.clear()

        hub.consolidate_callback(ready_fds=[2])
        if {6, 8, 9} <= hub.readers.keys() and not {8, 9} <= hub.writers.keys():
            assert True
        else:
            assert False

        assert 6 in hub.writers

    @t.skip.if_pypy
    def test_schedule_writes_hub_remove_fd_only_from_writers_when_write_job_is_done(self):
        pool = asynpool.AsynPool(threads=False)
        hub = self._get_hub()

        writer = Mock(name='writer')
        reader = Mock(name='reader')

        # add one writer and one reader with the same fd
        hub.add_writer(2, writer, 2)
        hub.add_reader(2, reader, 2)
        assert 2 in hub.writers

        # For test purposes to reach _write_job in schedule writes
        pool._all_inqueues = {2}
        worker = Mock("worker")
        # this lambda need to return a number higher than 4
        # to pass the while loop in _write_job function and to reach the hub.remove_writer
        worker.send_job_offset = lambda header, HW: 5

        pool._fileno_to_inq[2] = worker
        pool._create_write_handlers(hub)

        result = ApplyResult({}, lambda x: True)
        result._payload = [None, None, -1]
        pool.outbound_buffer.appendleft(result)

        hub.consolidate_callback(ready_fds=[2])
        assert 2 not in hub.writers
        assert 2 in hub.readers

    @t.skip.if_pypy
    def test_register_with_event_loop__no_on_tick_dupes(self):
        """Ensure AsynPool's register_with_event_loop only registers
        on_poll_start in the event loop the first time it's called. This
        prevents a leak when the Consumer is restarted.
        """
        pool = asynpool.AsynPool(threads=False)
        hub = Mock(name='hub')
        pool.register_with_event_loop(hub)
        pool.register_with_event_loop(hub)
        hub.on_tick.add.assert_called_once()

    @t.skip.if_pypy
    @patch('billiard.pool.Pool._create_worker_process')
    def test_before_create_process_signal(self, create_process):
        from celery import signals
        on_worker_before_create_process = Mock()
        signals.worker_before_create_process.connect(on_worker_before_create_process)
        pool = asynpool.AsynPool(processes=1, threads=False)
        create_process.assert_called_once_with(0)
        on_worker_before_create_process.assert_any_call(
            signal=signals.worker_before_create_process,
            sender=pool,
        )

    def test_untrack_child_process_without_sentinel_poll(self):
        """_untrack_child_process must not raise when proc lacks _sentinel_poll.

        Race condition during cold shutdown can cause _untrack_child_process to
        be called with a process that never had _sentinel_poll set or had it
        cleared. Use getattr for safe access.
        """
        pytest.importorskip('multiprocessing')
        pool = asynpool.AsynPool(processes=1, threads=False)
        hub = Mock(name='hub')
        proc = object()  # No _sentinel_poll attribute
        pool._untrack_child_process(proc, hub)  # Should not raise AttributeError
        hub.remove.assert_not_called()

    def test_untrack_child_process_with_sentinel_poll(self):
        """_untrack_child_process cleans up when proc has _sentinel_poll set."""
        pytest.importorskip('multiprocessing')
        pool = asynpool.AsynPool(processes=1, threads=False)
        hub = Mock(name='hub')
        fd = os.open(os.devnull, os.O_RDONLY)
        proc = Mock(_sentinel_poll=fd)
        pool._untrack_child_process(proc, hub)
        hub.remove.assert_called_once_with(fd)
        assert proc._sentinel_poll is None

    @t.skip.if_pypy
    def test_flush_no_synack_discards_unaccepted_jobs(self):
        """flush() should discard unaccepted jobs when synack is disabled.

        Previously, flush() only handled the synack case. Without synack,
        unaccepted jobs were never cleaned from the cache, leading to stale
        entries.
        """
        pool = asynpool.AsynPool(processes=1, synack=False, threads=False)
        pool._state = asynpool.RUN
        pool.maintain_pool = Mock(name='maintain_pool')

        job1 = Mock(name='job1')
        job1._accepted = False
        job1._writer.return_value = None
        job2 = Mock(name='job2')
        job2._accepted = True
        job2._writer.return_value = None

        pool._cache = {1: job1, 2: job2}
        pool.outbound_buffer.clear()
        pool._active_writers.clear()

        pool.flush()

        job1.discard.assert_called_once()
        job2.discard.assert_not_called()

    @t.skip.if_pypy
    def test_flush_synack_cancels_unaccepted_jobs(self):
        """flush() should call _cancel() on unaccepted jobs when synack is enabled."""
        pool = asynpool.AsynPool(processes=1, synack=True, threads=False)
        pool._state = asynpool.RUN
        pool.maintain_pool = Mock(name='maintain_pool')

        job1 = Mock(name='job1')
        job1._accepted = False
        job1._writer.return_value = None
        job2 = Mock(name='job2')
        job2._accepted = True
        job2._writer.return_value = None

        pool._cache = {1: job1, 2: job2}
        pool.outbound_buffer.clear()
        pool._active_writers.clear()

        pool.flush()

        job1._cancel.assert_called_once()
        job1.discard.assert_not_called()
        job2._cancel.assert_not_called()

    @t.skip.if_pypy
    @patch('billiard.pool.Pool._create_worker_process')
    def test_flush_dead_process_discards_active_writer(self, _create_worker_process):
        """flush() must discard generator from _active_writers when process is dead.

        Previously, when a process was dead, the generator was never removed
        from _active_writers, causing an infinite loop in the while loop.
        """
        pool = asynpool.AsynPool(processes=1, synack=False, threads=False)
        pool._state = asynpool.RUN
        pool.maintain_pool = Mock(name='maintain_pool')

        # Create a mock generator (already started, so not gen_not_started)
        gen = Mock(name='gen')
        gen.__name__ = '_write_job'
        # Simulate a started generator
        with patch.object(asynpool, 'gen_not_started', return_value=False):
            proc = Mock(name='proc')
            proc._is_alive.return_value = False  # Process is dead

            job = Mock(name='job')
            job._accepted = True
            job._write_to = proc
            job._writer.return_value = gen

            pool._cache = {1: job}
            pool._active_writers = {gen}
            pool.outbound_buffer.clear()

            pool.flush()

        # Generator should have been removed from active_writers
        assert gen not in pool._active_writers
        # Job should have been discarded since process is dead
        job.discard.assert_called()

    @t.skip.if_pypy
    @patch('billiard.pool.Pool._create_worker_process')
    def test_flush_alive_process_flushes_writer(self, _create_worker_process):
        """flush() should call _flush_writer when process is still alive."""
        pool = asynpool.AsynPool(processes=1, synack=False, threads=False)
        pool._state = asynpool.RUN
        pool.maintain_pool = Mock(name='maintain_pool')

        gen = Mock(name='gen')
        gen.__name__ = '_write_job'

        with patch.object(asynpool, 'gen_not_started', return_value=False):
            proc = Mock(name='proc')
            proc._is_alive.return_value = True

            job = Mock(name='job')
            job._accepted = True
            job._write_to = proc
            job._writer.return_value = gen

            pool._cache = {1: job}
            pool._active_writers = {gen}
            pool.outbound_buffer.clear()

            with patch.object(pool, '_flush_writer') as mock_flush:
                # _flush_writer removes from _active_writers in its finally
                def side_effect(p, g):
                    pool._active_writers.discard(g)
                mock_flush.side_effect = side_effect

                pool.flush()

            mock_flush.assert_called_once_with(proc, gen)

    @t.skip.if_pypy
    @patch('billiard.pool.Pool._create_worker_process')
    def test_flush_write_ack_coroutine_is_advanced_not_dropped(self, _create_worker_process):
        """flush() must not silently drop _write_ack generators (synack mode).

        _write_ack coroutines are added to _active_writers by send_ack() but
        are NOT mapped in owned_by (which is built from _cache job writers only).
        Dropping them mid-write leaves a partially-written ack on the synq pipe
        and hangs the worker process waiting for the ack that never arrives.
        flush() must advance them to completion instead.
        """
        pool = asynpool.AsynPool(processes=1, synack=True, threads=False)
        pool._state = asynpool.RUN
        pool.maintain_pool = Mock(name='maintain_pool')

        # Simulate a _write_ack generator (name != '_write_job', not in owned_by)
        ack_gen = Mock(name='ack_gen')
        ack_gen.__name__ = '_write_ack'
        # First call to next() returns normally (still writing),
        # second raises StopIteration (write complete).
        ack_gen.__next__ = Mock(side_effect=[None, StopIteration()])

        pool._cache = {}
        pool._active_writers = {ack_gen}
        pool.outbound_buffer.clear()

        pool.flush()

        # Generator should have been advanced (not just silently discarded)
        assert ack_gen.__next__.called
        # And removed once it signalled completion
        assert ack_gen not in pool._active_writers

    @t.skip.if_pypy
    @patch('billiard.pool.Pool._create_worker_process')
    def test_flush_write_ack_coroutine_handles_oserror(self, _create_worker_process):
        """flush() should discard the coroutine if OSError is raised during next()."""
        pool = asynpool.AsynPool(processes=1, synack=True, threads=False)
        pool._state = asynpool.RUN
        pool.maintain_pool = Mock(name='maintain_pool')

        ack_gen = MagicMock(name='ack_gen')
        ack_gen.__name__ = '_write_ack'
        ack_gen.__next__.side_effect = OSError()

        pool._cache = {}
        pool._active_writers = {ack_gen}
        pool.outbound_buffer.clear()
        pool.flush()

        assert ack_gen not in pool._active_writers

    @t.skip.if_pypy
    @patch('billiard.pool.Pool._create_worker_process')
    def test_flush_write_ack_coroutine_handles_eoferror(self, _create_worker_process):
        """flush() should discard the coroutine if EOFError is raised during next()."""
        pool = asynpool.AsynPool(processes=1, synack=True, threads=False)
        pool._state = asynpool.RUN
        pool.maintain_pool = Mock(name='maintain_pool')

        ack_gen = MagicMock(name='ack_gen')
        ack_gen.__name__ = '_write_ack'
        ack_gen.__next__.side_effect = EOFError()

        pool._cache = {}
        pool._active_writers = {ack_gen}
        pool.outbound_buffer.clear()
        pool.flush()

        assert ack_gen not in pool._active_writers

    @t.skip.if_pypy
    @patch('billiard.pool.Pool._create_worker_process')
    def test_flush_not_started_write_job_is_discarded(self, _create_worker_process):
        """flush() should discard a _write_job generator that has not started yet.

        When gen_not_started() returns True the job has not been written to the
        pipe at all, so it is safe to discard it and let the broker redeliver.
        """
        pool = asynpool.AsynPool(processes=1, synack=False, threads=False)
        pool._state = asynpool.RUN
        pool.maintain_pool = Mock(name='maintain_pool')

        gen = Mock(name='gen')
        gen.__name__ = '_write_job'

        with patch.object(asynpool, 'gen_not_started', return_value=True):
            job = Mock(name='job')
            job._accepted = True
            job._writer.return_value = gen

            pool._cache = {1: job}
            pool._active_writers = {gen}
            pool.outbound_buffer.clear()

            pool.flush()

        job.discard.assert_called_once()
        assert gen not in pool._active_writers

    def test_process_result(self):
        x = asynpool.ResultHandler(
            Mock(), Mock(), {}, Mock(),
            Mock(), Mock(), Mock(), Mock(),
            fileno_to_outq={},
            on_process_alive=Mock(),
            on_job_ready=Mock(),
        )
        assert x
        hub = Mock(name='hub')
        recv = x._recv_message = Mock(name='recv_message')
        recv.return_value = iter([])
        x.on_state_change = Mock()
        x.register_with_event_loop(hub)
        proc = x.fileno_to_outq[3] = Mock()
        reader = proc.outq._reader
        reader.poll.return_value = False
        x.handle_event(6)  # KeyError
        x.handle_event(3)
        x._recv_message.assert_called_with(
            hub.add_reader, 3, x.on_state_change,
        )


class test_TaskPool:

    def test_start(self):
        pool = TaskPool(10)
        pool.start()
        assert pool._pool.started
        assert pool._pool._state == asynpool.RUN

        _pool = pool._pool
        pool.stop()
        assert _pool.closed
        assert _pool.joined
        pool.stop()

        pool.start()
        _pool = pool._pool
        pool.terminate()
        pool.terminate()
        assert _pool.terminated

    def test_restart(self):
        pool = TaskPool(10)
        pool._pool = Mock(name='pool')
        pool.restart()
        pool._pool.restart.assert_called_with()
        pool._pool.apply_async.assert_called_with(mp.noop)

    def test_did_start_ok(self):
        pool = TaskPool(10)
        pool._pool = Mock(name='pool')
        assert pool.did_start_ok() is pool._pool.did_start_ok()

    def test_register_with_event_loop(self):
        pool = TaskPool(10)
        pool._pool = Mock(name='pool')
        loop = Mock(name='loop')
        pool.register_with_event_loop(loop)
        pool._pool.register_with_event_loop.assert_called_with(loop)

    def test_on_close(self):
        pool = TaskPool(10)
        pool._pool = Mock(name='pool')
        pool._pool._state = mp.RUN
        pool.on_close()
        pool._pool.close.assert_called_with()

    def test_on_close__pool_not_running(self):
        pool = TaskPool(10)
        pool._pool = Mock(name='pool')
        pool._pool._state = mp.CLOSE
        pool.on_close()
        pool._pool.close.assert_not_called()

    @patch('celery.concurrency.prefork.get_event_loop')
    @patch('celery.concurrency.prefork.threading.Thread')
    def test_on_stop_with_hub_fires_timers(self, mock_thread, mock_get_event_loop):
        pool = TaskPool(10)
        mock_pool = Mock(name='pool')
        mock_pool._state = mp.RUN
        pool._pool = mock_pool

        mock_hub = Mock(name='hub')
        mock_get_event_loop.return_value = mock_hub
        mock_timer_thread = Mock(name='timer_thread')
        mock_thread.return_value = mock_timer_thread

        pool.on_stop()

        mock_pool.close.assert_called_with()
        mock_pool.join.assert_called_with()
        mock_get_event_loop.assert_called_once()
        mock_thread.assert_called_once()
        assert mock_thread.call_args[1]['daemon'] is True
        mock_timer_thread.start.assert_called_once()
        mock_timer_thread.join.assert_called_once_with(timeout=1.0)

    @patch('celery.concurrency.prefork.get_event_loop')
    @patch('celery.concurrency.prefork.threading.Thread')
    @patch('celery.concurrency.prefork.threading.Event')
    def test_on_stop_timer_thread_handles_exceptions(
        self,
        mock_event_class,
        mock_thread,
        mock_get_event_loop,
    ):
        pool = TaskPool(10)
        mock_pool = Mock(name='pool')
        mock_pool._state = mp.RUN
        pool._pool = mock_pool

        mock_hub = Mock(name='hub')
        mock_hub.fire_timers.side_effect = [Exception("Hub error"), None]
        mock_get_event_loop.return_value = mock_hub

        mock_shutdown_event = Mock(name='shutdown_event')
        # Simulate two loop iterations and then shutdown
        mock_shutdown_event.is_set.side_effect = [False, False, True]
        mock_event_class.return_value = mock_shutdown_event

        thread_target = None

        def capture_thread(*args, **kwargs):
            nonlocal thread_target
            thread_target = kwargs['target']
            mock_timer_thread = Mock(name='timer_thread')
            return mock_timer_thread

        mock_thread.side_effect = capture_thread

        pool.on_stop()

        with patch('celery.concurrency.prefork.time.sleep'):
            thread_target()

        # Should match number of loop iterations allowed by mock_shutdown_event.is_set.side_effect
        assert mock_hub.fire_timers.call_count == 2

    @patch('celery.concurrency.prefork.get_event_loop')
    def test_on_stop_no_hub(self, mock_get_event_loop):
        pool = TaskPool(10)
        mock_pool = Mock(name='pool')
        mock_pool._state = mp.RUN
        pool._pool = mock_pool

        mock_get_event_loop.return_value = None

        pool.on_stop()

        mock_pool.close.assert_called_with()
        mock_pool.join.assert_called_with()
        mock_get_event_loop.assert_called_once()

    def test_apply_async(self):
        pool = TaskPool(10)
        pool.start()
        pool.apply_async(lambda x: x, (2,), {})

    def test_grow_shrink(self):
        pool = TaskPool(10)
        pool.start()
        assert pool._pool._processes == 10
        pool.grow()
        assert pool._pool._processes == 11
        pool.shrink(2)
        assert pool._pool._processes == 9

    def test_info(self):
        pool = TaskPool(10)
        procs = [Bunch(pid=i) for i in range(pool.limit)]

        class _Pool:
            _pool = procs
            _maxtasksperchild = None
            timeout = 10
            soft_timeout = 5

            def human_write_stats(self, *args, **kwargs):
                return {}
        pool._pool = _Pool()
        info = pool.info
        assert info['max-concurrency'] == pool.limit
        assert info['max-tasks-per-child'] == 'N/A'
        assert info['timeouts'] == (5, 10)

    def test_num_processes(self):
        pool = TaskPool(7)
        pool.start()
        assert pool.num_processes == 7

    @patch('billiard.forking_enable')
    def test_on_start_proc_alive_timeout_default(self, __forking_enable):
        app = Mock(conf=AttributeDict(DEFAULTS))
        pool = TaskPool(4, app=app)
        pool.on_start()
        assert pool._pool._proc_alive_timeout == 4.0

    @patch('billiard.forking_enable')
    def test_on_start_proc_alive_timeout_custom(self, __forking_enable):
        app = Mock(conf=AttributeDict(DEFAULTS))
        app.conf.worker_proc_alive_timeout = 8.0
        pool = TaskPool(4, app=app)
        pool.on_start()
        assert pool._pool._proc_alive_timeout == 8.0
