test_worker.py 34 KB

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