From 23d96d277cb3645c62ca550337401bedc78318a1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Johan=20H=C3=A4gg?= <johan.hagg@pingdom.com>
Date: Wed, 17 Aug 2016 14:28:07 +0200
Subject: [PATCH] (GITHUB-89) Do not throw ValueError on authentication error
 during POST

---
 librato/__init__.py            | 31 ++++++++++++++++++++---
 tests/mock_connection.py       |  4 +++
 tests/test_process_response.py | 46 ++++++++++++++++++++++++++++++++++
 3 files changed, 77 insertions(+), 4 deletions(-)
 create mode 100644 tests/test_process_response.py

diff --git a/librato/__init__.py b/librato/__init__.py
index 7828cad..2908c46 100644
--- a/librato/__init__.py
+++ b/librato/__init__.py
@@ -162,10 +162,7 @@ def _process_response(self, resp, backoff):
         not_a_server_error = resp.status < 500
 
         if not_a_server_error:
-            body = resp.read()
-            if body:
-                resp_data = json.loads(body.decode(_getcharset(resp)))
-            log.info("body(<-): %s" % body)
+            resp_data = _decode_body(resp)
             a_client_error = resp.status >= 400
             if a_client_error:
                 raise exceptions.get(resp.status, resp_data)
@@ -568,6 +565,25 @@ def connect(username, api_key, hostname=HOSTNAME, base_path=BASE_PATH, sanitizer
     return LibratoConnection(username, api_key, hostname, base_path, sanitizer=sanitizer, protocol=protocol)
 
 
+def _decode_body(resp):
+    """
+    Read and decode HTTPResponse body based on charset and content-type
+    """
+    body = resp.read()
+    log.info("body(<-): %s" % body)
+    if not body:
+        return None
+
+    decoded_body = body.decode(_getcharset(resp))
+    content_type = _get_content_type(resp)
+
+    if content_type == "application/json":
+        resp_data = json.loads(decoded_body)
+    else:
+        resp_data = decoded_body
+
+    return resp_data
+
 def _getcharset(resp, default='utf-8'):
     """
     Extract the charset from an HTTPResponse.
@@ -581,3 +597,10 @@ def _getcharset(resp, default='utf-8'):
         m = email.message.Message()
         m['content-type'] = resp.getheader('content-type')
         return m.get_content_charset(default)
+
+def _get_content_type(resp):
+    """
+    Get Content-Type header ignoring parameters
+    """
+    parts = resp.getheader('content-type', "application/json").split(";")
+    return parts[0]
diff --git a/tests/mock_connection.py b/tests/mock_connection.py
index 1312d01..465618b 100644
--- a/tests/mock_connection.py
+++ b/tests/mock_connection.py
@@ -446,12 +446,16 @@ class MockResponse(object):
     def __init__(self, request, fake_failure=False):
         self.request = request
         self.status = 500 if fake_failure else 200
+        self._headers = {'content-type': "application/json;charset=utf-8"}
 
     class headers(object):
         @staticmethod
         def get_content_charset(default):
             return 'utf-8'
 
+    def getheader(self, name, default=None):
+        return self._headers.get(name.lower(), default)
+
     def read(self):
         return self._json_body_based_on_request()
 
diff --git a/tests/test_process_response.py b/tests/test_process_response.py
new file mode 100644
index 0000000..99ed42f
--- /dev/null
+++ b/tests/test_process_response.py
@@ -0,0 +1,46 @@
+import logging
+import unittest
+try:
+    from unittest.mock import create_autospec, PropertyMock
+except ImportError:
+    from mock import create_autospec, PropertyMock
+import librato
+from mock_connection import MockConnect, server
+from six.moves.http_client import HTTPResponse
+
+#logging.basicConfig(level=logging.DEBUG)
+# Mock the server
+librato.HTTPSConnection = MockConnect
+
+class TestLibrato(unittest.TestCase):
+    def setUp(self):
+        self.conn = librato.connect('user_test', 'key_test')
+        server.clean()
+
+    def test_get_authentication_failure(self):
+        """
+        fails with Unauthorized on 401 during GET
+        """
+        mock_response = create_autospec(HTTPResponse, spec_set=True, instance=True)
+        mock_response.mock_add_spec(['status'], spec_set=True)
+        mock_response.status = 401
+        # GET 401 responds with JSON
+        mock_response.getheader.return_value = "application/json;charset=utf-8"
+        mock_response.read.return_value = '{"errors":{"request":["Authorization Required"]}}'.encode('utf-8')
+
+        with self.assertRaises(librato.exceptions.Unauthorized):
+            self.conn._process_response(mock_response, 1)
+
+    def test_post_authentication_failure(self):
+        """
+        fails with Unauthorized on 401 during POST
+        """
+        mock_response = create_autospec(HTTPResponse, spec_set=True, instance=True)
+        mock_response.mock_add_spec(['status'], spec_set=True)
+        mock_response.status = 401
+        # POST 401 responds with text
+        mock_response.getheader.return_value = "text/plain"
+        mock_response.read.return_value = 'Credentials are required to access this resource.'.encode('utf-8')
+
+        with self.assertRaises(librato.exceptions.Unauthorized):
+            self.conn._process_response(mock_response, 1)