[7.x] Add HTTP compression to RequestsHttpConnection

test_elasticsearch/test_connection.py

@@ -1,6 +1,7 @@
-import sys
 import re
 import ssl
+import gzip
+import io
 from mock import Mock, patch
 import urllib3
 import warnings
@@ -14,11 +15,32 @@ from elasticsearch.exceptions import (
     NotFoundError,
 )
 from elasticsearch.connection import RequestsHttpConnection, Urllib3HttpConnection
-from elasticsearch.exceptions import ImproperlyConfigured
 from elasticsearch import __versionstr__
 from .test_cases import TestCase, SkipTest
 
+
+def gzip_decompress(data):
+    buf = gzip.GzipFile(fileobj=io.BytesIO(data), mode="rb")
+    return buf.read()
+
+
 class TestUrllib3Connection(TestCase):
+    def _get_mock_connection(
+        self, connection_params={}, response_body=b"{}"
+    ):
+        con = Urllib3HttpConnection(**connection_params)
+
+        def _dummy_urlopen(*args, **kwargs):
+            dummy_response = Mock()
+            dummy_response.headers = {}
+            dummy_response.status = 200
+            dummy_response.data = response_body
+            _dummy_urlopen.call_args = (args, kwargs)
+            return dummy_response
+
+        con.pool.urlopen = _dummy_urlopen
+        return con
+
     def test_ssl_context(self):
         try:
             context = ssl.create_default_context()
@@ -59,10 +81,43 @@ class TestUrllib3Connection(TestCase):
         )
         self.assertEquals(con.headers["authorization"], "ApiKey ZWxhc3RpYzpjaGFuZ2VtZTI=")
 
+    def test_no_http_compression(self):
+        con = self._get_mock_connection()
+        self.assertFalse(con.http_compress)
+        self.assertNotIn("accept-encoding", con.headers)
+
+        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"])
+
     def test_http_compression(self):
-        con = Urllib3HttpConnection(http_compress=True)
+        con = self._get_mock_connection({"http_compress": True})
         self.assertTrue(con.http_compress)
-        self.assertEquals(con.headers["content-encoding"], "gzip")
+        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)
+
+        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")
+
+        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_default_user_agent(self):
         con = Urllib3HttpConnection()
@@ -171,6 +226,17 @@ class TestUrllib3Connection(TestCase):
                     str(w[0].message),
                 )
 
+    @patch("elasticsearch.connection.base.logger")
+    def test_uncompressed_body_logged(self, logger):
+        con = self._get_mock_connection(connection_params={"http_compress": True})
+        con.perform_request("GET", "/", body=b"{\"example\": \"body\"}")
+
+        self.assertEquals(2, logger.debug.call_count)
+        req, resp = logger.debug.call_args_list
+        print(req, resp)
+        self.assertEquals('> {"example": "body"}', req[0][0] % req[0][1:])
+        self.assertEquals('< {}', resp[0][0] % resp[0][1:])
+
 
 class TestRequestsConnection(TestCase):
     def _get_mock_connection(
@@ -239,6 +305,41 @@ class TestRequestsConnection(TestCase):
         )
         self.assertEquals(con.session.headers["authorization"], "ApiKey ZWxhc3RpYzpjaGFuZ2VtZTI=")
 
+    def test_no_http_compression(self):
+        con = self._get_mock_connection()
+
+        self.assertFalse(con.http_compress)
+        self.assertNotIn("content-encoding", con.session.headers)
+
+        con.perform_request("GET", "/")
+
+        req = con.session.send.call_args[0][0]
+        self.assertNotIn("content-encoding", req.headers)
+        self.assertNotIn("accept-encoding", req.headers)
+
+    def test_http_compression(self):
+        con = self._get_mock_connection(
+            {"http_compress": True},
+        )
+
+        self.assertTrue(con.http_compress)
+
+        # 'content-encoding' shouldn't be set at a session level.
+        # Should be applied only if the request is sent with a body.
+        self.assertNotIn("content-encoding", con.session.headers)
+
+        con.perform_request("GET", "/", body=b"{}")
+
+        req = con.session.send.call_args[0][0]
+        self.assertEqual(req.headers["content-encoding"], "gzip")
+        self.assertEqual(req.headers["accept-encoding"], "gzip,deflate")
+
+        con.perform_request("GET", "/")
+
+        req = con.session.send.call_args[0][0]
+        self.assertNotIn("content-encoding", req.headers)
+        self.assertEqual(req.headers["accept-encoding"], "gzip,deflate")
+
     def test_uses_https_if_verify_certs_is_off(self):
         with warnings.catch_warnings(record=True) as w:
             con = self._get_mock_connection(
@@ -400,6 +501,16 @@ class TestRequestsConnection(TestCase):
         self.assertEquals('> {"question": "what\'s that?"}', req[0][0] % req[0][1:])
         self.assertEquals('< {"answer": "that\'s it!"}', resp[0][0] % resp[0][1:])
 
+    @patch("elasticsearch.connection.base.logger")
+    def test_uncompressed_body_logged(self, logger):
+        con = self._get_mock_connection(connection_params={"http_compress": True})
+        con.perform_request("GET", "/", body=b"{\"example\": \"body\"}")
+
+        self.assertEquals(2, logger.debug.call_count)
+        req, resp = logger.debug.call_args_list
+        self.assertEquals('> {"example": "body"}', req[0][0] % req[0][1:])
+        self.assertEquals('< {}', resp[0][0] % resp[0][1:])
+
     def test_defaults(self):
         con = self._get_mock_connection()
         request = self._get_request(con, "GET", "/")