"""A minimal in-process memcached server (text protocol) for offline tests. It implements the subset of the memcached text protocol used by ``aiomcache`` (and thus by this project): ``get``, ``set``/``add``/ ``replace``/``append``/``prepend``, ``touch``, ``delete``, ``version``, ``flush_all`` and ``quit``, including lazy expiration. """ import asyncio from time import time from typing import Dict, List, Optional, Tuple _THIRTY_DAYS = 30 * 24 * 60 * 60 class FakeMemcachedServer: def __init__(self) -> None: self._data: Dict[bytes, Tuple[int, bytes]] = {} self._server: Optional[asyncio.AbstractServer] = None self.port: int = 0 async def start(self) -> None: self._server = await asyncio.start_server(self._handle, '127.0.0.1', 0) assert self._server.sockets is not None self.port = self._server.sockets[0].getsockname()[1] async def stop(self) -> None: if self._server is not None: self._server.close() await self._server.wait_closed() @staticmethod def _expiry_timestamp(exptime: int, now: float) -> int: if exptime == 0: return 0 if exptime > _THIRTY_DAYS: return exptime return int(now) + exptime def _expired(self, key: bytes, now: float) -> bool: entry = self._data.get(key) if entry is None: return False expiry, _ = entry return 0 < expiry <= now async def _handle(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: try: while True: line = await reader.readline() if not line: break parts = line.rstrip(b'\r\n').split() if not parts: continue command = parts[0] if command == b'get': self._handle_get(parts[1:], writer) elif command in (b'set', b'add', b'replace', b'append', b'prepend'): await self._handle_store(command, parts, reader, writer) elif command == b'touch': self._handle_touch(parts, writer) elif command == b'delete': self._handle_delete(parts, writer) elif command == b'version': writer.write(b'VERSION 1.6.23\r\n') elif command == b'flush_all': self._data.clear() writer.write(b'OK\r\n') elif command == b'quit': break else: writer.write(b'ERROR\r\n') await writer.drain() finally: writer.close() await writer.wait_closed() def _handle_get(self, keys: List[bytes], writer: asyncio.StreamWriter) -> None: now = time() for key in keys: if key in self._data and not self._expired(key, now): _, value = self._data[key] writer.write( b'VALUE ' + key + b' 0 ' + str(len(value)).encode('utf-8') + b'\r\n' ) writer.write(value + b'\r\n') writer.write(b'END\r\n') async def _handle_store( self, command: bytes, parts: List[bytes], reader: asyncio.StreamReader, writer: asyncio.StreamWriter, ) -> None: key = parts[1] exptime = int(parts[3]) length = int(parts[4]) value = (await reader.readexactly(length + 2))[:-2] now = time() stored = False if command == b'set': self._data[key] = (self._expiry_timestamp(exptime, now), value) stored = True elif command == b'add': if not (key in self._data and not self._expired(key, now)): self._data[key] = (self._expiry_timestamp(exptime, now), value) stored = True elif command == b'replace': if key in self._data and not self._expired(key, now): self._data[key] = (self._expiry_timestamp(exptime, now), value) stored = True elif command in (b'append', b'prepend'): if key in self._data and not self._expired(key, now): expiry, existing = self._data[key] self._data[key] = ( expiry, existing + value if command == b'append' else value + existing, ) stored = True writer.write(b'STORED\r\n' if stored else b'NOT_STORED\r\n') def _handle_touch(self, parts: List[bytes], writer: asyncio.StreamWriter) -> None: key = parts[1] exptime = int(parts[2]) if key in self._data and not self._expired(key, time()): _, value = self._data[key] self._data[key] = (self._expiry_timestamp(exptime, time()), value) writer.write(b'TOUCHED\r\n') else: writer.write(b'NOT_FOUND\r\n') def _handle_delete(self, parts: List[bytes], writer: asyncio.StreamWriter) -> None: key = parts[1] if key in self._data: del self._data[key] writer.write(b'DELETED\r\n') else: writer.write(b'NOT_FOUND\r\n') @property def data(self) -> Dict[bytes, Tuple[int, bytes]]: return self._data