utils.py 7.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225
  1. """
  2. Useful testing utilities.
  3. This module is considered public API.
  4. """
  5. import time
  6. import six
  7. import socket
  8. from library.pymemcache.exceptions import MemcacheClientError, MemcacheIllegalInputError
  9. from library.pymemcache.serde import LegacyWrappingSerde
  10. from library.pymemcache.client.base import check_key_helper
  11. class MockMemcacheClient(object):
  12. """
  13. A (partial) in-memory mock for Clients.
  14. """
  15. def __init__(self,
  16. server=None,
  17. serde=None,
  18. serializer=None,
  19. deserializer=None,
  20. connect_timeout=None,
  21. timeout=None,
  22. no_delay=False,
  23. ignore_exc=False,
  24. socket_module=None,
  25. default_noreply=True,
  26. allow_unicode_keys=False,
  27. encoding='ascii',
  28. tls_context=None):
  29. self._contents = {}
  30. self.serde = serde or LegacyWrappingSerde(serializer, deserializer)
  31. self.allow_unicode_keys = allow_unicode_keys
  32. # Unused, but present for interface compatibility
  33. self.server = server
  34. self.connect_timeout = connect_timeout
  35. self.timeout = timeout
  36. self.no_delay = no_delay
  37. self.ignore_exc = ignore_exc
  38. self.socket_module = socket
  39. self.sock = None
  40. self.encoding = encoding
  41. self.tls_context = tls_context
  42. def check_key(self, key):
  43. """Checks key and add key_prefix."""
  44. return check_key_helper(key, allow_unicode_keys=self.allow_unicode_keys)
  45. def clear(self):
  46. """Method used to clear/reset mock cache"""
  47. self._contents.clear()
  48. def get(self, key, default=None):
  49. key = self.check_key(key)
  50. if key not in self._contents:
  51. return default
  52. expire, value, flags = self._contents[key]
  53. if expire and expire < time.time():
  54. del self._contents[key]
  55. return default
  56. return self.serde.deserialize(key, value, flags)
  57. def get_many(self, keys):
  58. out = {}
  59. for key in keys:
  60. value = self.get(key)
  61. if value is not None:
  62. out[key] = value
  63. return out
  64. get_multi = get_many
  65. def set(self, key, value, expire=0, noreply=True, flags=None):
  66. key = self.check_key(key)
  67. if (isinstance(value, six.string_types) and
  68. not isinstance(value, six.binary_type)):
  69. try:
  70. value = value.encode(self.encoding)
  71. except (UnicodeEncodeError, UnicodeDecodeError):
  72. raise MemcacheIllegalInputError
  73. value, flags = self.serde.serialize(key, value)
  74. if expire:
  75. expire += time.time()
  76. self._contents[key] = expire, value, flags
  77. return True
  78. def set_many(self, values, expire=0, noreply=True, flags=None):
  79. result = []
  80. for key, value in six.iteritems(values):
  81. ret = self.set(key, value, expire, noreply, flags=flags)
  82. if not ret:
  83. result.append(key)
  84. return [] if noreply else result
  85. set_multi = set_many
  86. def incr(self, key, value, noreply=False):
  87. current = self.get(key)
  88. present = current is not None
  89. if present:
  90. self.set(key, current + value, noreply=noreply)
  91. return None if noreply or not present else current + value
  92. def decr(self, key, value, noreply=False):
  93. current = self.get(key)
  94. present = current is not None
  95. if present:
  96. self.set(key, current - value, noreply=noreply)
  97. return None if noreply or not present else current - value
  98. def add(self, key, value, expire=0, noreply=True, flags=None):
  99. current = self.get(key)
  100. present = current is not None
  101. if not present:
  102. self.set(key, value, expire, noreply, flags=flags)
  103. return noreply or not present
  104. def delete(self, key, noreply=True):
  105. key = self.check_key(key)
  106. current = self._contents.pop(key, None)
  107. present = current is not None
  108. return noreply or present
  109. def delete_many(self, keys, noreply=True):
  110. for key in keys:
  111. self.delete(key, noreply)
  112. return True
  113. def prepend(self, key, value, expire=0, noreply=True, flags=None):
  114. current = self.get(key)
  115. if current is not None:
  116. if (isinstance(value, six.string_types) and
  117. not isinstance(value, six.binary_type)):
  118. try:
  119. value = value.encode(self.encoding)
  120. except (UnicodeEncodeError, UnicodeDecodeError):
  121. raise MemcacheIllegalInputError
  122. self.set(key, value + current, expire, noreply, flags=flags)
  123. return True
  124. def append(self, key, value, expire=0, noreply=True, flags=None):
  125. current = self.get(key)
  126. if current is not None:
  127. if (isinstance(value, six.string_types) and
  128. not isinstance(value, six.binary_type)):
  129. try:
  130. value = value.encode(self.encoding)
  131. except (UnicodeEncodeError, UnicodeDecodeError):
  132. raise MemcacheIllegalInputError
  133. self.set(key, current + value, expire, noreply, flags=flags)
  134. return True
  135. delete_multi = delete_many
  136. def stats(self, *_args):
  137. # I make no claim that these values make any sense, but the format
  138. # of the output is the same as for pymemcache.client.Client.stats()
  139. return {
  140. "version": "MockMemcacheClient",
  141. "rusage_user": 1.0,
  142. "rusage_system": 1.0,
  143. "hash_is_expanding": False,
  144. "slab_reassign_running": False,
  145. "inter": "in-memory",
  146. "evictions": False,
  147. "growth_factor": 1.0,
  148. "stat_key_prefix": "",
  149. "umask": 0o644,
  150. "detail_enabled": False,
  151. "cas_enabled": False,
  152. "auth_enabled_sasl": False,
  153. "maxconns_fast": False,
  154. "slab_reassign": False,
  155. "slab_automove": False,
  156. }
  157. def replace(self, key, value, expire=0, noreply=True, flags=None):
  158. current = self.get(key)
  159. present = current is not None
  160. if present:
  161. self.set(key, value, expire, noreply, flags=flags)
  162. return noreply or present
  163. def cas(self, key, value, cas, expire=0, noreply=False, flags=None):
  164. raise MemcacheClientError('CAS is not enabled for this instance')
  165. def touch(self, key, expire=0, noreply=True):
  166. current = self.get(key)
  167. present = current is not None
  168. if present:
  169. self.set(key, current, expire, noreply=noreply)
  170. return True if noreply or present else False
  171. def cache_memlimit(self, memlimit):
  172. return True
  173. def version(self):
  174. return 'MockMemcacheClient'
  175. def flush_all(self, delay=0, noreply=True):
  176. self.clear()
  177. return noreply or self._contents == {}
  178. def quit(self):
  179. pass
  180. def close(self):
  181. pass