Update tooling. something could be tweaked/ removed from this PR
updated dockerfile to add python 3.12. Update tooling. something could be tweaked/ removed from this PR
import sys
from time import monotonic
from unittest.mock import Mock, patch

import pytest

from celery.concurrency.base import BasePool
from celery.utils.objects import Bunch
from celery.worker import autoscale, state


class MockPool(BasePool):

    shrink_raises_exception = False
    shrink_raises_ValueError = False

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._pool = Bunch(_processes=self.limit)

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

    def shrink(self, n=1):
        if self.shrink_raises_exception:
            raise KeyError('foo')
        if self.shrink_raises_ValueError:
            raise ValueError('foo')
        self._pool._processes -= n

    @property
    def num_processes(self):
        return self._pool._processes


class test_WorkerComponent:

    def test_register_with_event_loop(self):
        parent = Mock(name='parent')
        parent.autoscale = True
        parent.consumer.on_task_message = set()
        w = autoscale.WorkerComponent(parent)
        assert parent.autoscaler is None
        assert w.enabled

        hub = Mock(name='hub')
        w.create(parent)
        w.register_with_event_loop(parent, hub)
        assert (parent.autoscaler.maybe_scale in
                parent.consumer.on_task_message)
        hub.call_repeatedly.assert_called_with(
            parent.autoscaler.keepalive, parent.autoscaler.maybe_scale,
        )

        parent.hub = hub
        hub.on_init = []
        w.instantiate = Mock()
        w.register_with_event_loop(parent, Mock(name='loop'))
        assert parent.consumer.on_task_message

    def test_info_without_event_loop(self):
        parent = Mock(name='parent')
        parent.autoscale = True
        parent.max_concurrency = '10'
        parent.min_concurrency = '2'
        parent.use_eventloop = False
        w = autoscale.WorkerComponent(parent)
        w.create(parent)
        info = w.info(parent)
        assert 'autoscaler' in info
        assert parent.autoscaler_cls().info.called


class test_Autoscaler:

    def setup_method(self):
        self.pool = MockPool(3)

    def test_stop(self):

        class Scaler(autoscale.Autoscaler):
            alive = True
            joined = False

            def is_alive(self):
                return self.alive

            def join(self, timeout=None):
                self.joined = True

        worker = Mock(name='worker')
        x = Scaler(self.pool, 10, 3, worker=worker)
        # Don't allow thread joining or event waiting to block the test
        with patch("threading.Thread.join"), patch("threading.Event.wait"):
            x.stop()
        assert x.joined
        x.joined = False
        x.alive = False
        with patch("threading.Thread.join"), patch("threading.Event.wait"):
            x.stop()
        assert not x.joined

    @pytest.mark.sleepdeprived_patched_module(autoscale)
    def test_body(self, sleepdeprived):
        worker = Mock(name='worker')
        x = autoscale.Autoscaler(self.pool, 10, 3, worker=worker)
        x.body()
        assert x.pool.num_processes == 3
        _keep = [Mock(name=f'req{i}') for i in range(20)]
        [state.task_reserved(m) for m in _keep]
        x.body()
        x.body()
        assert x.pool.num_processes == 10
        state.reserved_requests.clear()
        x.body()
        assert x.pool.num_processes == 10
        x._last_scale_up = monotonic() - 10000
        x.body()
        assert x.pool.num_processes == 3

    def test_run(self):

        class Scaler(autoscale.Autoscaler):
            scale_called = False

            def body(self):
                self.scale_called = True
                getattr(self, "_bgThread__is_shutdown").set()

        worker = Mock(name='worker')
        x = Scaler(self.pool, 10, 3, worker=worker)
        x.run()
        assert getattr(x, "_bgThread__is_shutdown").is_set()
        assert getattr(x, "_bgThread__is_stopped").is_set()
        assert x.scale_called

    def test_shrink_raises_exception(self):
        worker = Mock(name='worker')
        x = autoscale.Autoscaler(self.pool, 10, 3, worker=worker)
        x.scale_up(3)
        x.pool.shrink_raises_exception = True
        x._shrink(1)

    @patch('celery.worker.autoscale.debug')
    def test_shrink_raises_ValueError(self, debug):
        worker = Mock(name='worker')
        x = autoscale.Autoscaler(self.pool, 10, 3, worker=worker)
        x.scale_up(3)
        x._last_scale_up = monotonic() - 10000
        x.pool.shrink_raises_ValueError = True
        x.scale_down(1)
        assert debug.call_count

    def test_update(self):
        worker = Mock(name='worker')
        x = autoscale.Autoscaler(self.pool, 10, 3, worker=worker)
        x.worker.consumer.prefetch_multiplier = 1
        x.keepalive = -1
        assert x.processes == 3
        x.scale_up(5)
        x.update(7, None)
        assert x.processes == 7
        assert x.max_concurrency == 7
        x.scale_down(4)
        x.update(None, 6)
        assert x.processes == 6
        assert x.min_concurrency == 6

        x.update(max=300, min=10)
        x.update(max=300, min=2)
        x.update(max=None, min=None)

    def test_prefetch_count_on_updates(self):
        worker = Mock(name='worker')
        x = autoscale.Autoscaler(self.pool, 10, 3, worker=worker)
        x.worker.consumer.prefetch_multiplier = 1
        x.update(5, None)
        worker.consumer._update_prefetch_count.assert_called_with(-5)
        x.update(15, 7)
        worker.consumer._update_prefetch_count.assert_called_with(10)

    def test_prefetch_count_on_updates_prefetch_multiplier_gt_one(self):
        worker = Mock(name='worker')
        x = autoscale.Autoscaler(self.pool, 10, 3, worker=worker)
        x.worker.consumer.prefetch_multiplier = 4
        x.update(5, None)
        worker.consumer._update_prefetch_count.assert_called_with(-5)
        x.update(15, 7)
        worker.consumer._update_prefetch_count.assert_called_with(10)

    def test_info(self):
        worker = Mock(name='worker')
        x = autoscale.Autoscaler(self.pool, 10, 3, worker=worker)
        info = x.info()
        assert info['max'] == 10
        assert info['min'] == 3
        assert info['current'] == 3

    @patch('os._exit')
    def test_thread_crash(self, _exit):

        class _Autoscaler(autoscale.Autoscaler):

            def body(self):
                getattr(self, "_bgThread__is_shutdown").set()
                raise OSError('foo')
        worker = Mock(name='worker')
        x = _Autoscaler(self.pool, 10, 3, worker=worker)

        stderr = Mock()
        p, sys.stderr = sys.stderr, stderr
        try:
            x.run()
        finally:
            sys.stderr = p
        _exit.assert_called_with(1)
        stderr.write.assert_called()

    @pytest.mark.sleepdeprived_patched_module(autoscale)
    def test_no_negative_scale(self, sleepdeprived):
        total_num_processes = []
        worker = Mock(name='worker')
        x = autoscale.Autoscaler(self.pool, 10, 3, worker=worker)
        x.body()  # the body func scales up or down

        _keep = [Mock(name=f'req{i}') for i in range(35)]
        for req in _keep:
            state.task_reserved(req)
            x.body()
            total_num_processes.append(self.pool.num_processes)

        for req in _keep:
            state.task_ready(req)
            x.body()
            total_num_processes.append(self.pool.num_processes)

        assert all(x.min_concurrency <= i <= x.max_concurrency
                   for i in total_num_processes)

import sys
from time import monotonic
from unittest.mock import Mock, patch

import pytest

from celery.concurrency.base import BasePool
from celery.utils.objects import Bunch
from celery.worker import autoscale, state


class MockPool(BasePool):

    shrink_raises_exception = False
    shrink_raises_ValueError = False

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._pool = Bunch(_processes=self.limit)

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

    def shrink(self, n=1):
        if self.shrink_raises_exception:
            raise KeyError('foo')
        if self.shrink_raises_ValueError:
            raise ValueError('foo')
        self._pool._processes -= n

    @property
    def num_processes(self):
        return self._pool._processes


class test_WorkerComponent:

    def test_register_with_event_loop(self):
        parent = Mock(name='parent')
        parent.autoscale = True
        parent.consumer.on_task_message = set()
        w = autoscale.WorkerComponent(parent)
        assert parent.autoscaler is None
        assert w.enabled

        hub = Mock(name='hub')
        w.create(parent)
        w.register_with_event_loop(parent, hub)
        assert (parent.autoscaler.maybe_scale in
                parent.consumer.on_task_message)
        hub.call_repeatedly.assert_called_with(
            parent.autoscaler.keepalive, parent.autoscaler.maybe_scale,
        )

        parent.hub = hub
        hub.on_init = []
        w.instantiate = Mock()
        w.register_with_event_loop(parent, Mock(name='loop'))
        assert parent.consumer.on_task_message

    def test_info_without_event_loop(self):
        parent = Mock(name='parent')
        parent.autoscale = True
        parent.max_concurrency = '10'
        parent.min_concurrency = '2'
        parent.use_eventloop = False
        w = autoscale.WorkerComponent(parent)
        w.create(parent)
        info = w.info(parent)
        assert 'autoscaler' in info
        assert parent.autoscaler_cls().info.called


class test_Autoscaler:

    def setup_method(self):
        self.pool = MockPool(3)

    def test_stop(self):

        class Scaler(autoscale.Autoscaler):
            alive = True
            joined = False

            def is_alive(self):
                return self.alive

            def join(self, timeout=None):
                self.joined = True

        worker = Mock(name='worker')
        x = Scaler(self.pool, 10, 3, worker=worker)
        # Don't allow thread joining or event waiting to block the test
        with patch("threading.Thread.join"), patch("threading.Event.wait"):
            x.stop()
        assert x.joined
        x.joined = False
        x.alive = False
        with patch("threading.Thread.join"), patch("threading.Event.wait"):
            x.stop()
        assert not x.joined

    @pytest.mark.sleepdeprived_patched_module(autoscale)
    def test_body(self, sleepdeprived):
        worker = Mock(name='worker')
        x = autoscale.Autoscaler(self.pool, 10, 3, worker=worker)
        x.body()
        assert x.pool.num_processes == 3
        _keep = [Mock(name=f'req{i}') for i in range(20)]
        [state.task_reserved(m) for m in _keep]
        x.body()
        x.body()
        assert x.pool.num_processes == 10
        state.reserved_requests.clear()
        x.body()
        assert x.pool.num_processes == 10
        x._last_scale_up = monotonic() - 10000
        x.body()
        assert x.pool.num_processes == 3

    def test_run(self):

        class Scaler(autoscale.Autoscaler):
            scale_called = False

            def body(self):
                self.scale_called = True
                getattr(self, "_bgThread__is_shutdown").set()

        worker = Mock(name='worker')
        x = Scaler(self.pool, 10, 3, worker=worker)
        x.run()
        assert getattr(x, "_bgThread__is_shutdown").is_set()
        assert getattr(x, "_bgThread__is_stopped").is_set()
        assert x.scale_called

    def test_shrink_raises_exception(self):
        worker = Mock(name='worker')
        x = autoscale.Autoscaler(self.pool, 10, 3, worker=worker)
        x.scale_up(3)
        x.pool.shrink_raises_exception = True
        x._shrink(1)

    @patch('celery.worker.autoscale.debug')
    def test_shrink_raises_ValueError(self, debug):
        worker = Mock(name='worker')
        x = autoscale.Autoscaler(self.pool, 10, 3, worker=worker)
        x.scale_up(3)
        x._last_scale_up = monotonic() - 10000
        x.pool.shrink_raises_ValueError = True
        x.scale_down(1)
        assert debug.call_count

    def test_update(self):
        worker = Mock(name='worker')
        x = autoscale.Autoscaler(self.pool, 10, 3, worker=worker)
        x.worker.consumer.prefetch_multiplier = 1
        x.keepalive = -1
        assert x.processes == 3
        x.scale_up(5)
        x.update(7, None)
        assert x.processes == 7
        assert x.max_concurrency == 7
        x.scale_down(4)
        x.update(None, 6)
        assert x.processes == 6
        assert x.min_concurrency == 6

        x.update(max=300, min=10)
        x.update(max=300, min=2)
        x.update(max=None, min=None)

    def test_prefetch_count_on_updates(self):
        worker = Mock(name='worker')
        x = autoscale.Autoscaler(self.pool, 10, 3, worker=worker)
        x.worker.consumer.prefetch_multiplier = 1
        x.update(5, None)
        worker.consumer._update_prefetch_count.assert_called_with(-5)
        x.update(15, 7)
        worker.consumer._update_prefetch_count.assert_called_with(10)

    def test_prefetch_count_on_updates_prefetch_multiplier_gt_one(self):
        worker = Mock(name='worker')
        x = autoscale.Autoscaler(self.pool, 10, 3, worker=worker)
        x.worker.consumer.prefetch_multiplier = 4
        x.update(5, None)
        worker.consumer._update_prefetch_count.assert_called_with(-5)
        x.update(15, 7)
        worker.consumer._update_prefetch_count.assert_called_with(10)

    def test_info(self):
        worker = Mock(name='worker')
        x = autoscale.Autoscaler(self.pool, 10, 3, worker=worker)
        info = x.info()
        assert info['max'] == 10
        assert info['min'] == 3
        assert info['current'] == 3

    @patch('os._exit')
    def test_thread_crash(self, _exit):

        class _Autoscaler(autoscale.Autoscaler):

            def body(self):
                getattr(self, "_bgThread__is_shutdown").set()
                raise OSError('foo')
        worker = Mock(name='worker')
        x = _Autoscaler(self.pool, 10, 3, worker=worker)

        stderr = Mock()
        p, sys.stderr = sys.stderr, stderr
        try:
            x.run()
        finally:
            sys.stderr = p
        _exit.assert_called_with(1)
        stderr.write.assert_called()

    @pytest.mark.sleepdeprived_patched_module(autoscale)
    def test_no_negative_scale(self, sleepdeprived):
        total_num_processes = []
        worker = Mock(name='worker')
        x = autoscale.Autoscaler(self.pool, 10, 3, worker=worker)
        x.body()  # the body func scales up or down

        _keep = [Mock(name=f'req{i}') for i in range(35)]
        for req in _keep:
            state.task_reserved(req)
            x.body()
            total_num_processes.append(self.pool.num_processes)

        for req in _keep:
            state.task_ready(req)
            x.body()
            total_num_processes.append(self.pool.num_processes)

        assert all(x.min_concurrency <= i <= x.max_concurrency
                   for i in total_num_processes)
