kaya-rbcs/tests/fake_memcached.py

"""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