def test_start(self):
try:
c = Mock()
c.app.connection_for_read = _amqp_connection()
mingle = Mingle(c)
self.assertTrue(mingle.enabled)
Aig = LimitedSet()
Big = LimitedSet()
Aig.add("Aig-1")
Aig.add("Aig-2")
Big.add("Big-1")
I = c.app.control.inspect.return_value = Mock()
I.hello.return_value = {
"*****@*****.**": {"clock": 312, "revoked": Aig._data},
"*****@*****.**": {"clock": 29, "revoked": Big._data},
"*****@*****.**": {"error": "unknown method"},
}
mingle.start(c)
I.hello.assert_called_with(c.hostname, worker_state.revoked._data)
c.app.clock.adjust.assert_has_calls([call(312), call(29)], any_order=True)
self.assertIn("Aig-1", worker_state.revoked)
self.assertIn("Aig-2", worker_state.revoked)
self.assertIn("Big-1", worker_state.revoked)
finally:
worker_state.revoked.clear()
def test_enter__exit(self):
hub = Hub()
P = hub.poller = Mock()
on_close = Mock()
hub.on_close.add(on_close)
try:
read_A = Mock()
read_B = Mock()
hub.add_reader(10, read_A)
hub.add_reader(File(11), read_B)
write_A = Mock()
write_B = Mock()
hub.add_writer(20, write_A)
hub.add_writer(File(21), write_B)
self.assertTrue(hub.readers)
self.assertTrue(hub.writers)
finally:
hub.close()
self.assertFalse(hub.readers)
self.assertFalse(hub.writers)
P.unregister.assert_has_calls([call(10), call(11), call(20), call(21)], any_order=True)
on_close.assert_called_with(hub)
def test_with_uid(
self, initgroups, setuid, setgid, getpwuid, parse_gid, parse_uid, getuid, geteuid, getgid, getegid
):
geteuid.return_value = 10
getuid.return_value = 10
class pw_struct(object):
pw_gid = 50001
def raise_on_second_call(*args, **kwargs):
setuid.side_effect = OSError()
setuid.side_effect.errno = errno.EPERM
setuid.side_effect = raise_on_second_call
getpwuid.return_value = pw_struct()
parse_uid.return_value = 5001
parse_gid.return_value = 5001
maybe_drop_privileges(uid="user")
parse_uid.assert_called_with("user")
getpwuid.assert_called_with(5001)
setgid.assert_called_with(50001)
initgroups.assert_called_with(5001, 50001)
setuid.assert_has_calls([call(5001), call(0)])
setuid.side_effect = raise_on_second_call
def to_root_on_second_call(mock, first):
return_value = [first]
def on_first_call(*args, **kwargs):
ret, return_value[0] = return_value[0], 0
return ret
mock.side_effect = on_first_call
to_root_on_second_call(geteuid, 10)
to_root_on_second_call(getuid, 10)
with self.assertRaises(AssertionError):
maybe_drop_privileges(uid="user")
getuid.return_value = getuid.side_effect = None
geteuid.return_value = geteuid.side_effect = None
getegid.return_value = 0
getgid.return_value = 0
setuid.side_effect = raise_on_second_call
with self.assertRaises(AssertionError):
maybe_drop_privileges(gid="group")
getuid.reset_mock()
geteuid.reset_mock()
setuid.reset_mock()
getuid.side_effect = geteuid.side_effect = None
def raise_on_second_call(*args, **kwargs):
setuid.side_effect = OSError()
setuid.side_effect.errno = errno.ENOENT
setuid.side_effect = raise_on_second_call
with self.assertRaises(OSError):
maybe_drop_privileges(uid="user")
def test_add_writers(self):
hub = Hub()
P = hub.poller = Mock()
write_A = Mock()
write_B = Mock()
hub.add_writer(20, write_A)
hub.add_writer(File(21), write_B)
P.register.assert_has_calls([
call(20, hub.WRITE),
call(21, hub.WRITE),
],
any_order=True)
self.assertEqual(hub.writers[20], (write_A, ()))
self.assertEqual(hub.writers[21], (write_B, ()))
hub.remove(20)
self.assertNotIn(20, hub.writers)
hub.remove(File(21))
self.assertNotIn(21, hub.writers)
P.unregister.assert_has_calls([
call(20),
call(21),
])
def test_add_remove_readers(self):
hub = Hub()
P = hub.poller = Mock()
read_A = Mock()
read_B = Mock()
hub.add_reader(10, read_A, 10)
hub.add_reader(File(11), read_B, 11)
P.register.assert_has_calls([
call(10, hub.READ | hub.ERR),
call(11, hub.READ | hub.ERR),
],
any_order=True)
self.assertEqual(hub.readers[10], (read_A, (10, )))
self.assertEqual(hub.readers[11], (read_B, (11, )))
hub.remove(10)
self.assertNotIn(10, hub.readers)
hub.remove(File(11))
self.assertNotIn(11, hub.readers)
P.unregister.assert_has_calls([
call(10),
call(11),
])
def test_pool_restart_import_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
assert panel.state.consumer.controller.consumer is consumer
_import = consumer.controller.app.loader.import_from_cwd = Mock()
_reload = Mock()
self.app.conf.CELERYD_POOL_RESTARTS = True
panel.handle('pool_restart', {'modules': ['foo', 'bar'],
'reloader': _reload})
self.assertTrue(consumer.controller.pool.restart.called)
consumer.reset_rate_limits.assert_called_with()
consumer.update_strategies.assert_called_with()
self.assertFalse(_reload.called)
self.assertItemsEqual(
[call('bar'), call('foo')],
_import.call_args_list,
)
def test_enter__exit(self):
hub = Hub()
P = hub.poller = Mock()
on_close = Mock()
hub.on_close.add(on_close)
try:
read_A = Mock()
read_B = Mock()
hub.add_reader(10, read_A)
hub.add_reader(File(11), read_B)
write_A = Mock()
write_B = Mock()
hub.add_writer(20, write_A)
hub.add_writer(File(21), write_B)
self.assertTrue(hub.readers)
self.assertTrue(hub.writers)
finally:
assert hub.poller
hub.close()
self.assertFalse(hub.readers)
self.assertFalse(hub.writers)
P.unregister.assert_has_calls([
call(10),
call(11),
call(20),
call(21),
],
any_order=True)
on_close.assert_called_with(hub)
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)
self.assertEqual(_import.call_count, 2)
def test_with_uid(self, initgroups, setuid, setgid, getpwuid,
parse_gid, parse_uid, getuid, geteuid, getgid,
getegid):
geteuid.return_value = 10
getuid.return_value = 10
class pw_struct(object):
pw_gid = 50001
def raise_on_second_call(*args, **kwargs):
setuid.side_effect = OSError()
setuid.side_effect.errno = errno.EPERM
setuid.side_effect = raise_on_second_call
getpwuid.return_value = pw_struct()
parse_uid.return_value = 5001
parse_gid.return_value = 5001
maybe_drop_privileges(uid='user')
parse_uid.assert_called_with('user')
getpwuid.assert_called_with(5001)
setgid.assert_called_with(50001)
initgroups.assert_called_with(5001, 50001)
setuid.assert_has_calls([call(5001), call(0)])
setuid.side_effect = raise_on_second_call
def to_root_on_second_call(mock, first):
return_value = [first]
def on_first_call(*args, **kwargs):
ret, return_value[0] = return_value[0], 0
return ret
mock.side_effect = on_first_call
to_root_on_second_call(geteuid, 10)
to_root_on_second_call(getuid, 10)
with self.assertRaises(AssertionError):
maybe_drop_privileges(uid='user')
getuid.return_value = getuid.side_effect = None
geteuid.return_value = geteuid.side_effect = None
getegid.return_value = 0
getgid.return_value = 0
setuid.side_effect = raise_on_second_call
with self.assertRaises(AssertionError):
maybe_drop_privileges(gid='group')
getuid.reset_mock()
geteuid.reset_mock()
setuid.reset_mock()
getuid.side_effect = geteuid.side_effect = None
def raise_on_second_call(*args, **kwargs):
setuid.side_effect = OSError()
setuid.side_effect.errno = errno.ENOENT
setuid.side_effect = raise_on_second_call
with self.assertRaises(OSError):
maybe_drop_privileges(uid='user')
def test_pool_restart_import_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
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
panel.handle('pool_restart', {
'modules': ['foo', 'bar'],
'reloader': _reload
})
self.assertTrue(consumer.controller.pool.restart.called)
consumer.reset_rate_limits.assert_called_with()
consumer.update_strategies.assert_called_with()
self.assertFalse(_reload.called)
self.assertItemsEqual(
[call('bar'), call('foo')],
_import.call_args_list,
)
def test_on_chord_part_return(self, restore):
b = self.Backend(app=self.app, new_join=True)
def create_task():
tid = uuid()
task = Mock(name='task-{0}'.format(tid))
task.name = 'foobarbaz'
self.app.tasks['foobarbaz'] = task
task.request.chord = signature(task)
task.request.id = tid
task.request.chord['chord_size'] = 10
task.request.group = 'group_id'
return task
tasks = [create_task() for i in range(10)]
for i in range(10):
b.on_chord_part_return(tasks[i].request, states.SUCCESS, i)
self.assertTrue(b.client.rpush.call_count)
b.client.rpush.reset_mock()
self.assertTrue(b.client.lrange.call_count)
jkey = b.get_key_for_group('group_id', '.j')
tkey = b.get_key_for_group('group_id', '.t')
b.client.delete.assert_has_calls([call(jkey), call(tkey)])
b.client.expire.assert_has_calls([
call(jkey, 86400), call(tkey, 86400),
])
def test_on_chord_part_return(self, restore):
b = self.Backend(app=self.app, new_join=True)
def create_task():
tid = uuid()
task = Mock(name='task-{0}'.format(tid))
task.name = 'foobarbaz'
self.app.tasks['foobarbaz'] = task
task.request.chord = signature(task)
task.request.id = tid
task.request.chord['chord_size'] = 10
task.request.group = 'group_id'
return task
tasks = [create_task() for i in range(10)]
for i in range(10):
b.on_chord_part_return(tasks[i], states.SUCCESS, i)
self.assertTrue(b.client.rpush.call_count)
b.client.rpush.reset_mock()
self.assertTrue(b.client.lrange.call_count)
jkey = b.get_key_for_group('group_id', '.j')
tkey = b.get_key_for_group('group_id', '.t')
b.client.delete.assert_has_calls([call(jkey), call(tkey)])
b.client.expire.assert_has_calls([
call(jkey, 86400),
call(tkey, 86400),
])
def test_on_tick(self, sleep):
on_tick = Mock(name='on_tick')
t = timer2.Timer(on_tick=on_tick)
ne = t._next_entry = Mock(name='_next_entry')
ne.return_value = 3.33
self.on_nth_call_do(ne, t._is_shutdown.set, 3)
t.run()
sleep.assert_called_with(3.33)
on_tick.assert_has_calls([call(3.33), call(3.33), call(3.33)])
def test_closes(self):
with patch("os.close") as _close:
with patch("celery.platforms.get_fdmax") as fdmax:
fdmax.return_value = 3
close_open_fds()
_close.assert_has_calls([call(2), call(1), call(0)])
_close.side_effect = OSError()
_close.side_effect.errno = errno.EBADF
close_open_fds()
def test_on_tick(self, sleep):
on_tick = Mock(name='on_tick')
t = timer2.Timer(on_tick=on_tick)
ne = t._next_entry = Mock(name='_next_entry')
ne.return_value = 3.33
self.on_nth_call_do(ne, t._is_shutdown.set, 3)
t.run()
sleep.assert_called_with(3.33)
on_tick.assert_has_calls([call(3.33), call(3.33), call(3.33)])
def test_receive_multi(self):
r = self.app.events.Receiver(Mock(name='connection'))
r.process = Mock(name='process')
efm = r.event_from_message = Mock(name='event_from_message')
def on_efm(*args):
return args
efm.side_effect = on_efm
r._receive([1, 2, 3], Mock())
r.process.assert_has_calls([call(1), call(2), call(3)])
def test_receive_multi(self):
r = self.app.events.Receiver(Mock(name='connection'))
r.process = Mock(name='process')
efm = r.event_from_message = Mock(name='event_from_message')
def on_efm(*args):
return args
efm.side_effect = on_efm
r._receive([1, 2, 3], Mock())
r.process.assert_has_calls([call(1), call(2), call(3)])
def test_closes(self):
with patch('os.close') as _close:
with patch('os.closerange', create=True) as closerange:
with patch('celery.platforms.get_fdmax') as fdmax:
fdmax.return_value = 3
close_open_fds()
if not closerange.called:
_close.assert_has_calls([call(2), call(1), call(0)])
_close.side_effect = OSError()
_close.side_effect.errno = errno.EBADF
close_open_fds()
def test_closes(self):
with patch('os.close') as _close:
with patch('os.closerange', create=True) as closerange:
with patch('celery.platforms.get_fdmax') as fdmax:
fdmax.return_value = 3
close_open_fds()
if not closerange.called:
_close.assert_has_calls([call(2), call(1), call(0)])
_close.side_effect = OSError()
_close.side_effect.errno = errno.EBADF
close_open_fds()
def test_on_chord_part_return(self, restore):
tasks = [self.create_task() for i in range(10)]
for i in range(10):
self.b.on_chord_part_return(tasks[i].request, states.SUCCESS, i)
self.assertTrue(self.b.client.rpush.call_count)
self.b.client.rpush.reset_mock()
self.assertTrue(self.b.client.lrange.call_count)
jkey = self.b.get_key_for_group('group_id', '.j')
tkey = self.b.get_key_for_group('group_id', '.t')
self.b.client.delete.assert_has_calls([call(jkey), call(tkey)])
self.b.client.expire.assert_has_calls([
call(jkey, 86400), call(tkey, 86400),
])
def test_on_chord_part_return(self, restore):
tasks = [self.create_task() for i in range(10)]
for i in range(10):
self.b.on_chord_part_return(tasks[i].request, states.SUCCESS, i)
self.assertTrue(self.b.client.rpush.call_count)
self.b.client.rpush.reset_mock()
self.assertTrue(self.b.client.lrange.call_count)
jkey = self.b.get_key_for_group('group_id', '.j')
tkey = self.b.get_key_for_group('group_id', '.t')
self.b.client.delete.assert_has_calls([call(jkey), call(tkey)])
self.b.client.expire.assert_has_calls([
call(jkey, 86400), call(tkey, 86400),
])
def test_send_buffer_group(self):
buf_received = [None]
producer = MockProducer()
producer.connection = self.app.connection_for_write()
connection = Mock()
connection.transport.driver_type = 'amqp'
eventer = self.app.events.Dispatcher(
connection, enabled=False,
buffer_group={'task'}, buffer_limit=2,
)
eventer.producer = producer
eventer.enabled = True
eventer._publish = Mock(name='_publish')
def on_eventer_publish(events, *args, **kwargs):
buf_received[0] = list(events)
eventer._publish.side_effect = on_eventer_publish
self.assertFalse(eventer._group_buffer['task'])
eventer.on_send_buffered = Mock(name='on_send_buffered')
eventer.send('task-received', uuid=1)
prev_buffer = eventer._group_buffer['task']
self.assertTrue(eventer._group_buffer['task'])
eventer.on_send_buffered.assert_called_with()
eventer.send('task-received', uuid=1)
self.assertFalse(eventer._group_buffer['task'])
eventer._publish.assert_has_calls([
call([], eventer.producer, 'task.multi'),
])
# clear in place
self.assertIs(eventer._group_buffer['task'], prev_buffer)
self.assertEqual(len(buf_received[0]), 2)
eventer.on_send_buffered = None
eventer.send('task-received', uuid=1)
def test_send_buffer_group(self):
buf_received = [None]
producer = MockProducer()
producer.connection = self.app.connection()
connection = Mock()
connection.transport.driver_type = 'amqp'
eventer = self.app.events.Dispatcher(
connection, enabled=False,
buffer_group={'task'}, buffer_limit=2,
)
eventer.producer = producer
eventer.enabled = True
eventer._publish = Mock(name='_publish')
def on_eventer_publish(events, *args, **kwargs):
buf_received[0] = list(events)
eventer._publish.side_effect = on_eventer_publish
self.assertFalse(eventer._group_buffer['task'])
eventer.on_send_buffered = Mock(name='on_send_buffered')
eventer.send('task-received', uuid=1)
prev_buffer = eventer._group_buffer['task']
self.assertTrue(eventer._group_buffer['task'])
eventer.on_send_buffered.assert_called_with()
eventer.send('task-received', uuid=1)
self.assertFalse(eventer._group_buffer['task'])
eventer._publish.assert_has_calls([
call([], eventer.producer, 'task.multi'),
])
# clear in place
self.assertIs(eventer._group_buffer['task'], prev_buffer)
self.assertEqual(len(buf_received[0]), 2)
eventer.on_send_buffered = None
eventer.send('task-received', uuid=1)
def test_with_uid(self, initgroups, setuid, setgid, getpwuid, parse_uid):
class pw_struct(object):
pw_gid = 50001
def raise_on_second_call(*args, **kwargs):
setuid.side_effect = OSError()
setuid.side_effect.errno = errno.EPERM
setuid.side_effect = raise_on_second_call
getpwuid.return_value = pw_struct()
parse_uid.return_value = 5001
maybe_drop_privileges(uid="user")
parse_uid.assert_called_with("user")
getpwuid.assert_called_with(5001)
setgid.assert_called_with(50001)
initgroups.assert_called_with(5001, 50001)
setuid.assert_has_calls([call(5001), call(0)])
def test_select(self):
ebadf = socket.error()
ebadf.errno = errno.EBADF
with patch("select.select") as select:
select.return_value = ([3], [], [])
self.assertEqual(asynpool._select({3}), ([3], [], 0))
select.return_value = ([], [], [3])
self.assertEqual(asynpool._select({3}, None, {3}), ([3], [], 0))
eintr = socket.error()
eintr.errno = errno.EINTR
select.side_effect = eintr
readers = {3}
self.assertEqual(asynpool._select(readers), ([], [], 1))
self.assertIn(3, readers)
with patch("select.select") as select:
select.side_effect = ebadf
readers = {3}
self.assertEqual(asynpool._select(readers), ([], [], 1))
select.assert_has_calls([call([3], [], [], 0)])
self.assertNotIn(3, readers)
with patch("select.select") as select:
select.side_effect = MemoryError()
with self.assertRaises(MemoryError):
asynpool._select({1})
with patch("select.select") as select:
def se(*args):
select.side_effect = MemoryError()
raise ebadf
select.side_effect = se
with self.assertRaises(MemoryError):
asynpool._select({3})
with patch("select.select") as select:
def se2(*args):
select.side_effect = socket.error()
select.side_effect.errno = 1321
raise ebadf
select.side_effect = se2
with self.assertRaises(socket.error):
asynpool._select({3})
with patch("select.select") as select:
select.side_effect = socket.error()
select.side_effect.errno = 34134
with self.assertRaises(socket.error):
asynpool._select({3})
def test_with_uid(self, initgroups, setuid, setgid, getpwuid,
parse_uid):
class pw_struct(object):
pw_gid = 50001
def raise_on_second_call(*args, **kwargs):
setuid.side_effect = OSError()
setuid.side_effect.errno = errno.EPERM
setuid.side_effect = raise_on_second_call
getpwuid.return_value = pw_struct()
parse_uid.return_value = 5001
maybe_drop_privileges(uid='user')
parse_uid.assert_called_with('user')
getpwuid.assert_called_with(5001)
setgid.assert_called_with(50001)
initgroups.assert_called_with(5001, 50001)
setuid.assert_has_calls([call(5001), call(0)])
def test_send_all(self):
nodes = [Mock(name='n1'), Mock(name='n2')]
self.cluster.getpids = Mock(name='getpids')
self.cluster.getpids.return_value = nodes
self.cluster.send_all(15)
self.cluster.on_node_signal.assert_has_calls(
call(node, 'TERM') for node in nodes
)
for node in nodes:
node.send.assert_called_with(15, self.cluster.on_node_signal_dead)
def test_send_all(self):
nodes = [Mock(name='n1'), Mock(name='n2')]
self.cluster.getpids = Mock(name='getpids')
self.cluster.getpids.return_value = nodes
self.cluster.send_all(15)
self.cluster.on_node_signal.assert_has_calls(
call(node, 'TERM') for node in nodes
)
for node in nodes:
node.send.assert_called_with(15, self.cluster.on_node_signal_dead)
def test_pool_restart_import_modules(self):
consumer = Consumer(self.app)
consumer.controller = _WC(app=self.app)
consumer.controller.pool.restart = Mock()
panel = self.create_panel(consumer=consumer)
panel.app = self.app
_import = consumer.controller.app.loader.import_from_cwd = Mock()
_reload = Mock()
self.app.conf.CELERYD_POOL_RESTARTS = True
panel.handle('pool_restart', {'modules': ['foo', 'bar'],
'reloader': _reload})
self.assertTrue(consumer.controller.pool.restart.called)
self.assertFalse(_reload.called)
self.assertItemsEqual(
[call('bar'), call('foo')],
_import.call_args_list,
)
def test_add_remove_readers(self):
hub = Hub()
P = hub.poller = Mock()
read_A = Mock()
read_B = Mock()
hub.add_reader(10, read_A, 10)
hub.add_reader(File(11), read_B, 11)
P.register.assert_has_calls([call(10, hub.READ | hub.ERR), call(File(11), hub.READ | hub.ERR)], any_order=True)
self.assertEqual(hub.readers[10], (read_A, (10,)))
self.assertEqual(hub.readers[11], (read_B, (11,)))
hub.remove(10)
self.assertNotIn(10, hub.readers)
hub.remove(File(11))
self.assertNotIn(11, hub.readers)
P.unregister.assert_has_calls([call(10), call(11)])
def test_add_writers(self):
hub = Hub()
P = hub.poller = Mock()
write_A = Mock()
write_B = Mock()
hub.add_writer(20, write_A)
hub.add_writer(File(21), write_B)
P.register.assert_has_calls([call(20, hub.WRITE), call(File(21), hub.WRITE)], any_order=True)
self.assertEqual(hub.writers[20], (write_A, ()))
self.assertEqual(hub.writers[21], (write_B, ()))
hub.remove(20)
self.assertNotIn(20, hub.writers)
hub.remove(File(21))
self.assertNotIn(21, hub.writers)
P.unregister.assert_has_calls([call(20), call(21)])
def test_remove_db(self, remove):
s = create_persistent_scheduler()[0](app=self.app, schedule_filename="schedule")
s._remove_db()
remove.assert_has_calls([call("schedule" + suffix) for suffix in s.known_suffixes])
err = OSError()
err.errno = errno.ENOENT
remove.side_effect = err
s._remove_db()
err.errno = errno.EPERM
with self.assertRaises(OSError):
s._remove_db()
def test_join_native(self):
backend = SimpleBackend()
results = [self.app.AsyncResult(uuid(), backend=backend) for i in range(10)]
ts = self.app.GroupResult(uuid(), results)
ts.app.backend = backend
backend.ids = [result.id for result in results]
res = ts.join_native()
self.assertEqual(res, list(range(10)))
callback = Mock(name="callback")
self.assertFalse(ts.join_native(callback=callback))
callback.assert_has_calls([call(r.id, i) for i, r in enumerate(ts.results)])
def test_pool_restart_import_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
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
panel.handle("pool_restart", {"modules": ["foo", "bar"], "reloader": _reload})
self.assertTrue(consumer.controller.pool.restart.called)
consumer.reset_rate_limits.assert_called_with()
consumer.update_strategies.assert_called_with()
self.assertFalse(_reload.called)
self.assertItemsEqual([call("bar"), call("foo")], _import.call_args_list)
def test_start(self):
try:
c = Mock()
c.app.connection_for_read = _amqp_connection()
mingle = Mingle(c)
self.assertTrue(mingle.enabled)
Aig = LimitedSet()
Big = LimitedSet()
Aig.add('Aig-1')
Aig.add('Aig-2')
Big.add('Big-1')
I = c.app.control.inspect.return_value = Mock()
I.hello.return_value = {
'*****@*****.**': {
'clock': 312,
'revoked': Aig._data,
},
'*****@*****.**': {
'clock': 29,
'revoked': Big._data,
},
'*****@*****.**': {
'error': 'unknown method',
},
}
mingle.start(c)
I.hello.assert_called_with(c.hostname, worker_state.revoked._data)
c.app.clock.adjust.assert_has_calls([
call(312),
call(29),
],
any_order=True)
self.assertIn('Aig-1', worker_state.revoked)
self.assertIn('Aig-2', worker_state.revoked)
self.assertIn('Big-1', worker_state.revoked)
finally:
worker_state.revoked.clear()
def test_with_guid(self, initgroups, setuid, setgid, parse_gid, parse_uid):
def raise_on_second_call(*args, **kwargs):
setuid.side_effect = OSError()
setuid.side_effect.errno = errno.EPERM
setuid.side_effect = raise_on_second_call
parse_uid.return_value = 5001
parse_gid.return_value = 50001
maybe_drop_privileges(uid="user", gid="group")
parse_uid.assert_called_with("user")
parse_gid.assert_called_with("group")
setgid.assert_called_with(50001)
initgroups.assert_called_with(5001, 50001)
setuid.assert_has_calls([call(5001), call(0)])
setuid.side_effect = None
with self.assertRaises(RuntimeError):
maybe_drop_privileges(uid="user", gid="group")
setuid.side_effect = OSError()
setuid.side_effect.errno = errno.EINVAL
with self.assertRaises(OSError):
maybe_drop_privileges(uid="user", gid="group")
def test_remove_db(self, remove):
s = create_persistent_scheduler()[0](app=self.app,
schedule_filename='schedule')
s._remove_db()
remove.assert_has_calls(
[call('schedule' + suffix) for suffix in s.known_suffixes])
err = OSError()
err.errno = errno.ENOENT
remove.side_effect = err
s._remove_db()
err.errno = errno.EPERM
with self.assertRaises(OSError):
s._remove_db()
def test_with_guid(self, initgroups, setuid, setgid, parse_gid, parse_uid):
def raise_on_second_call(*args, **kwargs):
setuid.side_effect = OSError()
setuid.side_effect.errno = errno.EPERM
setuid.side_effect = raise_on_second_call
parse_uid.return_value = 5001
parse_gid.return_value = 50001
maybe_drop_privileges(uid='user', gid='group')
parse_uid.assert_called_with('user')
parse_gid.assert_called_with('group')
setgid.assert_called_with(50001)
initgroups.assert_called_with(5001, 50001)
setuid.assert_has_calls([call(5001), call(0)])
setuid.side_effect = None
with self.assertRaises(RuntimeError):
maybe_drop_privileges(uid='user', gid='group')
setuid.side_effect = OSError()
setuid.side_effect.errno = errno.EINVAL
with self.assertRaises(OSError):
maybe_drop_privileges(uid='user', gid='group')
def test_start(self):
try:
c = Mock()
c.app.connection = _amqp_connection()
mingle = Mingle(c)
self.assertTrue(mingle.enabled)
Aig = LimitedSet()
Big = LimitedSet()
Aig.add('Aig-1')
Aig.add('Aig-2')
Big.add('Big-1')
I = c.app.control.inspect.return_value = Mock()
I.hello.return_value = {
'*****@*****.**': {
'clock': 312,
'revoked': Aig._data,
},
'*****@*****.**': {
'clock': 29,
'revoked': Big._data,
},
'*****@*****.**': {
'error': 'unknown method',
},
}
mingle.start(c)
I.hello.assert_called_with(c.hostname, worker_state.revoked._data)
c.app.clock.adjust.assert_has_calls([
call(312), call(29),
], any_order=True)
self.assertIn('Aig-1', worker_state.revoked)
self.assertIn('Aig-2', worker_state.revoked)
self.assertIn('Big-1', worker_state.revoked)
finally:
worker_state.revoked.clear()
def test_join_native(self):
backend = SimpleBackend()
results = [
self.app.AsyncResult(uuid(), backend=backend) for i in range(10)
]
ts = self.app.GroupResult(uuid(), results)
ts.app.backend = backend
backend.ids = [result.id for result in results]
res = ts.join_native()
self.assertEqual(res, list(range(10)))
callback = Mock(name='callback')
self.assertFalse(ts.join_native(callback=callback))
callback.assert_has_calls(
[call(r.id, i) for i, r in enumerate(ts.results)])
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)
self.assertEqual(_import.call_count, 2)
def test_get_many(self):
for is_dict in True, False:
self.b.mget_returns_dict = is_dict
ids = {uuid(): i for i in range(10)}
for id, i in items(ids):
self.b.mark_as_done(id, i)
it = self.b.get_many(list(ids))
for i, (got_id, got_state) in enumerate(it):
self.assertEqual(got_state["result"], ids[got_id])
self.assertEqual(i, 9)
self.assertTrue(list(self.b.get_many(list(ids))))
self.b._cache.clear()
callback = Mock(name="callback")
it = self.b.get_many(list(ids), on_message=callback)
for i, (got_id, got_state) in enumerate(it):
self.assertEqual(got_state["result"], ids[got_id])
self.assertEqual(i, 9)
self.assertTrue(list(self.b.get_many(list(ids))))
callback.assert_has_calls([call(ANY) for id in ids])
def test_get_many(self):
for is_dict in True, False:
self.b.mget_returns_dict = is_dict
ids = {uuid(): i for i in range(10)}
for id, i in items(ids):
self.b.mark_as_done(id, i)
it = self.b.get_many(list(ids))
for i, (got_id, got_state) in enumerate(it):
self.assertEqual(got_state['result'], ids[got_id])
self.assertEqual(i, 9)
self.assertTrue(list(self.b.get_many(list(ids))))
self.b._cache.clear()
callback = Mock(name='callback')
it = self.b.get_many(list(ids), on_message=callback)
for i, (got_id, got_state) in enumerate(it):
self.assertEqual(got_state['result'], ids[got_id])
self.assertEqual(i, 9)
self.assertTrue(list(self.b.get_many(list(ids))))
callback.assert_has_calls([call(ANY) for id in ids])
def test_select(self):
ebadf = socket.error()
ebadf.errno = errno.EBADF
with patch('select.select') as select:
select.return_value = ([3], [], [])
self.assertEqual(
asynpool._select(set([3])),
([3], [], 0),
)
select.return_value = ([], [], [3])
self.assertEqual(
asynpool._select(set([3]), None, set([3])),
([3], [], 0),
)
eintr = socket.error()
eintr.errno = errno.EINTR
select.side_effect = eintr
readers = set([3])
self.assertEqual(asynpool._select(readers), ([], [], 1))
self.assertIn(3, readers)
with patch('select.select') as select:
select.side_effect = ebadf
readers = set([3])
self.assertEqual(asynpool._select(readers), ([], [], 1))
select.assert_has_calls([call([3], [], [], 0)])
self.assertNotIn(3, readers)
with patch('select.select') as select:
select.side_effect = MemoryError()
with self.assertRaises(MemoryError):
asynpool._select(set([1]))
with patch('select.select') as select:
def se(*args):
select.side_effect = MemoryError()
raise ebadf
select.side_effect = se
with self.assertRaises(MemoryError):
asynpool._select(set([3]))
with patch('select.select') as select:
def se2(*args):
select.side_effect = socket.error()
select.side_effect.errno = 1321
raise ebadf
select.side_effect = se2
with self.assertRaises(socket.error):
asynpool._select(set([3]))
with patch('select.select') as select:
select.side_effect = socket.error()
select.side_effect.errno = 34134
with self.assertRaises(socket.error):
asynpool._select(set([3]))
def test_start(self):
self.cluster.start_node = Mock(name='start_node')
self.cluster.start()
self.cluster.start_node.assert_has_calls(
call(node) for node in self.cluster
)
def test_start(self):
self.cluster.start_node = Mock(name='start_node')
self.cluster.start()
self.cluster.start_node.assert_has_calls(
call(node) for node in self.cluster
)