Fix ClientSession.close() was never awaited (#39)

Author: ciscorn

.gitignore

@@ -7,6 +7,7 @@ dist
 build
 *.egg
 coverage.xml
+.pytest_cache
 junit.xml
 test_elasticsearch_async/htmlcov
 docs/_build

README

@@ -21,8 +21,8 @@ Example::
 
  loop = asyncio.get_event_loop()
  loop.run_until_complete(print_info())
+ loop.run_until_complete(client.transport.close())
  loop.close()
- client.transport.close()
 
 
 Example with SSL Context::
@@ -46,8 +46,8 @@ Example with SSL Context::
 
  loop = asyncio.get_event_loop()
  loop.run_until_complete(print_info())
+ loop.run_until_complete(client.transport.close())
  loop.close()
- client.transport.close()
 
 
 

elasticsearch_async/connection.py

@@ -80,8 +80,9 @@ def __init__(self, host='localhost', port=9200, http_auth=None,
  host, port, self.url_prefix
  )
 
+ @asyncio.coroutine
  def close(self):
- return self.session.close()
+ yield from self.session.close()
 
  @asyncio.coroutine
  def perform_request(self, method, url, params=None, body=None, timeout=None, ignore=(), headers=None):

elasticsearch_async/connection_pool.py

@@ -0,0 +1,18 @@
+import asyncio
+
+from elasticsearch.connection_pool import ConnectionPool, DummyConnectionPool
+
+
+class AsyncConnectionPool(ConnectionPool):
+ def __init__(self, connections, loop, **kwargs):
+ self.loop = loop
+ super().__init__(connections, **kwargs)
+
+ async def close(self):
+ await asyncio.gather(*[conn.close() for conn in self.orig_connections],
+ loop=self.loop)
+
+
+class AsyncDummyConnectionPool(DummyConnectionPool):
+ async def close(self):
+ await self.connection.close()

elasticsearch_async/transport.py

@@ -4,19 +4,23 @@
 from itertools import chain
 
 from elasticsearch import Transport, TransportError, ConnectionTimeout, ConnectionError, SerializationError
+from elasticsearch.connection_pool import DummyConnectionPool
 
 from .connection import AIOHttpConnection
+from .connection_pool import AsyncConnectionPool, AsyncDummyConnectionPool
 from .helpers import ensure_future
 
 logger = logging.getLogger('elasticsearch')
 
 class AsyncTransport(Transport):
  def __init__(self, hosts, connection_class=AIOHttpConnection, loop=None,
+ connection_pool_class=AsyncConnectionPool,
  sniff_on_start=False, raise_on_sniff_error=True, **kwargs):
  self.raise_on_sniff_error = raise_on_sniff_error
  self.loop = asyncio.get_event_loop() if loop is None else loop
  kwargs['loop'] = self.loop
- super().__init__(hosts, connection_class=connection_class, sniff_on_start=False, **kwargs)
+ super().__init__(hosts, connection_class=connection_class, sniff_on_start=False, 
+ connection_pool_class=connection_pool_class, **kwargs)
 
  self.sniffing_task = None
  if sniff_on_start:
@@ -42,10 +46,16 @@ def initiate_sniff(self, initial=False):
  if self.sniffing_task is None:
  self.sniffing_task = ensure_future(self.sniff_hosts(initial), loop=self.loop)
 
+ @asyncio.coroutine
  def close(self):
  if self.sniffing_task:
  self.sniffing_task.cancel()
- super().close()
+ yield from self.connection_pool.close()
+
+ def set_connections(self, hosts):
+ super().set_connections(hosts)
+ if isinstance(self.connection_pool, DummyConnectionPool):
+ self.connection_pool = AsyncDummyConnectionPool(self.connection_pool.connection_opts)
 
  def get_connection(self):
  if self.sniffer_timeout:

setup.py

@@ -20,7 +20,6 @@
  'pytest',
  'pytest-asyncio',
  'pytest-cov',
- 'pytest-catchlog',
 ]
 
 setup(

test_elasticsearch_async/conftest.py

@@ -1,5 +1,6 @@
 import asyncio
 import json
+import logging
 
 from pytest import yield_fixture, fixture
 
@@ -11,7 +12,7 @@
 def connection(event_loop, server, port):
  connection = AIOHttpConnection(port=port, loop=event_loop)
  yield connection
- connection.close()
+ event_loop.run_until_complete(connection.close())
 
 
 class DummyElasticsearch(aiohttp.web.Server):
@@ -61,17 +62,19 @@ def port():
 
 @yield_fixture
 def server(event_loop, port):
- server = DummyElasticsearch(debug=True, keep_alive=75)
+ server = DummyElasticsearch(debug=True, loop=event_loop)
  f = event_loop.create_server(server, '127.0.0.1', port)
  event_loop.run_until_complete(f)
  yield server
  event_loop.run_until_complete(server.shutdown(timeout=.5))
 
 @yield_fixture
 def client(event_loop, server, port):
+ logger = logging.getLogger('elasticsearch')
+ logger.setLevel(logging.DEBUG)
  c = AsyncElasticsearch([{'host': '127.0.0.1','port': port}], loop=event_loop)
  yield c
- c.transport.close()
+ event_loop.run_until_complete(c.transport.close())
 
 @fixture
 def sniff_data():

test_elasticsearch_async/test_connection.py

@@ -30,7 +30,7 @@ def test_auth_is_set_correctly(event_loop):
 def test_ssl_context_is_correctly(event_loop):
  context = create_ssl_context(cafile="test_elasticsearch_async/ca.crt")
  connection = AIOHttpConnection(ssl_context=context, http_auth=('user', 'secret'), loop=event_loop)
- assert connection.session.connector.ssl_context.get_ca_certs() == [{
+ assert connection.session.connector._ssl.get_ca_certs() == [{
  'subject': ((('commonName', 'Elastic Certificate Tool Autogenerated CA'),),),
  'issuer': ((('commonName', 'Elastic Certificate Tool Autogenerated CA'),),),
  'version': 3,

test_elasticsearch_async/test_connection_pool.py

@@ -0,0 +1,36 @@
+from pytest import mark
+
+from elasticsearch_async import AsyncElasticsearch
+from elasticsearch_async.connection_pool import \
+ AsyncConnectionPool, AsyncDummyConnectionPool
+
+
+@mark.asyncio
+def test_single_host_makes_async_dummy_pool(server, client, event_loop, port):
+ client = AsyncElasticsearch(port=port, loop=event_loop)
+ assert isinstance(client.transport.connection_pool, AsyncDummyConnectionPool)
+ yield from client.transport.close()
+
+@mark.asyncio
+def test_multiple_hosts_make_async_pool(server, event_loop, port):
+ client = AsyncElasticsearch(
+ hosts=['localhost', 'localhost'], port=port, loop=event_loop)
+ assert isinstance(client.transport.connection_pool, AsyncConnectionPool)
+ assert len(client.transport.connection_pool.orig_connections) == 2
+ yield from client.transport.close()
+
+@mark.asyncio
+def test_async_dummy_pool_is_closed_properly(server, event_loop, port):
+ client = AsyncElasticsearch(port=port, loop=event_loop)
+ assert isinstance(client.transport.connection_pool, AsyncDummyConnectionPool)
+ yield from client.transport.close()
+ assert client.transport.connection_pool.connection.session.closed
+
+@mark.asyncio
+def test_async_pool_is_closed_properly(server, event_loop, port):
+ client = AsyncElasticsearch(
+ hosts=['localhost', 'localhost'], port=port, loop=event_loop)
+ assert isinstance(client.transport.connection_pool, AsyncConnectionPool)
+ yield from client.transport.close()
+ for conn in client.transport.connection_pool.orig_connections:
+ assert conn.session.closed

test_elasticsearch_async/test_transport.py

@@ -19,12 +19,12 @@ def test_sniff_on_start_sniffs(server, event_loop, port, sniff_data):
 
  assert 1 == len(connections)
  assert 'http://node1:9200' == connections[0].host
- client.transport.close()
+ yield from client.transport.close()
 
 @mark.asyncio
 def test_retry_will_work(port, server, event_loop):
  client = AsyncElasticsearch(hosts=['not-an-es-host', 'localhost'], port=port, loop=event_loop, randomize_hosts=False)
 
  data = yield from client.info()
  assert {'body': '', 'method': 'GET', 'params': {}, 'path': '/'} == data
- client.transport.close()
+ yield from client.transport.close()