def test_acquire_release(self):
x = LaxBoundedSemaphore(2)
c1 = Mock()
x.acquire(c1, 1)
assert x.value == 1
c1.assert_called_with(1)
c2 = Mock()
x.acquire(c2, 2)
assert x.value == 0
c2.assert_called_with(2)
c3 = Mock()
x.acquire(c3, 3)
assert x.value == 0
c3.assert_not_called()
x.release()
assert x.value == 0
x.release()
assert x.value == 1
x.release()
assert x.value == 2
c3.assert_called_with(3)
def test_pool_restart(self):
consumer = Consumer(self.app)
consumer.controller = _WC(app=self.app)
consumer.controller.consumer = consumer
consumer.controller.pool.restart = Mock()
consumer.reset_rate_limits = Mock(name='reset_rate_limits()')
consumer.update_strategies = Mock(name='update_strategies()')
consumer.event_dispatcher = Mock(name='evd')
panel = self.create_panel(consumer=consumer)
assert panel.state.consumer.controller.consumer is consumer
panel.app = self.app
_import = panel.app.loader.import_from_cwd = Mock()
_reload = Mock()
with pytest.raises(ValueError):
panel.handle('pool_restart', {'reloader': _reload})
self.app.conf.worker_pool_restarts = True
panel.handle('pool_restart', {'reloader': _reload})
consumer.controller.pool.restart.assert_called()
consumer.reset_rate_limits.assert_called_with()
consumer.update_strategies.assert_called_with()
_reload.assert_not_called()
_import.assert_not_called()
consumer.controller.pool.restart.side_effect = NotImplementedError()
panel.handle('pool_restart', {'reloader': _reload})
consumer.controller.consumer = None
panel.handle('pool_restart', {'reloader': _reload})
def test_pool_restart_import_modules(self, _debug):
consumer = Consumer(self.app)
consumer.controller = _WC(app=self.app)
consumer.controller.consumer = consumer
consumer.controller.pool.restart = Mock()
consumer.reset_rate_limits = Mock(name='reset_rate_limits()')
consumer.update_strategies = Mock(name='update_strategies()')
panel = self.create_panel(consumer=consumer)
panel.app = self.app
assert panel.state.consumer.controller.consumer is consumer
_import = consumer.controller.app.loader.import_from_cwd = Mock()
_reload = Mock()
self.app.conf.worker_pool_restarts = True
with patch('sys.modules'):
panel.handle('pool_restart', {
'modules': ['foo', 'bar'],
'reloader': _reload,
})
consumer.controller.pool.restart.assert_called()
consumer.reset_rate_limits.assert_called_with()
consumer.update_strategies.assert_called_with()
_reload.assert_not_called()
_import.assert_has_calls([call('bar'), call('foo')], any_order=True)
assert _import.call_count == 2
def test_pool_restart_reload_modules(self):
consumer = Consumer(self.app)
consumer.controller = _WC(app=self.app)
consumer.controller.consumer = consumer
consumer.controller.pool.restart = Mock()
consumer.reset_rate_limits = Mock(name="reset_rate_limits()")
consumer.update_strategies = Mock(name="update_strategies()")
panel = self.create_panel(consumer=consumer)
panel.app = self.app
_import = panel.app.loader.import_from_cwd = Mock()
_reload = Mock()
self.app.conf.worker_pool_restarts = True
with patch.dict(sys.modules, {"foo": None}):
panel.handle("pool_restart", {"modules": ["foo"], "reload": False, "reloader": _reload})
consumer.controller.pool.restart.assert_called()
_reload.assert_not_called()
_import.assert_not_called()
_import.reset_mock()
_reload.reset_mock()
consumer.controller.pool.restart.reset_mock()
panel.handle("pool_restart", {"modules": ["foo"], "reload": True, "reloader": _reload})
consumer.controller.pool.restart.assert_called()
_reload.assert_called()
_import.assert_not_called()
class test_frame_handler:
@pytest.fixture(autouse=True)
def setup_conn(self):
self.conn = Mock(name='connection')
self.conn.bytes_recv = 0
self.callback = Mock(name='callback')
self.g = frame_handler(self.conn, self.callback)
def test_header(self):
buf = pack('>HH', 60, 51)
assert self.g((1, 1, buf))
self.callback.assert_called_with(1, (60, 51), buf, None)
assert self.conn.bytes_recv
def test_header_message_empty_body(self):
assert not self.g((1, 1, pack('>HH', *spec.Basic.Deliver)))
self.callback.assert_not_called()
with pytest.raises(UnexpectedFrame):
self.g((1, 1, pack('>HH', *spec.Basic.Deliver)))
m = Message()
m.properties = {}
buf = pack('>HxxQ', m.CLASS_ID, 0)
buf += m._serialize_properties()
assert self.g((2, 1, buf))
self.callback.assert_called()
msg = self.callback.call_args[0][3]
self.callback.assert_called_with(
1, msg.frame_method, msg.frame_args, msg,
)
def test_header_message_content(self):
assert not self.g((1, 1, pack('>HH', *spec.Basic.Deliver)))
self.callback.assert_not_called()
m = Message()
m.properties = {}
buf = pack('>HxxQ', m.CLASS_ID, 16)
buf += m._serialize_properties()
assert not self.g((2, 1, buf))
self.callback.assert_not_called()
assert not self.g((3, 1, b'thequick'))
self.callback.assert_not_called()
assert self.g((3, 1, b'brownfox'))
self.callback.assert_called()
msg = self.callback.call_args[0][3]
self.callback.assert_called_with(
1, msg.frame_method, msg.frame_args, msg,
)
assert msg.body == b'thequickbrownfox'
def test_heartbeat_frame(self):
assert not self.g((8, 1, ''))
self.callback.assert_not_called()
assert self.conn.bytes_recv
def test_apply_target__raises_BaseException_raises_else(self, reraise):
target = Mock(name='target')
callback = Mock(name='callback')
reraise.side_effect = KeyError()
target.side_effect = BaseException()
with pytest.raises(KeyError):
apply_target(target, callback=callback)
callback.assert_not_called()
def test_filter(self):
callback = Mock()
filt = filter_callback(callback, ["add", "mul"])
t1 = {"task": "add"}
t2 = {"task": "div"}
message = Mock()
filt(t2, message)
callback.assert_not_called()
filt(t1, message)
callback.assert_called_with(t1, message)
def test_filter(self):
callback = Mock()
filt = filter_callback(callback, ['add', 'mul'])
t1 = {'task': 'add'}
t2 = {'task': 'div'}
message = Mock()
filt(t2, message)
callback.assert_not_called()
filt(t1, message)
callback.assert_called_with(t1, message)
def test_event_from_message_localize_disabled(self):
r = self.app.events.Receiver(Mock(), node_id='celery.tests')
r.adjust_clock = Mock()
ts_adjust = Mock()
r.event_from_message(
{'type': 'worker-online', 'clock': 313},
localize=False,
adjust_timestamp=ts_adjust,
)
ts_adjust.assert_not_called()
r.adjust_clock.assert_called_with(313)
def test_migrate(self, app, name='testcelery'):
connection_kwargs = {
'transport_options': {'polling_interval': 0.01}
}
x = Connection('memory://foo', **connection_kwargs)
y = Connection('memory://foo', **connection_kwargs)
# use separate state
x.default_channel.queues = {}
y.default_channel.queues = {}
ex = Exchange(name, 'direct')
q = Queue(name, exchange=ex, routing_key=name)
q(x.default_channel).declare()
Producer(x).publish('foo', exchange=name, routing_key=name)
Producer(x).publish('bar', exchange=name, routing_key=name)
Producer(x).publish('baz', exchange=name, routing_key=name)
assert x.default_channel.queues
assert not y.default_channel.queues
migrate_tasks(x, y, accept=['text/plain'], app=app)
yq = q(y.default_channel)
assert yq.get().body == ensure_bytes('foo')
assert yq.get().body == ensure_bytes('bar')
assert yq.get().body == ensure_bytes('baz')
Producer(x).publish('foo', exchange=name, routing_key=name)
callback = Mock()
migrate_tasks(x, y,
callback=callback, accept=['text/plain'], app=app)
callback.assert_called()
migrate = Mock()
Producer(x).publish('baz', exchange=name, routing_key=name)
migrate_tasks(x, y, callback=callback,
migrate=migrate, accept=['text/plain'], app=app)
migrate.assert_called()
with patch('kombu.transport.virtual.Channel.queue_declare') as qd:
def effect(*args, **kwargs):
if kwargs.get('passive'):
raise ChannelError('some channel error')
return 0, 3, 0
qd.side_effect = effect
migrate_tasks(x, y, app=app)
x = Connection('memory://', **connection_kwargs)
x.default_channel.queues = {}
y.default_channel.queues = {}
callback = Mock()
migrate_tasks(x, y,
callback=callback, accept=['text/plain'], app=app)
callback.assert_not_called()
def test_accept__content_disallowed(self):
conn = Connection('memory://')
q = Queue('foo', exchange=self.exchange)
p = conn.Producer()
p.publish(
{'complex': object()},
declare=[q], exchange=self.exchange, serializer='pickle',
)
callback = Mock(name='callback')
with conn.Consumer(queues=[q], callbacks=[callback]) as consumer:
with pytest.raises(consumer.ContentDisallowed):
conn.drain_events(timeout=1)
callback.assert_not_called()
def test_migrate(self, app, name="testcelery"):
x = Connection("memory://foo")
y = Connection("memory://foo")
# use separate state
x.default_channel.queues = {}
y.default_channel.queues = {}
ex = Exchange(name, "direct")
q = Queue(name, exchange=ex, routing_key=name)
q(x.default_channel).declare()
Producer(x).publish("foo", exchange=name, routing_key=name)
Producer(x).publish("bar", exchange=name, routing_key=name)
Producer(x).publish("baz", exchange=name, routing_key=name)
assert x.default_channel.queues
assert not y.default_channel.queues
migrate_tasks(x, y, accept=["text/plain"], app=app)
yq = q(y.default_channel)
assert yq.get().body == ensure_bytes("foo")
assert yq.get().body == ensure_bytes("bar")
assert yq.get().body == ensure_bytes("baz")
Producer(x).publish("foo", exchange=name, routing_key=name)
callback = Mock()
migrate_tasks(x, y, callback=callback, accept=["text/plain"], app=app)
callback.assert_called()
migrate = Mock()
Producer(x).publish("baz", exchange=name, routing_key=name)
migrate_tasks(x, y, callback=callback, migrate=migrate, accept=["text/plain"], app=app)
migrate.assert_called()
with patch("kombu.transport.virtual.Channel.queue_declare") as qd:
def effect(*args, **kwargs):
if kwargs.get("passive"):
raise ChannelError("some channel error")
return 0, 3, 0
qd.side_effect = effect
migrate_tasks(x, y, app=app)
x = Connection("memory://")
x.default_channel.queues = {}
y.default_channel.queues = {}
callback = Mock()
migrate_tasks(x, y, callback=callback, accept=["text/plain"], app=app)
callback.assert_not_called()
def test_loop__todo(self):
deferred = Mock(name='cb_deferred')
def defer():
self.hub.call_soon(deferred)
callbacks = [Mock(name='cb1', wraps=defer), Mock(name='cb2')]
for cb in callbacks:
self.hub.call_soon(cb)
self.hub._ready.add(None)
next(self.hub.loop)
callbacks[0].assert_called_once_with()
callbacks[1].assert_called_once_with()
deferred.assert_not_called()
def test_pool_restart_import_modules(self, _debug):
consumer = Consumer(self.app)
consumer.controller = _WC(app=self.app)
consumer.controller.consumer = consumer
consumer.controller.pool.restart = Mock()
consumer.reset_rate_limits = Mock(name="reset_rate_limits()")
consumer.update_strategies = Mock(name="update_strategies()")
panel = self.create_panel(consumer=consumer)
panel.app = self.app
assert panel.state.consumer.controller.consumer is consumer
_import = consumer.controller.app.loader.import_from_cwd = Mock()
_reload = Mock()
self.app.conf.worker_pool_restarts = True
with patch("sys.modules"):
panel.handle("pool_restart", {"modules": ["foo", "bar"], "reloader": _reload})
consumer.controller.pool.restart.assert_called()
consumer.reset_rate_limits.assert_called_with()
consumer.update_strategies.assert_called_with()
_reload.assert_not_called()
_import.assert_has_calls([call("bar"), call("foo")], any_order=True)
assert _import.call_count == 2
def test_grow_shrink(self):
x = LaxBoundedSemaphore(1)
assert x.initial_value == 1
cb1 = Mock()
x.acquire(cb1, 1)
cb1.assert_called_with(1)
assert x.value == 0
cb2 = Mock()
x.acquire(cb2, 2)
cb2.assert_not_called()
assert x.value == 0
cb3 = Mock()
x.acquire(cb3, 3)
cb3.assert_not_called()
x.grow(2)
cb2.assert_called_with(2)
cb3.assert_called_with(3)
assert x.value == 2
assert x.initial_value == 3
assert not x._waiting
x.grow(3)
for i in range(x.initial_value):
assert x.acquire(Mock())
assert not x.acquire(Mock())
x.clear()
x.shrink(3)
for i in range(x.initial_value):
assert x.acquire(Mock())
assert not x.acquire(Mock())
assert x.value == 0
for i in range(100):
x.release()
assert x.value == x.initial_value
def test_pool_restart_reload_modules(self):
consumer = Consumer(self.app)
consumer.controller = _WC(app=self.app)
consumer.controller.consumer = consumer
consumer.controller.pool.restart = Mock()
consumer.reset_rate_limits = Mock(name='reset_rate_limits()')
consumer.update_strategies = Mock(name='update_strategies()')
panel = self.create_panel(consumer=consumer)
panel.app = self.app
_import = panel.app.loader.import_from_cwd = Mock()
_reload = Mock()
self.app.conf.worker_pool_restarts = True
with patch.dict(sys.modules, {'foo': None}):
panel.handle('pool_restart', {
'modules': ['foo'],
'reload': False,
'reloader': _reload,
})
consumer.controller.pool.restart.assert_called()
_reload.assert_not_called()
_import.assert_not_called()
_import.reset_mock()
_reload.reset_mock()
consumer.controller.pool.restart.reset_mock()
panel.handle('pool_restart', {
'modules': ['foo'],
'reload': True,
'reloader': _reload,
})
consumer.controller.pool.restart.assert_called()
_reload.assert_called()
_import.assert_not_called()
def test_repr(self):
obj = Mock(name='cb')
assert 'promise' in repr(ChannelPromise(obj))
obj.assert_not_called()
