test_worker.py 34 KB

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