Add tests for AIOHttpConnection

sethmlarson · sethmlarson · commit aaed0d4fadd1 · 2020-04-29T16:19:42.000-05:00
diff --git a/elasticsearch/_async/client/__init__.py b/elasticsearch/_async/client/__init__.py
@@ -229,6 +229,9 @@ async def __aenter__(self):
         return self
 
     async def __aexit__(self, *_):
+        await self.close()
+
+    async def close(self):
         await self.transport.close()
 
     # AUTO-GENERATED-API-DEFINITIONS #
diff --git a/elasticsearch/_async/http_aiohttp.py b/elasticsearch/_async/http_aiohttp.py
@@ -5,6 +5,7 @@
 import asyncio
 import ssl
 import os
+import urllib3
 import warnings
 
 import aiohttp
@@ -25,6 +26,7 @@
 # This is used to detect if a user is passing in a value
 # for SSL kwargs if also using an SSLContext.
 VERIFY_CERTS_DEFAULT = object()
+SSL_SHOW_WARN_DEFAULT = object()
 
 CA_CERTS = None
 
@@ -43,7 +45,8 @@ def __init__(
         port=None,
         http_auth=None,
         use_ssl=False,
-        verify_certs=True,
+        verify_certs=VERIFY_CERTS_DEFAULT,
+        ssl_show_warn=SSL_SHOW_WARN_DEFAULT,
         ca_certs=None,
         client_cert=None,
         client_key=None,
@@ -74,15 +77,14 @@ def __init__(
         )
 
         if http_auth is not None:
-            if isinstance(http_auth, str):
-                http_auth = tuple(http_auth.split(":", 1))
-
             if isinstance(http_auth, (tuple, list)):
-                http_auth = aiohttp.BasicAuth(*http_auth)
+                http_auth = ":".join(http_auth)
+            self.headers.update(urllib3.make_headers(basic_auth=http_auth))
 
         # if providing an SSL context, raise error if any other SSL related flag is used
         if ssl_context and (
             (verify_certs is not VERIFY_CERTS_DEFAULT)
+            or (ssl_show_warn is not SSL_SHOW_WARN_DEFAULT)
             or ca_certs
             or client_cert
             or client_key
@@ -100,6 +102,8 @@ def __init__(
             # values if not using an SSLContext.
             if verify_certs is VERIFY_CERTS_DEFAULT:
                 verify_certs = True
+            if ssl_show_warn is SSL_SHOW_WARN_DEFAULT:
+                ssl_show_warn = True
 
             ca_certs = CA_CERTS if ca_certs is None else ca_certs
             if verify_certs:
@@ -109,6 +113,13 @@ def __init__(
                         "validation. Either pass them in using the ca_certs parameter or "
                         "install certifi to use it automatically."
                     )
+            else:
+                if ssl_show_warn:
+                    warnings.warn(
+                        "Connecting to %s using SSL with verify_certs=False is insecure."
+                        % self.host
+                    )
+
             if os.path.isfile(ca_certs):
                 ssl_context.load_verify_locations(cafile=ca_certs)
             elif os.path.isdir(ca_certs):
@@ -180,7 +191,7 @@ async def perform_request(
                     await response.release()
                     raw_data = ""
                 else:
-                    raw_data = await response.text()
+                    raw_data = (await response.read()).decode("utf-8", "surrogatepass")
                 duration = self.loop.time() - start
 
         # We want to reraise a cancellation.
diff --git a/test_elasticsearch/test_async/__init__.py b/test_elasticsearch/test_async/__init__.py
@@ -0,0 +1,4 @@
+# Licensed to Elasticsearch B.V under one or more agreements.
+# Elasticsearch B.V licenses this file to you under the Apache 2.0 License.
+# See the LICENSE file in the project root for more information
+
diff --git a/test_elasticsearch/test_async/test_connection.py b/test_elasticsearch/test_async/test_connection.py
@@ -0,0 +1,298 @@
+# -*- coding: utf-8 -*-
+# Licensed to Elasticsearch B.V under one or more agreements.
+# Elasticsearch B.V licenses this file to you under the Apache 2.0 License.
+# See the LICENSE file in the project root for more information
+
+import re
+import ssl
+from mock import Mock, patch
+import warnings
+from platform import python_version
+
+from elasticsearch.exceptions import (
+    TransportError,
+    ConflictError,
+    RequestError,
+    NotFoundError,
+)
+from elasticsearch import AIOHttpConnection
+from elasticsearch import __versionstr__
+from ..test_cases import TestCase, SkipTest
+
+
+class TestAIOHttpConnection(TestCase):
+    async def _get_mock_connection(self, connection_params={}, response_body=b"{}"):
+        con = AIOHttpConnection(**connection_params)
+
+        async def _dummy_request(*args, **kwargs):
+            async def read():
+                return response_body
+
+            dummy_response = Mock()
+            dummy_response.headers = {}
+            dummy_response.status_code = 200
+            dummy_response.read = read
+            _dummy_request.call_args = (args, kwargs)
+            return dummy_response
+
+        con.session.request = _dummy_request
+        return con
+
+    async def test_ssl_context(self):
+        try:
+            context = ssl.create_default_context()
+        except AttributeError:
+            # if create_default_context raises an AttributeError Exception
+            # it means SSLContext is not available for that version of python
+            # and we should skip this test.
+            raise SkipTest(
+                "Test test_ssl_context is skipped cause SSLContext is not available for this version of ptyhon"
+            )
+
+        con = AIOHttpConnection(use_ssl=True, ssl_context=context)
+        await con._create_aiohttp_session()
+        self.assertTrue(con.use_ssl)
+        self.assertEqual(con.session.ssl_context, context)
+
+    def test_opaque_id(self):
+        con = AIOHttpConnection(opaque_id="app-1")
+        self.assertEqual(con.headers["x-opaque-id"], "app-1")
+
+    def test_http_cloud_id(self):
+        con = AIOHttpConnection(
+            cloud_id="cluster:dXMtZWFzdC0xLmF3cy5mb3VuZC5pbyQ0ZmE4ODIxZTc1NjM0MDMyYmVkMWNmMjIxMTBlMmY5NyQ0ZmE4ODIxZTc1NjM0MDMyYmVkMWNmMjIxMTBlMmY5Ng=="
+        )
+        self.assertTrue(con.use_ssl)
+        self.assertEqual(
+            con.host, "https://4fa8821e75634032bed1cf22110e2f97.us-east-1.aws.found.io"
+        )
+        self.assertEqual(con.port, None)
+        self.assertEqual(
+            con.hostname, "4fa8821e75634032bed1cf22110e2f97.us-east-1.aws.found.io"
+        )
+        self.assertTrue(con.http_compress)
+
+        con = AIOHttpConnection(
+            cloud_id="cluster:dXMtZWFzdC0xLmF3cy5mb3VuZC5pbyQ0ZmE4ODIxZTc1NjM0MDMyYmVkMWNmMjIxMTBlMmY5NyQ0ZmE4ODIxZTc1NjM0MDMyYmVkMWNmMjIxMTBlMmY5Ng==",
+            port=9243,
+        )
+        self.assertEqual(
+            con.host,
+            "https://4fa8821e75634032bed1cf22110e2f97.us-east-1.aws.found.io:9243",
+        )
+        self.assertEqual(con.port, 9243)
+        self.assertEqual(
+            con.hostname, "4fa8821e75634032bed1cf22110e2f97.us-east-1.aws.found.io"
+        )
+
+    def test_api_key_auth(self):
+        # test with tuple
+        con = AIOHttpConnection(
+            cloud_id="cluster:dXMtZWFzdC0xLmF3cy5mb3VuZC5pbyQ0ZmE4ODIxZTc1NjM0MDMyYmVkMWNmMjIxMTBlMmY5NyQ0ZmE4ODIxZTc1NjM0MDMyYmVkMWNmMjIxMTBlMmY5Ng==",
+            api_key=("elastic", "changeme1"),
+        )
+        self.assertEqual(
+            con.headers["authorization"], "ApiKey ZWxhc3RpYzpjaGFuZ2VtZTE="
+        )
+        self.assertEqual(
+            con.host, "https://4fa8821e75634032bed1cf22110e2f97.us-east-1.aws.found.io"
+        )
+
+        # test with base64 encoded string
+        con = AIOHttpConnection(
+            cloud_id="cluster:dXMtZWFzdC0xLmF3cy5mb3VuZC5pbyQ0ZmE4ODIxZTc1NjM0MDMyYmVkMWNmMjIxMTBlMmY5NyQ0ZmE4ODIxZTc1NjM0MDMyYmVkMWNmMjIxMTBlMmY5Ng==",
+            api_key="ZWxhc3RpYzpjaGFuZ2VtZTI=",
+        )
+        self.assertEqual(
+            con.headers["authorization"], "ApiKey ZWxhc3RpYzpjaGFuZ2VtZTI="
+        )
+        self.assertEqual(
+            con.host, "https://4fa8821e75634032bed1cf22110e2f97.us-east-1.aws.found.io"
+        )
+
+    async def test_no_http_compression(self):
+        con = await self._get_mock_connection()
+        self.assertFalse(con.http_compress)
+        self.assertNotIn("accept-encoding", con.headers)
+
+        await con.perform_request("GET", "/")
+
+        (_, _, req_body), kwargs = con.pool.urlopen.call_args
+
+        self.assertFalse(req_body)
+        self.assertNotIn("accept-encoding", kwargs["headers"])
+        self.assertNotIn("content-encoding", kwargs["headers"])
+
+    async def test_http_compression(self):
+        con = await self._get_mock_connection({"http_compress": True})
+        self.assertTrue(con.http_compress)
+        self.assertEqual(con.headers["accept-encoding"], "gzip,deflate")
+
+        # 'content-encoding' shouldn't be set at a connection level.
+        # Should be applied only if the request is sent with a body.
+        self.assertNotIn("content-encoding", con.headers)
+
+        await con.perform_request("GET", "/", body=b"{}")
+
+        (_, _, req_body), kwargs = con.pool.urlopen.call_args
+
+        self.assertEqual(gzip_decompress(req_body), b"{}")
+        self.assertEqual(kwargs["headers"]["accept-encoding"], "gzip,deflate")
+        self.assertEqual(kwargs["headers"]["content-encoding"], "gzip")
+
+        await con.perform_request("GET", "/")
+
+        (_, _, req_body), kwargs = con.pool.urlopen.call_args
+
+        self.assertFalse(req_body)
+        self.assertEqual(kwargs["headers"]["accept-encoding"], "gzip,deflate")
+        self.assertNotIn("content-encoding", kwargs["headers"])
+
+    def test_cloud_id_http_compress_override(self):
+        # 'http_compress' will be 'True' by default for connections with
+        # 'cloud_id' set but should prioritize user-defined values.
+        con = AIOHttpConnection(
+            cloud_id="cluster:dXMtZWFzdC0xLmF3cy5mb3VuZC5pbyQ0ZmE4ODIxZTc1NjM0MDMyYmVkMWNmMjIxMTBlMmY5NyQ0ZmE4ODIxZTc1NjM0MDMyYmVkMWNmMjIxMTBlMmY5Ng==",
+        )
+        self.assertEqual(con.http_compress, True)
+
+        con = AIOHttpConnection(
+            cloud_id="cluster:dXMtZWFzdC0xLmF3cy5mb3VuZC5pbyQ0ZmE4ODIxZTc1NjM0MDMyYmVkMWNmMjIxMTBlMmY5NyQ0ZmE4ODIxZTc1NjM0MDMyYmVkMWNmMjIxMTBlMmY5Ng==",
+            http_compress=False,
+        )
+        self.assertEqual(con.http_compress, False)
+
+        con = AIOHttpConnection(
+            cloud_id="cluster:dXMtZWFzdC0xLmF3cy5mb3VuZC5pbyQ0ZmE4ODIxZTc1NjM0MDMyYmVkMWNmMjIxMTBlMmY5NyQ0ZmE4ODIxZTc1NjM0MDMyYmVkMWNmMjIxMTBlMmY5Ng==",
+            http_compress=True,
+        )
+        self.assertEqual(con.http_compress, True)
+
+    def test_default_user_agent(self):
+        con = AIOHttpConnection()
+        self.assertEqual(
+            con._get_default_user_agent(),
+            "elasticsearch-py/%s (Python %s)" % (__versionstr__, python_version()),
+        )
+
+    def test_timeout_set(self):
+        con = AIOHttpConnection(timeout=42)
+        self.assertEqual(42, con.timeout)
+
+    def test_keep_alive_is_on_by_default(self):
+        con = AIOHttpConnection()
+        self.assertEqual(
+            {
+                "connection": "keep-alive",
+                "content-type": "application/json",
+                "user-agent": con._get_default_user_agent(),
+            },
+            con.headers,
+        )
+
+    def test_http_auth(self):
+        con = AIOHttpConnection(http_auth="username:secret")
+        self.assertEqual(
+            {
+                "authorization": "Basic dXNlcm5hbWU6c2VjcmV0",
+                "connection": "keep-alive",
+                "content-type": "application/json",
+                "user-agent": con._get_default_user_agent(),
+            },
+            con.headers,
+        )
+
+    def test_http_auth_tuple(self):
+        con = AIOHttpConnection(http_auth=("username", "secret"))
+        self.assertEqual(
+            {
+                "authorization": "Basic dXNlcm5hbWU6c2VjcmV0",
+                "content-type": "application/json",
+                "connection": "keep-alive",
+                "user-agent": con._get_default_user_agent(),
+            },
+            con.headers,
+        )
+
+    def test_http_auth_list(self):
+        con = AIOHttpConnection(http_auth=["username", "secret"])
+        self.assertEqual(
+            {
+                "authorization": "Basic dXNlcm5hbWU6c2VjcmV0",
+                "content-type": "application/json",
+                "connection": "keep-alive",
+                "user-agent": con._get_default_user_agent(),
+            },
+            con.headers,
+        )
+
+    def test_uses_https_if_verify_certs_is_off(self):
+        with warnings.catch_warnings(record=True) as w:
+            con = AIOHttpConnection(use_ssl=True, verify_certs=False)
+            self.assertEqual(1, len(w))
+            self.assertEqual(
+                "Connecting to https://localhost:9200 using SSL with verify_certs=False is insecure.",
+                str(w[0].message),
+            )
+
+        self.assertTrue(con.use_ssl)
+        self.assertEqual(con.scheme, "https")
+        self.assertEqual(con.host, "https://localhost:9200")
+
+    def nowarn_when_test_uses_https_if_verify_certs_is_off(self):
+        with warnings.catch_warnings(record=True) as w:
+            con = Urllib3HttpConnection(
+                use_ssl=True, verify_certs=False, ssl_show_warn=False
+            )
+            self.assertEqual(0, len(w))
+
+        self.assertIsInstance(con.pool, urllib3.HTTPSConnectionPool)
+
+    def test_doesnt_use_https_if_not_specified(self):
+        con = AIOHttpConnection()
+        self.assertFalse(con.use_ssl)
+
+    def test_no_warning_when_using_ssl_context(self):
+        ctx = ssl.create_default_context()
+        with warnings.catch_warnings(record=True) as w:
+            AIOHttpConnection(ssl_context=ctx)
+            self.assertEqual(0, len(w), str([x.message for x in w]))
+
+    def test_warns_if_using_non_default_ssl_kwargs_with_ssl_context(self):
+        for kwargs in (
+            {"ssl_show_warn": False},
+            {"ssl_show_warn": True},
+            {"verify_certs": True},
+            {"verify_certs": False},
+            {"ca_certs": "/path/to/certs"},
+            {"ssl_show_warn": True, "ca_certs": "/path/to/certs"},
+        ):
+            kwargs["ssl_context"] = ssl.create_default_context()
+
+            with warnings.catch_warnings(record=True) as w:
+                warnings.simplefilter("always")
+
+                AIOHttpConnection(**kwargs)
+
+                self.assertEqual(1, len(w))
+                self.assertEqual(
+                    "When using `ssl_context`, all other SSL related kwargs are ignored",
+                    str(w[0].message),
+                )
+
+    @patch("elasticsearch.connection.base.logger")
+    async def test_uncompressed_body_logged(self, logger):
+        con = await self._get_mock_connection(connection_params={"http_compress": True})
+        await con.perform_request("GET", "/", body=b'{"example": "body"}')
+
+        self.assertEqual(2, logger.debug.call_count)
+        req, resp = logger.debug.call_args_list
+
+        self.assertEqual('> {"example": "body"}', req[0][0] % req[0][1:])
+        self.assertEqual("< {}", resp[0][0] % resp[0][1:])
+
+    async def test_surrogatepass_into_bytes(self):
+        buf = b"\xe4\xbd\xa0\xe5\xa5\xbd\xed\xa9\xaa"
+        con = await self._get_mock_connection(response_body=buf)
+        status, headers, data = await con.perform_request("GET", "/")
+        self.assertEqual(u"你好\uda6a", data)
diff --git a/utils/templates/overrides/cluster/stats b/utils/templates/overrides/cluster/stats
@@ -1,5 +1,5 @@
 {% extends "base" %}
 {% block request %}
-        return self.transport.perform_request("{{ api.method }}", "/_cluster/stats" if node_id in SKIP_IN_PATH else _make_path("_cluster/stats/nodes", node_id), params=params, headers=headers)
+        return self.transport.perform_request("{{ api.method }}", "/_cluster/stats" if node_id in SKIP_IN_PATH else _make_path("_cluster", "stats", "nodes", node_id), params=params, headers=headers)
 {% endblock%}
 
