mocks.py 4.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172
  1. from __future__ import absolute_import, unicode_literals
  2. from itertools import count
  3. from case import ContextMock, Mock
  4. from kombu.transport import base
  5. from kombu.utils import json
  6. def PromiseMock(*args, **kwargs):
  7. m = Mock(*args, **kwargs)
  8. def on_throw(exc=None, *args, **kwargs):
  9. if exc:
  10. raise exc
  11. raise
  12. m.throw.side_effect = on_throw
  13. m.set_error_state.side_effect = on_throw
  14. m.throw1.side_effect = on_throw
  15. return m
  16. class MockPool(object):
  17. def __init__(self, value=None):
  18. self.value = value or ContextMock()
  19. def acquire(self, **kwargs):
  20. return self.value
  21. class Message(base.Message):
  22. def __init__(self, *args, **kwargs):
  23. self.throw_decode_error = kwargs.get('throw_decode_error', False)
  24. super(Message, self).__init__(*args, **kwargs)
  25. def decode(self):
  26. if self.throw_decode_error:
  27. raise ValueError("can't decode message")
  28. return super(Message, self).decode()
  29. class Channel(base.StdChannel):
  30. open = True
  31. throw_decode_error = False
  32. _ids = count(1)
  33. def __init__(self, connection):
  34. self.connection = connection
  35. self.called = []
  36. self.deliveries = count(1)
  37. self.to_deliver = []
  38. self.events = {'basic_return': set()}
  39. self.channel_id = next(self._ids)
  40. def _called(self, name):
  41. self.called.append(name)
  42. def __contains__(self, key):
  43. return key in self.called
  44. def exchange_declare(self, *args, **kwargs):
  45. self._called('exchange_declare')
  46. def prepare_message(self, body, priority=0, content_type=None,
  47. content_encoding=None, headers=None, properties={}):
  48. self._called('prepare_message')
  49. return {'body': body,
  50. 'headers': headers,
  51. 'properties': properties,
  52. 'priority': priority,
  53. 'content_type': content_type,
  54. 'content_encoding': content_encoding}
  55. def basic_publish(self, message, exchange='', routing_key='',
  56. mandatory=False, immediate=False, **kwargs):
  57. self._called('basic_publish')
  58. return message, exchange, routing_key
  59. def exchange_delete(self, *args, **kwargs):
  60. self._called('exchange_delete')
  61. def queue_declare(self, *args, **kwargs):
  62. self._called('queue_declare')
  63. def queue_bind(self, *args, **kwargs):
  64. self._called('queue_bind')
  65. def queue_unbind(self, *args, **kwargs):
  66. self._called('queue_unbind')
  67. def queue_delete(self, queue, if_unused=False, if_empty=False, **kwargs):
  68. self._called('queue_delete')
  69. def basic_get(self, *args, **kwargs):
  70. self._called('basic_get')
  71. try:
  72. return self.to_deliver.pop()
  73. except IndexError:
  74. pass
  75. def queue_purge(self, *args, **kwargs):
  76. self._called('queue_purge')
  77. def basic_consume(self, *args, **kwargs):
  78. self._called('basic_consume')
  79. def basic_cancel(self, *args, **kwargs):
  80. self._called('basic_cancel')
  81. def basic_ack(self, *args, **kwargs):
  82. self._called('basic_ack')
  83. def basic_recover(self, requeue=False):
  84. self._called('basic_recover')
  85. def exchange_bind(self, *args, **kwargs):
  86. self._called('exchange_bind')
  87. def exchange_unbind(self, *args, **kwargs):
  88. self._called('exchange_unbind')
  89. def close(self):
  90. self._called('close')
  91. def message_to_python(self, message, *args, **kwargs):
  92. self._called('message_to_python')
  93. return Message(body=json.dumps(message),
  94. channel=self,
  95. delivery_tag=next(self.deliveries),
  96. throw_decode_error=self.throw_decode_error,
  97. content_type='application/json',
  98. content_encoding='utf-8')
  99. def flow(self, active):
  100. self._called('flow')
  101. def basic_reject(self, delivery_tag, requeue=False):
  102. if requeue:
  103. return self._called('basic_reject:requeue')
  104. return self._called('basic_reject')
  105. def basic_qos(self, prefetch_size=0, prefetch_count=0,
  106. apply_global=False):
  107. self._called('basic_qos')
  108. class Connection(object):
  109. connected = True
  110. def __init__(self, client):
  111. self.client = client
  112. def channel(self):
  113. return Channel(self)
  114. class Transport(base.Transport):
  115. def establish_connection(self):
  116. return Connection(self.client)
  117. def create_channel(self, connection):
  118. return connection.channel()
  119. def drain_events(self, connection, **kwargs):
  120. return 'event'
  121. def close_connection(self, connection):
  122. connection.connected = False