test_worker.py 34 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077
  1. from __future__ import absolute_import, print_function, unicode_literals
  2. import os
  3. import pytest
  4. import socket
  5. import sys
  6. from collections import deque
  7. from datetime import datetime, timedelta
  8. from functools import partial
  9. from threading import Event
  10. from amqp import ChannelError
  11. from case import Mock, patch, skip
  12. from kombu import Connection
  13. from kombu.common import QoS, ignore_errors
  14. from kombu.transport.base import Message
  15. from kombu.transport.memory import Transport
  16. from kombu.utils.uuid import uuid
  17. from celery.bootsteps import RUN, CLOSE, TERMINATE, StartStopStep
  18. from celery.concurrency.base import BasePool
  19. from celery.exceptions import (
  20. WorkerShutdown, WorkerTerminate, TaskRevokedError,
  21. InvalidTaskError, ImproperlyConfigured,
  22. )
  23. from celery.five import Empty, range, Queue as FastQueue
  24. from celery.platforms import EX_FAILURE
  25. from celery.worker import worker as worker_module
  26. from celery.worker import components
  27. from celery.worker import consumer
  28. from celery.worker import state
  29. from celery.worker.consumer import Consumer
  30. from celery.worker.pidbox import gPidbox
  31. from celery.worker.request import Request
  32. from celery.utils.nodenames import worker_direct
  33. from celery.utils.serialization import pickle
  34. from celery.utils.timer2 import Timer
  35. def MockStep(step=None):
  36. if step is None:
  37. step = Mock(name='step')
  38. else:
  39. step.blueprint = Mock(name='step.blueprint')
  40. step.blueprint.name = 'MockNS'
  41. step.name = 'MockStep(%s)' % (id(step),)
  42. return step
  43. def mock_event_dispatcher():
  44. evd = Mock(name='event_dispatcher')
  45. evd.groups = ['worker']
  46. evd._outbound_buffer = deque()
  47. return evd
  48. def find_step(obj, typ):
  49. return obj.blueprint.steps[typ.name]
  50. def create_message(channel, **data):
  51. data.setdefault('id', uuid())
  52. m = Message(channel, body=pickle.dumps(dict(**data)),
  53. content_type='application/x-python-serialize',
  54. content_encoding='binary',
  55. delivery_info={'consumer_tag': 'mock'})
  56. m.accept = ['application/x-python-serialize']
  57. return m
  58. class ConsumerCase:
  59. def create_task_message(self, channel, *args, **kwargs):
  60. m = self.TaskMessage(*args, **kwargs)
  61. m.channel = channel
  62. m.delivery_info = {'consumer_tag': 'mock'}
  63. return m
  64. class test_Consumer(ConsumerCase):
  65. def setup(self):
  66. self.buffer = FastQueue()
  67. self.timer = Timer()
  68. @self.app.task(shared=False)
  69. def foo_task(x, y, z):
  70. return x * y * z
  71. self.foo_task = foo_task
  72. def teardown(self):
  73. self.timer.stop()
  74. def LoopConsumer(self, buffer=None, controller=None, timer=None, app=None,
  75. without_mingle=True, without_gossip=True,
  76. without_heartbeat=True, **kwargs):
  77. if controller is None:
  78. controller = Mock(name='.controller')
  79. buffer = buffer if buffer is not None else self.buffer.put
  80. timer = timer if timer is not None else self.timer
  81. app = app if app is not None else self.app
  82. c = Consumer(
  83. buffer,
  84. timer=timer,
  85. app=app,
  86. controller=controller,
  87. without_mingle=without_mingle,
  88. without_gossip=without_gossip,
  89. without_heartbeat=without_heartbeat,
  90. **kwargs
  91. )
  92. c.task_consumer = Mock(name='.task_consumer')
  93. c.qos = QoS(c.task_consumer.qos, 10)
  94. c.connection = Mock(name='.connection')
  95. c.controller = c.app.WorkController()
  96. c.heart = Mock(name='.heart')
  97. c.controller.consumer = c
  98. c.pool = c.controller.pool = Mock(name='.controller.pool')
  99. c.node = Mock(name='.node')
  100. c.event_dispatcher = mock_event_dispatcher()
  101. return c
  102. def NoopConsumer(self, *args, **kwargs):
  103. c = self.LoopConsumer(*args, **kwargs)
  104. c.loop = Mock(name='.loop')
  105. return c
  106. def test_info(self):
  107. c = self.NoopConsumer()
  108. c.connection.info.return_value = {'foo': 'bar'}
  109. c.controller.pool.info.return_value = [Mock(), Mock()]
  110. info = c.controller.stats()
  111. assert info['prefetch_count'] == 10
  112. assert info['broker']
  113. def test_start_when_closed(self):
  114. c = self.NoopConsumer()
  115. c.blueprint.state = CLOSE
  116. c.start()
  117. def test_connection(self):
  118. c = self.NoopConsumer()
  119. c.blueprint.start(c)
  120. assert isinstance(c.connection, Connection)
  121. c.blueprint.state = RUN
  122. c.event_dispatcher = None
  123. c.blueprint.restart(c)
  124. assert c.connection
  125. c.blueprint.state = RUN
  126. c.shutdown()
  127. assert c.connection is None
  128. assert c.task_consumer is None
  129. c.blueprint.start(c)
  130. assert isinstance(c.connection, Connection)
  131. c.blueprint.restart(c)
  132. c.stop()
  133. c.shutdown()
  134. assert c.connection is None
  135. assert c.task_consumer is None
  136. def test_close_connection(self):
  137. c = self.NoopConsumer()
  138. c.blueprint.state = RUN
  139. step = find_step(c, consumer.Connection)
  140. connection = c.connection
  141. step.shutdown(c)
  142. connection.close.assert_called()
  143. assert c.connection is None
  144. def test_close_connection__heart_shutdown(self):
  145. c = self.NoopConsumer()
  146. event_dispatcher = c.event_dispatcher
  147. heart = c.heart
  148. c.event_dispatcher.enabled = True
  149. c.blueprint.state = RUN
  150. Events = find_step(c, consumer.Events)
  151. Events.shutdown(c)
  152. Heart = find_step(c, consumer.Heart)
  153. Heart.shutdown(c)
  154. event_dispatcher.close.assert_called()
  155. heart.stop.assert_called_with()
  156. @patch('celery.worker.consumer.consumer.warn')
  157. def test_receive_message_unknown(self, warn):
  158. c = self.LoopConsumer()
  159. c.blueprint.state = RUN
  160. c.steps.pop()
  161. channel = Mock(name='.channeol')
  162. m = create_message(channel, unknown={'baz': '!!!'})
  163. callback = self._get_on_message(c)
  164. callback(m)
  165. warn.assert_called()
  166. @patch('celery.worker.strategy.to_timestamp')
  167. def test_receive_message_eta_OverflowError(self, to_timestamp):
  168. to_timestamp.side_effect = OverflowError()
  169. c = self.LoopConsumer()
  170. c.blueprint.state = RUN
  171. c.steps.pop()
  172. m = self.create_task_message(
  173. Mock(), self.foo_task.name,
  174. args=('2, 2'), kwargs={},
  175. eta=datetime.now().isoformat(),
  176. )
  177. c.update_strategies()
  178. callback = self._get_on_message(c)
  179. callback(m)
  180. assert m.acknowledged
  181. @patch('celery.worker.consumer.consumer.error')
  182. def test_receive_message_InvalidTaskError(self, error):
  183. c = self.LoopConsumer()
  184. c.blueprint.state = RUN
  185. c.steps.pop()
  186. m = self.create_task_message(
  187. Mock(), self.foo_task.name,
  188. args=(1, 2), kwargs='foobarbaz', id=1)
  189. c.update_strategies()
  190. strat = c.strategies[self.foo_task.name] = Mock(name='strategy')
  191. strat.side_effect = InvalidTaskError()
  192. callback = self._get_on_message(c)
  193. callback(m)
  194. error.assert_called()
  195. assert 'Received invalid task message' in error.call_args[0][0]
  196. @patch('celery.worker.consumer.consumer.crit')
  197. def test_on_decode_error(self, crit):
  198. c = self.LoopConsumer()
  199. class MockMessage(Mock):
  200. content_type = 'application/x-msgpack'
  201. content_encoding = 'binary'
  202. body = 'foobarbaz'
  203. message = MockMessage()
  204. c.on_decode_error(message, KeyError('foo'))
  205. assert message.ack.call_count
  206. assert "Can't decode message body" in crit.call_args[0][0]
  207. def _get_on_message(self, c):
  208. if c.qos is None:
  209. c.qos = Mock()
  210. c.task_consumer = Mock()
  211. c.event_dispatcher = mock_event_dispatcher()
  212. c.connection = Mock(name='.connection')
  213. c.connection.get_heartbeat_interval.return_value = 0
  214. c.connection.drain_events.side_effect = WorkerShutdown()
  215. with pytest.raises(WorkerShutdown):
  216. c.loop(*c.loop_args())
  217. assert c.task_consumer.on_message
  218. return c.task_consumer.on_message
  219. def test_receieve_message(self):
  220. c = self.LoopConsumer()
  221. c.blueprint.state = RUN
  222. m = self.create_task_message(
  223. Mock(), self.foo_task.name,
  224. args=[2, 4, 8], kwargs={},
  225. )
  226. c.update_strategies()
  227. callback = self._get_on_message(c)
  228. callback(m)
  229. in_bucket = self.buffer.get_nowait()
  230. assert isinstance(in_bucket, Request)
  231. assert in_bucket.name == self.foo_task.name
  232. assert in_bucket.execute() == 2 * 4 * 8
  233. assert self.timer.empty()
  234. def test_start_channel_error(self):
  235. c = self.NoopConsumer(task_events=False, pool=BasePool())
  236. c.loop.on_nth_call_do_raise(KeyError('foo'), SyntaxError('bar'))
  237. c.channel_errors = (KeyError,)
  238. try:
  239. with pytest.raises(KeyError):
  240. c.start()
  241. finally:
  242. c.timer and c.timer.stop()
  243. def test_start_connection_error(self):
  244. c = self.NoopConsumer(task_events=False, pool=BasePool())
  245. c.loop.on_nth_call_do_raise(KeyError('foo'), SyntaxError('bar'))
  246. c.connection_errors = (KeyError,)
  247. try:
  248. with pytest.raises(SyntaxError):
  249. c.start()
  250. finally:
  251. c.timer and c.timer.stop()
  252. def test_loop_ignores_socket_timeout(self):
  253. class Connection(self.app.connection_for_read().__class__):
  254. obj = None
  255. def drain_events(self, **kwargs):
  256. self.obj.connection = None
  257. raise socket.timeout(10)
  258. c = self.NoopConsumer()
  259. c.connection = Connection()
  260. c.connection.obj = c
  261. c.qos = QoS(c.task_consumer.qos, 10)
  262. c.loop(*c.loop_args())
  263. def test_loop_when_socket_error(self):
  264. class Connection(self.app.connection_for_read().__class__):
  265. obj = None
  266. def drain_events(self, **kwargs):
  267. self.obj.connection = None
  268. raise socket.error('foo')
  269. c = self.LoopConsumer()
  270. c.blueprint.state = RUN
  271. conn = c.connection = Connection()
  272. c.connection.obj = c
  273. c.qos = QoS(c.task_consumer.qos, 10)
  274. with pytest.raises(socket.error):
  275. c.loop(*c.loop_args())
  276. c.blueprint.state = CLOSE
  277. c.connection = conn
  278. c.loop(*c.loop_args())
  279. def test_loop(self):
  280. class Connection(self.app.connection_for_read().__class__):
  281. obj = None
  282. def drain_events(self, **kwargs):
  283. self.obj.connection = None
  284. @property
  285. def supports_heartbeats(self):
  286. return False
  287. c = self.LoopConsumer()
  288. c.blueprint.state = RUN
  289. c.connection = Connection()
  290. c.connection.obj = c
  291. c.qos = QoS(c.task_consumer.qos, 10)
  292. c.loop(*c.loop_args())
  293. c.loop(*c.loop_args())
  294. assert c.task_consumer.consume.call_count
  295. c.task_consumer.qos.assert_called_with(prefetch_count=10)
  296. assert c.qos.value == 10
  297. c.qos.decrement_eventually()
  298. assert c.qos.value == 9
  299. c.qos.update()
  300. assert c.qos.value == 9
  301. c.task_consumer.qos.assert_called_with(prefetch_count=9)
  302. def test_ignore_errors(self):
  303. c = self.NoopConsumer()
  304. c.connection_errors = (AttributeError, KeyError,)
  305. c.channel_errors = (SyntaxError,)
  306. ignore_errors(c, Mock(side_effect=AttributeError('foo')))
  307. ignore_errors(c, Mock(side_effect=KeyError('foo')))
  308. ignore_errors(c, Mock(side_effect=SyntaxError('foo')))
  309. with pytest.raises(IndexError):
  310. ignore_errors(c, Mock(side_effect=IndexError('foo')))
  311. def test_apply_eta_task(self):
  312. c = self.NoopConsumer()
  313. c.qos = QoS(None, 10)
  314. task = Mock(name='task', id='1234213')
  315. qos = c.qos.value
  316. c.apply_eta_task(task)
  317. assert task in state.reserved_requests
  318. assert c.qos.value == qos - 1
  319. assert self.buffer.get_nowait() is task
  320. def test_receieve_message_eta_isoformat(self):
  321. c = self.LoopConsumer()
  322. c.blueprint.state = RUN
  323. c.steps.pop()
  324. m = self.create_task_message(
  325. Mock(), self.foo_task.name,
  326. eta=(datetime.now() + timedelta(days=1)).isoformat(),
  327. args=[2, 4, 8], kwargs={},
  328. )
  329. c.qos = QoS(c.task_consumer.qos, 1)
  330. current_pcount = c.qos.value
  331. c.event_dispatcher.enabled = False
  332. c.update_strategies()
  333. callback = self._get_on_message(c)
  334. callback(m)
  335. c.timer.stop()
  336. c.timer.join(1)
  337. items = [entry[2] for entry in self.timer.queue]
  338. found = 0
  339. for item in items:
  340. if item.args[0].name == self.foo_task.name:
  341. found = True
  342. assert found
  343. assert c.qos.value > current_pcount
  344. c.timer.stop()
  345. def test_pidbox_callback(self):
  346. c = self.NoopConsumer()
  347. con = find_step(c, consumer.Control).box
  348. con.node = Mock()
  349. con.reset = Mock()
  350. con.on_message('foo', 'bar')
  351. con.node.handle_message.assert_called_with('foo', 'bar')
  352. con.node = Mock()
  353. con.node.handle_message.side_effect = KeyError('foo')
  354. con.on_message('foo', 'bar')
  355. con.node.handle_message.assert_called_with('foo', 'bar')
  356. con.node = Mock()
  357. con.node.handle_message.side_effect = ValueError('foo')
  358. con.on_message('foo', 'bar')
  359. con.node.handle_message.assert_called_with('foo', 'bar')
  360. con.reset.assert_called()
  361. def test_revoke(self):
  362. c = self.LoopConsumer()
  363. c.blueprint.state = RUN
  364. c.steps.pop()
  365. channel = Mock(name='channel')
  366. id = uuid()
  367. t = self.create_task_message(
  368. channel, self.foo_task.name,
  369. args=[2, 4, 8], kwargs={}, id=id,
  370. )
  371. state.revoked.add(id)
  372. callback = self._get_on_message(c)
  373. callback(t)
  374. assert self.buffer.empty()
  375. def test_receieve_message_not_registered(self):
  376. c = self.LoopConsumer()
  377. c.blueprint.state = RUN
  378. c.steps.pop()
  379. channel = Mock(name='channel')
  380. m = self.create_task_message(
  381. channel, 'x.X.31x', args=[2, 4, 8], kwargs={},
  382. )
  383. callback = self._get_on_message(c)
  384. assert not callback(m)
  385. with pytest.raises(Empty):
  386. self.buffer.get_nowait()
  387. assert self.timer.empty()
  388. @patch('celery.worker.consumer.consumer.warn')
  389. @patch('celery.worker.consumer.consumer.logger')
  390. def test_receieve_message_ack_raises(self, logger, warn):
  391. c = self.LoopConsumer()
  392. c.blueprint.state = RUN
  393. channel = Mock(name='channel')
  394. m = self.create_task_message(
  395. channel, self.foo_task.name,
  396. args=[2, 4, 8], kwargs={},
  397. )
  398. m.headers = None
  399. c.update_strategies()
  400. c.connection_errors = (socket.error,)
  401. m.reject = Mock()
  402. m.reject.side_effect = socket.error('foo')
  403. callback = self._get_on_message(c)
  404. assert not callback(m)
  405. warn.assert_called()
  406. with pytest.raises(Empty):
  407. self.buffer.get_nowait()
  408. assert self.timer.empty()
  409. m.reject_log_error.assert_called_with(logger, c.connection_errors)
  410. def test_receive_message_eta(self):
  411. if os.environ.get('C_DEBUG_TEST'):
  412. pp = partial(print, file=sys.__stderr__)
  413. else:
  414. def pp(*args, **kwargs):
  415. pass
  416. pp('TEST RECEIVE MESSAGE ETA')
  417. pp('+CREATE MYKOMBUCONSUMER')
  418. c = self.LoopConsumer()
  419. pp('-CREATE MYKOMBUCONSUMER')
  420. c.steps.pop()
  421. channel = Mock(name='channel')
  422. pp('+ CREATE MESSAGE')
  423. m = self.create_task_message(
  424. channel, self.foo_task.name,
  425. args=[2, 4, 8], kwargs={},
  426. eta=(datetime.now() + timedelta(days=1)).isoformat(),
  427. )
  428. pp('- CREATE MESSAGE')
  429. try:
  430. pp('+ BLUEPRINT START 1')
  431. c.blueprint.start(c)
  432. pp('- BLUEPRINT START 1')
  433. p = c.app.conf.broker_connection_retry
  434. c.app.conf.broker_connection_retry = False
  435. pp('+ BLUEPRINT START 2')
  436. c.blueprint.start(c)
  437. pp('- BLUEPRINT START 2')
  438. c.app.conf.broker_connection_retry = p
  439. pp('+ BLUEPRINT RESTART')
  440. c.blueprint.restart(c)
  441. pp('- BLUEPRINT RESTART')
  442. pp('+ GET ON MESSAGE')
  443. callback = self._get_on_message(c)
  444. pp('- GET ON MESSAGE')
  445. pp('+ CALLBACK')
  446. callback(m)
  447. pp('- CALLBACK')
  448. finally:
  449. pp('+ STOP TIMER')
  450. c.timer.stop()
  451. pp('- STOP TIMER')
  452. try:
  453. pp('+ JOIN TIMER')
  454. c.timer.join()
  455. pp('- JOIN TIMER')
  456. except RuntimeError:
  457. pass
  458. in_hold = c.timer.queue[0]
  459. assert len(in_hold) == 3
  460. eta, priority, entry = in_hold
  461. task = entry.args[0]
  462. assert isinstance(task, Request)
  463. assert task.name == self.foo_task.name
  464. assert task.execute() == 2 * 4 * 8
  465. with pytest.raises(Empty):
  466. self.buffer.get_nowait()
  467. def test_reset_pidbox_node(self):
  468. c = self.NoopConsumer()
  469. con = find_step(c, consumer.Control).box
  470. con.node = Mock()
  471. chan = con.node.channel = Mock()
  472. chan.close.side_effect = socket.error('foo')
  473. c.connection_errors = (socket.error,)
  474. con.reset()
  475. chan.close.assert_called_with()
  476. def test_reset_pidbox_node_green(self):
  477. c = self.NoopConsumer(pool=Mock(is_green=True))
  478. con = find_step(c, consumer.Control)
  479. assert isinstance(con.box, gPidbox)
  480. con.start(c)
  481. c.pool.spawn_n.assert_called_with(con.box.loop, c)
  482. def test_green_pidbox_node(self):
  483. pool = Mock()
  484. pool.is_green = True
  485. c = self.NoopConsumer(pool=Mock(is_green=True))
  486. controller = find_step(c, consumer.Control)
  487. class BConsumer(Mock):
  488. def __enter__(self):
  489. self.consume()
  490. return self
  491. def __exit__(self, *exc_info):
  492. self.cancel()
  493. controller.box.node.listen = BConsumer()
  494. connections = []
  495. class Connection(object):
  496. calls = 0
  497. def __init__(self, obj):
  498. connections.append(self)
  499. self.obj = obj
  500. self.default_channel = self.channel()
  501. self.closed = False
  502. def __enter__(self):
  503. return self
  504. def __exit__(self, *exc_info):
  505. self.close()
  506. def channel(self):
  507. return Mock()
  508. def as_uri(self):
  509. return 'dummy://'
  510. def drain_events(self, **kwargs):
  511. if not self.calls:
  512. self.calls += 1
  513. raise socket.timeout()
  514. self.obj.connection = None
  515. controller.box._node_shutdown.set()
  516. def close(self):
  517. self.closed = True
  518. c.connect = lambda: Connection(obj=c)
  519. controller = find_step(c, consumer.Control)
  520. controller.box.loop(c)
  521. controller.box.node.listen.assert_called()
  522. assert controller.box.consumer
  523. controller.box.consumer.consume.assert_called_with()
  524. assert c.connection is None
  525. assert connections[0].closed
  526. @patch('kombu.connection.Connection._establish_connection')
  527. @patch('kombu.utils.functional.sleep')
  528. def test_connect_errback(self, sleep, connect):
  529. c = self.NoopConsumer()
  530. Transport.connection_errors = (ChannelError,)
  531. connect.on_nth_call_do(ChannelError('error'), n=1)
  532. c.connect()
  533. connect.assert_called_with()
  534. def test_stop_pidbox_node(self):
  535. c = self.NoopConsumer()
  536. cont = find_step(c, consumer.Control)
  537. cont._node_stopped = Event()
  538. cont._node_shutdown = Event()
  539. cont._node_stopped.set()
  540. cont.stop(c)
  541. def test_start__loop(self):
  542. class _QoS(object):
  543. prev = 3
  544. value = 4
  545. def update(self):
  546. self.prev = self.value
  547. init_callback = Mock(name='init_callback')
  548. c = self.NoopConsumer(init_callback=init_callback)
  549. c.qos = _QoS()
  550. c.connection = Connection()
  551. c.iterations = 0
  552. def raises_KeyError(*args, **kwargs):
  553. c.iterations += 1
  554. if c.qos.prev != c.qos.value:
  555. c.qos.update()
  556. if c.iterations >= 2:
  557. raise KeyError('foo')
  558. c.loop = raises_KeyError
  559. with pytest.raises(KeyError):
  560. c.start()
  561. assert c.iterations == 2
  562. assert c.qos.prev == c.qos.value
  563. init_callback.reset_mock()
  564. c = self.NoopConsumer(task_events=False, init_callback=init_callback)
  565. c.qos = _QoS()
  566. c.connection = Connection()
  567. c.loop = Mock(side_effect=socket.error('foo'))
  568. with pytest.raises(socket.error):
  569. c.start()
  570. c.loop.assert_called()
  571. def test_reset_connection_with_no_node(self):
  572. c = self.NoopConsumer()
  573. c.steps.pop()
  574. c.blueprint.start(c)
  575. class test_WorkController(ConsumerCase):
  576. def setup(self):
  577. self.worker = self.create_worker()
  578. self._logger = worker_module.logger
  579. self._comp_logger = components.logger
  580. self.logger = worker_module.logger = Mock()
  581. self.comp_logger = components.logger = Mock()
  582. @self.app.task(shared=False)
  583. def foo_task(x, y, z):
  584. return x * y * z
  585. self.foo_task = foo_task
  586. def teardown(self):
  587. worker_module.logger = self._logger
  588. components.logger = self._comp_logger
  589. def create_worker(self, **kw):
  590. worker = self.app.WorkController(concurrency=1, loglevel=0, **kw)
  591. worker.blueprint.shutdown_complete.set()
  592. return worker
  593. def test_on_consumer_ready(self):
  594. self.worker.on_consumer_ready(Mock())
  595. def test_setup_queues_worker_direct(self):
  596. self.app.conf.worker_direct = True
  597. self.app.amqp.__dict__['queues'] = Mock()
  598. self.worker.setup_queues({})
  599. self.app.amqp.queues.select_add.assert_called_with(
  600. worker_direct(self.worker.hostname),
  601. )
  602. def test_setup_queues__missing_queue(self):
  603. self.app.amqp.queues.select = Mock(name='select')
  604. self.app.amqp.queues.deselect = Mock(name='deselect')
  605. self.app.amqp.queues.select.side_effect = KeyError()
  606. self.app.amqp.queues.deselect.side_effect = KeyError()
  607. with pytest.raises(ImproperlyConfigured):
  608. self.worker.setup_queues('x,y', exclude='foo,bar')
  609. self.app.amqp.queues.select = Mock(name='select')
  610. with pytest.raises(ImproperlyConfigured):
  611. self.worker.setup_queues('x,y', exclude='foo,bar')
  612. def test_send_worker_shutdown(self):
  613. with patch('celery.signals.worker_shutdown') as ws:
  614. self.worker._send_worker_shutdown()
  615. ws.send.assert_called_with(sender=self.worker)
  616. @skip.todo('unstable test')
  617. def test_process_shutdown_on_worker_shutdown(self):
  618. from celery.concurrency.prefork import process_destructor
  619. from celery.concurrency.asynpool import Worker
  620. with patch('celery.signals.worker_process_shutdown') as ws:
  621. with patch('os._exit') as _exit:
  622. worker = Worker(None, None, on_exit=process_destructor)
  623. worker._do_exit(22, 3.1415926)
  624. ws.send.assert_called_with(
  625. sender=None, pid=22, exitcode=3.1415926,
  626. )
  627. _exit.assert_called_with(3.1415926)
  628. def test_process_task_revoked_release_semaphore(self):
  629. self.worker._quick_release = Mock()
  630. req = Mock()
  631. req.execute_using_pool.side_effect = TaskRevokedError
  632. self.worker._process_task(req)
  633. self.worker._quick_release.assert_called_with()
  634. delattr(self.worker, '_quick_release')
  635. self.worker._process_task(req)
  636. def test_shutdown_no_blueprint(self):
  637. self.worker.blueprint = None
  638. self.worker._shutdown()
  639. @patch('celery.worker.worker.create_pidlock')
  640. def test_use_pidfile(self, create_pidlock):
  641. create_pidlock.return_value = Mock()
  642. worker = self.create_worker(pidfile='pidfilelockfilepid')
  643. worker.steps = []
  644. worker.start()
  645. create_pidlock.assert_called()
  646. worker.stop()
  647. worker.pidlock.release.assert_called()
  648. def test_attrs(self):
  649. worker = self.worker
  650. assert worker.timer is not None
  651. assert isinstance(worker.timer, Timer)
  652. assert worker.pool is not None
  653. assert worker.consumer is not None
  654. assert worker.steps
  655. def test_with_embedded_beat(self):
  656. worker = self.app.WorkController(concurrency=1, loglevel=0, beat=True)
  657. assert worker.beat
  658. assert worker.beat in [w.obj for w in worker.steps]
  659. def test_with_autoscaler(self):
  660. worker = self.create_worker(
  661. autoscale=[10, 3], send_events=False,
  662. timer_cls='celery.utils.timer2.Timer',
  663. )
  664. assert worker.autoscaler
  665. def test_dont_stop_or_terminate(self):
  666. worker = self.app.WorkController(concurrency=1, loglevel=0)
  667. worker.stop()
  668. assert worker.blueprint.state != CLOSE
  669. worker.terminate()
  670. assert worker.blueprint.state != CLOSE
  671. sigsafe, worker.pool.signal_safe = worker.pool.signal_safe, False
  672. try:
  673. worker.blueprint.state = RUN
  674. worker.stop(in_sighandler=True)
  675. assert worker.blueprint.state != CLOSE
  676. worker.terminate(in_sighandler=True)
  677. assert worker.blueprint.state != CLOSE
  678. finally:
  679. worker.pool.signal_safe = sigsafe
  680. def test_on_timer_error(self):
  681. worker = self.app.WorkController(concurrency=1, loglevel=0)
  682. try:
  683. raise KeyError('foo')
  684. except KeyError as exc:
  685. components.Timer(worker).on_timer_error(exc)
  686. msg, args = self.comp_logger.error.call_args[0]
  687. assert 'KeyError' in msg % args
  688. def test_on_timer_tick(self):
  689. worker = self.app.WorkController(concurrency=1, loglevel=10)
  690. components.Timer(worker).on_timer_tick(30.0)
  691. xargs = self.comp_logger.debug.call_args[0]
  692. fmt, arg = xargs[0], xargs[1]
  693. assert arg == 30.0
  694. assert 'Next ETA %s secs' in fmt
  695. def test_process_task(self):
  696. worker = self.worker
  697. worker.pool = Mock()
  698. channel = Mock()
  699. m = self.create_task_message(
  700. channel, self.foo_task.name,
  701. args=[4, 8, 10], kwargs={},
  702. )
  703. task = Request(m, app=self.app)
  704. worker._process_task(task)
  705. assert worker.pool.apply_async.call_count == 1
  706. worker.pool.stop()
  707. def test_process_task_raise_base(self):
  708. worker = self.worker
  709. worker.pool = Mock()
  710. worker.pool.apply_async.side_effect = KeyboardInterrupt('Ctrl+C')
  711. channel = Mock()
  712. m = self.create_task_message(
  713. channel, self.foo_task.name,
  714. args=[4, 8, 10], kwargs={},
  715. )
  716. task = Request(m, app=self.app)
  717. worker.steps = []
  718. worker.blueprint.state = RUN
  719. with pytest.raises(KeyboardInterrupt):
  720. worker._process_task(task)
  721. def test_process_task_raise_WorkerTerminate(self):
  722. worker = self.worker
  723. worker.pool = Mock()
  724. worker.pool.apply_async.side_effect = WorkerTerminate()
  725. channel = Mock()
  726. m = self.create_task_message(
  727. channel, self.foo_task.name,
  728. args=[4, 8, 10], kwargs={},
  729. )
  730. task = Request(m, app=self.app)
  731. worker.steps = []
  732. worker.blueprint.state = RUN
  733. with pytest.raises(SystemExit):
  734. worker._process_task(task)
  735. def test_process_task_raise_regular(self):
  736. worker = self.worker
  737. worker.pool = Mock()
  738. worker.pool.apply_async.side_effect = KeyError('some exception')
  739. channel = Mock()
  740. m = self.create_task_message(
  741. channel, self.foo_task.name,
  742. args=[4, 8, 10], kwargs={},
  743. )
  744. task = Request(m, app=self.app)
  745. worker._process_task(task)
  746. worker.pool.stop()
  747. def test_start_catches_base_exceptions(self):
  748. worker1 = self.create_worker()
  749. worker1.blueprint.state = RUN
  750. stc = MockStep()
  751. stc.start.side_effect = WorkerTerminate()
  752. worker1.steps = [stc]
  753. worker1.start()
  754. stc.start.assert_called_with(worker1)
  755. assert stc.terminate.call_count
  756. worker2 = self.create_worker()
  757. worker2.blueprint.state = RUN
  758. sec = MockStep()
  759. sec.start.side_effect = WorkerShutdown()
  760. sec.terminate = None
  761. worker2.steps = [sec]
  762. worker2.start()
  763. assert sec.stop.call_count
  764. def test_statedb(self):
  765. from celery.worker import state
  766. Persistent = state.Persistent
  767. state.Persistent = Mock()
  768. try:
  769. worker = self.create_worker(statedb='statefilename')
  770. assert worker._persistence
  771. finally:
  772. state.Persistent = Persistent
  773. def test_process_task_sem(self):
  774. worker = self.worker
  775. worker._quick_acquire = Mock()
  776. req = Mock()
  777. worker._process_task_sem(req)
  778. worker._quick_acquire.assert_called_with(worker._process_task, req)
  779. def test_signal_consumer_close(self):
  780. worker = self.worker
  781. worker.consumer = Mock()
  782. worker.signal_consumer_close()
  783. worker.consumer.close.assert_called_with()
  784. worker.consumer.close.side_effect = AttributeError()
  785. worker.signal_consumer_close()
  786. def test_rusage__no_resource(self):
  787. from celery.worker import worker
  788. prev, worker.resource = worker.resource, None
  789. try:
  790. self.worker.pool = Mock(name='pool')
  791. with pytest.raises(NotImplementedError):
  792. self.worker.rusage()
  793. self.worker.stats()
  794. finally:
  795. worker.resource = prev
  796. def test_repr(self):
  797. assert repr(self.worker)
  798. def test_str(self):
  799. assert str(self.worker) == self.worker.hostname
  800. def test_start__stop(self):
  801. worker = self.worker
  802. worker.blueprint.shutdown_complete.set()
  803. worker.steps = [MockStep(StartStopStep(self)) for _ in range(4)]
  804. worker.blueprint.state = RUN
  805. worker.blueprint.started = 4
  806. for w in worker.steps:
  807. w.start = Mock()
  808. w.close = Mock()
  809. w.stop = Mock()
  810. worker.start()
  811. for w in worker.steps:
  812. w.start.assert_called()
  813. worker.consumer = Mock()
  814. worker.stop(exitcode=3)
  815. for stopstep in worker.steps:
  816. stopstep.close.assert_called()
  817. stopstep.stop.assert_called()
  818. # Doesn't close pool if no pool.
  819. worker.start()
  820. worker.pool = None
  821. worker.stop()
  822. # test that stop of None is not attempted
  823. worker.steps[-1] = None
  824. worker.start()
  825. worker.stop()
  826. def test_start__KeyboardInterrupt(self):
  827. worker = self.worker
  828. worker.blueprint = Mock(name='blueprint')
  829. worker.blueprint.start.side_effect = KeyboardInterrupt()
  830. worker.stop = Mock(name='stop')
  831. worker.start()
  832. worker.stop.assert_called_with(exitcode=EX_FAILURE)
  833. def test_register_with_event_loop(self):
  834. worker = self.worker
  835. hub = Mock(name='hub')
  836. worker.blueprint = Mock(name='blueprint')
  837. worker.register_with_event_loop(hub)
  838. worker.blueprint.send_all.assert_called_with(
  839. worker, 'register_with_event_loop', args=(hub,),
  840. description='hub.register',
  841. )
  842. def test_step_raises(self):
  843. worker = self.worker
  844. step = Mock()
  845. worker.steps = [step]
  846. step.start.side_effect = TypeError()
  847. worker.stop = Mock()
  848. worker.start()
  849. worker.stop.assert_called_with(exitcode=EX_FAILURE)
  850. def test_state(self):
  851. assert self.worker.state
  852. def test_start__terminate(self):
  853. worker = self.worker
  854. worker.blueprint.shutdown_complete.set()
  855. worker.blueprint.started = 5
  856. worker.blueprint.state = RUN
  857. worker.steps = [MockStep() for _ in range(5)]
  858. worker.start()
  859. for w in worker.steps[:3]:
  860. w.start.assert_called()
  861. assert worker.blueprint.started == len(worker.steps)
  862. assert worker.blueprint.state == RUN
  863. worker.terminate()
  864. for step in worker.steps:
  865. step.terminate.assert_called()
  866. worker.blueprint.state = TERMINATE
  867. worker.terminate()
  868. def test_Hub_create(self):
  869. w = Mock()
  870. x = components.Hub(w)
  871. x.create(w)
  872. assert w.timer.max_interval
  873. def test_Pool_create_threaded(self):
  874. w = Mock()
  875. w._conninfo.connection_errors = w._conninfo.channel_errors = ()
  876. w.pool_cls = Mock()
  877. w.use_eventloop = False
  878. pool = components.Pool(w)
  879. pool.create(w)
  880. def test_Pool_pool_no_sem(self):
  881. w = Mock()
  882. w.pool_cls.uses_semaphore = False
  883. components.Pool(w).create(w)
  884. assert w.process_task is w._process_task
  885. def test_Pool_create(self):
  886. from kombu.async.semaphore import LaxBoundedSemaphore
  887. w = Mock()
  888. w._conninfo.connection_errors = w._conninfo.channel_errors = ()
  889. w.hub = Mock()
  890. PoolImp = Mock()
  891. poolimp = PoolImp.return_value = Mock()
  892. poolimp._pool = [Mock(), Mock()]
  893. poolimp._cache = {}
  894. poolimp._fileno_to_inq = {}
  895. poolimp._fileno_to_outq = {}
  896. from celery.concurrency.prefork import TaskPool as _TaskPool
  897. class MockTaskPool(_TaskPool):
  898. Pool = PoolImp
  899. @property
  900. def timers(self):
  901. return {Mock(): 30}
  902. w.pool_cls = MockTaskPool
  903. w.use_eventloop = True
  904. w.consumer.restart_count = -1
  905. pool = components.Pool(w)
  906. pool.create(w)
  907. pool.register_with_event_loop(w, w.hub)
  908. if sys.platform != 'win32':
  909. assert isinstance(w.semaphore, LaxBoundedSemaphore)
  910. P = w.pool
  911. P.start()