Repository: aio-libs/aiomcache

Issue:
Encryption via TLS
Will I be able to use TLS encryption in future releases?

Discussion:

Patch:
diff --git a/aiomcache/client.py b/aiomcache/client.py
index d31232d..1b61a23 100644
--- a/aiomcache/client.py
+++ b/aiomcache/client.py
@@ -1,8 +1,8 @@
 import functools
 import re
 import sys
-from typing import (Any, Awaitable, Callable, Dict, Generic, Optional, Tuple, TypeVar,
-                    Union, overload)
+from typing import (Any, Awaitable, Callable, Dict, Generic, Mapping, Optional, Tuple,
+                    TypeVar, Union, overload)
 
 from . import constants as const
 from .exceptions import ClientException, ValidationException
@@ -52,6 +52,7 @@ def acquire(
 class FlagClient(Generic[_T]):
     def __init__(self, host: str, port: int = 11211, *,
                  pool_size: int = 2, pool_minsize: Optional[int] = None,
+                 conn_args: Optional[Mapping[str, Any]] = None,
                  get_flag_handler: Optional[_GetFlagHandler[_T]] = None,
                  set_flag_handler: Optional[_SetFlagHandler[_T]] = None):
         """
@@ -61,6 +62,9 @@ class FlagClient(Generic[_T]):
         :param port: memcached port
         :param pool_size: max connection pool size
         :param pool_minsize: min connection pool size
+        :param conn_args: extra arguments passed to
+            asyncio.open_connection(). For details, see:
+            https://docs.python.org/3/library/asyncio-stream.html#asyncio.open_connection.
         :param get_flag_handler: async method to call to convert flagged
             values. Method takes tuple: (value, flags) and should return
             processed value or raise ClientException if not supported.
@@ -72,7 +76,8 @@ class FlagClient(Generic[_T]):
             pool_minsize = pool_size
 
         self._pool = MemcachePool(
-            host, port, minsize=pool_minsize, maxsize=pool_size)
+            host, port, minsize=pool_minsize, maxsize=pool_size,
+            conn_args=conn_args)
 
         self._get_flag_handler = get_flag_handler
         self._set_flag_handler = set_flag_handler
@@ -493,6 +498,8 @@ class FlagClient(Generic[_T]):
 
 class Client(FlagClient[bytes]):
     def __init__(self, host: str, port: int = 11211, *,
-                 pool_size: int = 2, pool_minsize: Optional[int] = None):
+                 pool_size: int = 2, pool_minsize: Optional[int] = None,
+                 conn_args: Optional[Mapping[str, Any]] = None):
         super().__init__(host, port, pool_size=pool_size, pool_minsize=pool_minsize,
+                         conn_args=conn_args,
                          get_flag_handler=None, set_flag_handler=None)
diff --git a/aiomcache/pool.py b/aiomcache/pool.py
index f9eb5a6..edee644 100644
--- a/aiomcache/pool.py
+++ b/aiomcache/pool.py
@@ -1,5 +1,5 @@
 import asyncio
-from typing import NamedTuple, Optional, Set
+from typing import Any, Mapping, NamedTuple, Optional, Set
 
 __all__ = ['MemcachePool']
 
@@ -10,11 +10,13 @@ class Connection(NamedTuple):
 
 
 class MemcachePool:
-    def __init__(self, host: str, port: int, *, minsize: int, maxsize: int):
+    def __init__(self, host: str, port: int, *, minsize: int, maxsize: int,
+                 conn_args: Optional[Mapping[str, Any]] = None):
         self._host = host
         self._port = port
         self._minsize = minsize
         self._maxsize = maxsize
+        self.conn_args = conn_args or {}
         self._pool: asyncio.Queue[Connection] = asyncio.Queue()
         self._in_use: Set[Connection] = set()
 
@@ -66,7 +68,7 @@ class MemcachePool:
     async def _create_new_conn(self) -> Optional[Connection]:
         if self.size() < self._maxsize:
             reader, writer = await asyncio.open_connection(
-                self._host, self._port)
+                self._host, self._port, **self.conn_args)
             if self.size() < self._maxsize:
                 return Connection(reader, writer)
             else:


Test patch:
diff --git a/tests/commands_test.py b/tests/commands_test.py
index bb41b41..ae5768e 100644
--- a/tests/commands_test.py
+++ b/tests/commands_test.py
@@ -19,7 +19,6 @@ from .flag_helper import FlagHelperDemo
     bytes("こんにちは", "utf-8"),
     bytes("안녕하세요", "utf-8"),
 ))
-@pytest.mark.asyncio
 async def test_valid_key(mcache: Client, key: bytes) -> None:
     assert mcache._validate_key(key) == key
 
@@ -41,13 +40,11 @@ async def test_valid_key(mcache: Client, key: bytes) -> None:
     "\u0080".encode(),
     "\u009F".encode(),
 ))
-@pytest.mark.asyncio
 async def test_invalid_key(mcache: Client, key: bytes) -> None:
     with pytest.raises(ValidationException, match="invalid key"):
         mcache._validate_key(key)
 
 
-@pytest.mark.asyncio
 async def test_version(mcache: Client) -> None:
     version = await mcache.version()
     stats = await mcache.stats()
@@ -64,7 +61,6 @@ async def test_version(mcache: Client) -> None:
             await mcache.version()
 
 
-@pytest.mark.asyncio
 async def test_flush_all(mcache: Client) -> None:
     key, value = b'key:flush_all', b'flush_all_value'
     await mcache.set(key, value)
@@ -85,7 +81,6 @@ async def test_flush_all(mcache: Client) -> None:
             await mcache.flush_all()
 
 
-@pytest.mark.asyncio
 async def test_set_get(mcache: Client) -> None:
     key, value = b'key:set', b'1'
     await mcache.set(key, value)
@@ -104,7 +99,6 @@ async def test_set_get(mcache: Client) -> None:
             await mcache.set(key, value)
 
 
-@pytest.mark.asyncio
 async def test_gets(mcache: Client) -> None:
     key, value = b'key:set', b'1'
     await mcache.set(key, value)
@@ -122,7 +116,6 @@ async def test_gets(mcache: Client) -> None:
     assert cas is None
 
 
-@pytest.mark.asyncio
 async def test_multi_get(mcache: Client) -> None:
     key1, value1 = b'key:multi_get:1', b'1'
     key2, value2 = b'key:multi_get:2', b'2'
@@ -137,7 +130,6 @@ async def test_multi_get(mcache: Client) -> None:
     assert test_value == ()
 
 
-@pytest.mark.asyncio
 async def test_multi_get_doubling_keys(mcache: Client) -> None:
     key, value = b'key:multi_get:3', b'1'
     await mcache.set(key, value)
@@ -146,7 +138,6 @@ async def test_multi_get_doubling_keys(mcache: Client) -> None:
         await mcache.multi_get(key, key)
 
 
-@pytest.mark.asyncio
 async def test_set_expire(mcache: Client) -> None:
     key, value = b'key:set', b'1'
     await mcache.set(key, value, exptime=1)
@@ -159,7 +150,6 @@ async def test_set_expire(mcache: Client) -> None:
     assert test_value is None
 
 
-@pytest.mark.asyncio
 async def test_set_errors(mcache: Client) -> None:
     key, value = b'key:set', b'1'
     await mcache.set(key, value, exptime=1)
@@ -171,7 +161,6 @@ async def test_set_errors(mcache: Client) -> None:
         await mcache.set(key, value, exptime=3.14)  # type: ignore[arg-type]
 
 
-@pytest.mark.asyncio
 async def test_gets_cas(mcache: Client) -> None:
     key, value = b'key:set', b'1'
     await mcache.set(key, value)
@@ -187,14 +176,12 @@ async def test_gets_cas(mcache: Client) -> None:
     assert stored is False
 
 
-@pytest.mark.asyncio
 async def test_cas_missing(mcache: Client) -> None:
     key, value = b'key:set', b'1'
     stored = await mcache.cas(key, value, 123)
     assert stored is False
 
 
-@pytest.mark.asyncio
 async def test_add(mcache: Client) -> None:
     key, value = b'key:add', b'1'
     await mcache.set(key, value)
@@ -209,7 +196,6 @@ async def test_add(mcache: Client) -> None:
     assert test_value3 == b"2"
 
 
-@pytest.mark.asyncio
 async def test_replace(mcache: Client) -> None:
     key, value = b'key:replace', b'1'
     await mcache.set(key, value)
@@ -227,7 +213,6 @@ async def test_replace(mcache: Client) -> None:
     assert test_value4 is None
 
 
-@pytest.mark.asyncio
 async def test_append(mcache: Client) -> None:
     key, value = b'key:append', b'1'
     await mcache.set(key, value)
@@ -246,7 +231,6 @@ async def test_append(mcache: Client) -> None:
     assert test_value4 is None
 
 
-@pytest.mark.asyncio
 async def test_prepend(mcache: Client) -> None:
     key, value = b'key:prepend', b'1'
     await mcache.set(key, value)
@@ -265,7 +249,6 @@ async def test_prepend(mcache: Client) -> None:
     assert test_value4 is None
 
 
-@pytest.mark.asyncio
 async def test_delete(mcache: Client) -> None:
     key, value = b'key:delete', b'value'
     await mcache.set(key, value)
@@ -289,13 +272,11 @@ async def test_delete(mcache: Client) -> None:
             await mcache.delete(key)
 
 
-@pytest.mark.asyncio
 async def test_delete_key_not_exists(mcache: Client) -> None:
     is_deleted = await mcache.delete(b"not:key")
     assert not is_deleted
 
 
-@pytest.mark.asyncio
 async def test_incr(mcache: Client) -> None:
     key, value = b'key:incr:1', b'1'
     await mcache.set(key, value)
@@ -308,7 +289,6 @@ async def test_incr(mcache: Client) -> None:
     assert test_value2 == b"3"
 
 
-@pytest.mark.asyncio
 async def test_incr_errors(mcache: Client) -> None:
     key, value = b'key:incr:2', b'string'
     await mcache.set(key, value)
@@ -320,7 +300,6 @@ async def test_incr_errors(mcache: Client) -> None:
         await mcache.incr(key, 3.14)  # type: ignore[arg-type]
 
 
-@pytest.mark.asyncio
 async def test_decr(mcache: Client) -> None:
     key, value = b'key:decr:1', b'17'
     await mcache.set(key, value)
@@ -335,7 +314,6 @@ async def test_decr(mcache: Client) -> None:
     assert test_value3 == 0
 
 
-@pytest.mark.asyncio
 async def test_decr_errors(mcache: Client) -> None:
     key, value = b'key:decr:2', b'string'
     await mcache.set(key, value)
@@ -347,13 +325,11 @@ async def test_decr_errors(mcache: Client) -> None:
         await mcache.decr(key, 3.14)  # type: ignore[arg-type]
 
 
-@pytest.mark.asyncio
 async def test_stats(mcache: Client) -> None:
     stats = await mcache.stats()
     assert b'pid' in stats
 
 
-@pytest.mark.asyncio
 async def test_touch(mcache: Client) -> None:
     key, value = b'key:touch:1', b'17'
     await mcache.set(key, value)
@@ -381,7 +357,6 @@ async def test_touch(mcache: Client) -> None:
             await mcache.touch(b"not:" + key, 1)
 
 
-@pytest.mark.asyncio
 async def test_close(mcache: Client) -> None:
     await mcache.close()
     assert mcache._pool.size() == 0
@@ -403,7 +378,6 @@ async def test_close(mcache: Client) -> None:
         bytes("안녕하세요", "utf-8"),
     ]
 )
-@pytest.mark.asyncio
 async def test_flag_helper(
         mcache_flag_client: FlagClient[Any], value: object) -> None:
     key = b"key:test_flag_helper"
@@ -413,7 +387,6 @@ async def test_flag_helper(
     assert v2 == value
 
 
-@pytest.mark.asyncio
 async def test_objects_not_supported_without_flag_handler(mcache: Client) -> None:
     key = b"key:test_objects_not_supported_without_flag_handler"
 
@@ -426,7 +399,6 @@ async def test_objects_not_supported_without_flag_handler(mcache: Client) -> Non
     assert result is None
 
 
-@pytest.mark.asyncio
 async def test_flag_handler_invoked_only_when_expected(
         mcache_flag_client: FlagClient[Any], demo_flag_helper: FlagHelperDemo) -> None:
     key = b"key:test_flag_handler_invoked_only_when_expected"
diff --git a/tests/conn_args_test.py b/tests/conn_args_test.py
new file mode 100644
index 0000000..f71e1f9
--- /dev/null
+++ b/tests/conn_args_test.py
@@ -0,0 +1,39 @@
+import ssl
+import sys
+from asyncio import StreamReader, StreamWriter
+from unittest import mock
+
+import pytest
+
+from aiomcache import Client
+from .conftest import McacheParams
+
+
+@pytest.mark.skipif(sys.version_info < (3, 8), reason="AsyncMock requires python3.8")
+async def test_params_forwarded_from_client() -> None:
+    client = Client("host", port=11211, conn_args={
+        "ssl": True, "ssl_handshake_timeout": 20
+    })
+
+    with mock.patch(
+        "asyncio.open_connection",
+        return_value=(
+            mock.create_autospec(StreamReader),
+            mock.create_autospec(StreamWriter),
+        ),
+        autospec=True,
+    ) as oc:
+        await client._pool.acquire()
+
+    oc.assert_called_with("host", 11211, ssl=True, ssl_handshake_timeout=20)
+
+
+async def test_ssl_client_fails_against_plaintext_server(
+    mcache_params: McacheParams,
+) -> None:
+    client = Client(**mcache_params, conn_args={"ssl": True})
+    # If SSL was correctly enabled, this should
+    # fail, since SSL isn't enabled on the memcache
+    # server.
+    with pytest.raises(ssl.SSLError):
+        await client.get(b"key")
diff --git a/tests/pool_test.py b/tests/pool_test.py
index 75e5ac8..1487efd 100644
--- a/tests/pool_test.py
+++ b/tests/pool_test.py
@@ -8,14 +8,12 @@ from aiomcache.pool import Connection, MemcachePool
 from .conftest import McacheParams
 
 
-@pytest.mark.asyncio
 async def test_pool_creation(mcache_params: McacheParams) -> None:
     pool = MemcachePool(minsize=1, maxsize=5, **mcache_params)
     assert pool.size() == 0
     assert pool._minsize == 1
 
 
-@pytest.mark.asyncio
 async def test_pool_acquire_release(mcache_params: McacheParams) -> None:
     pool = MemcachePool(minsize=1, maxsize=5, **mcache_params)
     conn = await pool.acquire()
@@ -25,7 +23,6 @@ async def test_pool_acquire_release(mcache_params: McacheParams) -> None:
     await pool.clear()
 
 
-@pytest.mark.asyncio
 async def test_pool_acquire_release2(mcache_params: McacheParams) -> None:
     pool = MemcachePool(minsize=1, maxsize=5, **mcache_params)
     reader, writer = await asyncio.open_connection(
@@ -42,7 +39,6 @@ async def test_pool_acquire_release2(mcache_params: McacheParams) -> None:
     await pool.clear()
 
 
-@pytest.mark.asyncio
 async def test_pool_clear(mcache_params: McacheParams) -> None:
     pool = MemcachePool(minsize=1, maxsize=5, **mcache_params)
     conn = await pool.acquire()
@@ -52,7 +48,6 @@ async def test_pool_clear(mcache_params: McacheParams) -> None:
     assert pool._pool.qsize() == 0
 
 
-@pytest.mark.asyncio
 async def test_acquire_dont_create_new_connection_if_have_conn_in_pool(
     mcache_params: McacheParams,
 ) -> None:
@@ -72,7 +67,6 @@ async def test_acquire_dont_create_new_connection_if_have_conn_in_pool(
     await pool.clear()
 
 
-@pytest.mark.asyncio
 async def test_acquire_limit_maxsize(mcache_params: McacheParams) -> None:
     pool = MemcachePool(minsize=1, maxsize=1, **mcache_params)
     assert pool.size() == 0
@@ -98,7 +92,6 @@ async def test_acquire_limit_maxsize(mcache_params: McacheParams) -> None:
     await pool.clear()
 
 
-@pytest.mark.asyncio
 async def test_acquire_task_cancellation(mcache_params: McacheParams) -> None:
 
     class TestClient(Client):
@@ -127,25 +120,24 @@ async def test_acquire_task_cancellation(mcache_params: McacheParams) -> None:
     await client._pool.clear()
 
 
-@pytest.mark.asyncio
 async def test_maxsize_greater_than_minsize(mcache_params: McacheParams) -> None:
     pool = MemcachePool(minsize=5, maxsize=1, **mcache_params)
     conn = await pool.acquire()
     assert isinstance(conn.reader, asyncio.StreamReader)
     assert isinstance(conn.writer, asyncio.StreamWriter)
     pool.release(conn)
+    await pool.clear()
 
 
-@pytest.mark.asyncio
 async def test_0_minsize(mcache_params: McacheParams) -> None:
     pool = MemcachePool(minsize=0, maxsize=5, **mcache_params)
     conn = await pool.acquire()
     assert isinstance(conn.reader, asyncio.StreamReader)
     assert isinstance(conn.writer, asyncio.StreamWriter)
     pool.release(conn)
+    await pool.clear()
 
 
-@pytest.mark.asyncio
 async def test_bad_connection(mcache_params: McacheParams) -> None:
     pool = MemcachePool(minsize=5, maxsize=1, **mcache_params)
     pool._host = "INVALID_HOST"