From 223cb20f72c3f5002b57363a947178cecee1e878 Mon Sep 17 00:00:00 2001
From: debugtalk <mail@debugtalk.com>
Date: Tue, 29 Aug 2017 14:44:44 +0800
Subject: [PATCH] bugfix #34: handle exception when response content is empty

---
 ate/__init__.py        |  2 +-
 ate/exception.py       |  3 +++
 ate/response.py        |  7 +++++--
 ate/utils.py           | 17 ++++++++++-------
 tests/test_response.py | 28 ++++++++++++++++++++++++++++
 tests/test_utils.py    | 11 +++++++++++
 6 files changed, 58 insertions(+), 10 deletions(-)

diff --git a/ate/__init__.py b/ate/__init__.py
index f93e0653..b4a839ec 100644
--- a/ate/__init__.py
+++ b/ate/__init__.py
@@ -1 +1 @@
-__version__ = '0.5.2'
\ No newline at end of file
+__version__ = '0.5.3'
\ No newline at end of file
diff --git a/ate/exception.py b/ate/exception.py
index 56768985..53b91138 100644
--- a/ate/exception.py
+++ b/ate/exception.py
@@ -6,6 +6,9 @@ class MyBaseError(BaseException):
 class ParamsError(MyBaseError):
     pass
 
+class ResponseError(MyBaseError):
+    pass
+
 class ParseResponseError(MyBaseError):
     pass
 
diff --git a/ate/response.py b/ate/response.py
index 34dc0f43..d97e4a89 100644
--- a/ate/response.py
+++ b/ate/response.py
@@ -31,10 +31,13 @@ class ResponseObject(object):
             "content.person.name.first_name"
         """
         try:
-            field += "."
             # string.split(sep=None, maxsplit=-1) -> list of strings
             # e.g. "content.person.name" => ["content", "person.name"]
-            top_query, sub_query = field.split(delimiter, 1)
+            try:
+                top_query, sub_query = field.split(delimiter, 1)
+            except ValueError:
+                top_query = field
+                sub_query = None
 
             if top_query in ["body", "content", "text"]:
                 json_content = self.parsed_body()
diff --git a/ate/utils.py b/ate/utils.py
index f2ad95cd..00ea94f5 100644
--- a/ate/utils.py
+++ b/ate/utils.py
@@ -6,9 +6,10 @@ import os.path
 import random
 import re
 import string
-import yaml
 
+import yaml
 from ate import exception
+from requests.structures import CaseInsensitiveDict
 
 try:
     string_type = basestring
@@ -131,15 +132,17 @@ def query_json(json_content, query, delimiter='.'):
         "person.cities.0"         =>  "Guangzhou"
     @return queried result
     """
-    stripped_query = query.strip(delimiter)
-    if not stripped_query:
-        return None
+    if json_content == "":
+        raise exception.ResponseError("response content is empty!")
 
     try:
-        for key in stripped_query.split(delimiter):
+        for key in query.split(delimiter):
             if isinstance(json_content, list):
-                key = int(key)
-            json_content = json_content[key]
+                json_content = json_content[int(key)]
+            elif isinstance(json_content, (dict, CaseInsensitiveDict)):
+                json_content = json_content[key]
+            else:
+                raise exception.ParseResponseError("response content is in text format! failed to query key {}!".format(key))
     except (KeyError, ValueError, IndexError):
         raise exception.ParseResponseError("failed to query json when extracting response!")
 
diff --git a/tests/test_response.py b/tests/test_response.py
index b290c947..8ad281bd 100644
--- a/tests/test_response.py
+++ b/tests/test_response.py
@@ -149,6 +149,34 @@ class TestResponse(ApiServerUnittest):
             "abc"
         )
 
+    def test_extract_response_empty(self):
+        resp = requests.post(
+            url="http://127.0.0.1:5000/customize-response",
+            json={
+                'headers': {
+                    'Content-Type': "application/json"
+                },
+                'body': ""
+            }
+        )
+
+        extract_binds_list = [
+            {"resp_content_body": "content"}
+        ]
+        resp_obj = response.ResponseObject(resp)
+        extract_binds_dict_list = resp_obj.extract_response(extract_binds_list)
+        self.assertEqual(
+            extract_binds_dict_list[0]["resp_content_body"],
+            ""
+        )
+
+        extract_binds_list = [
+            {"resp_content_body": "content.abc"}
+        ]
+        resp_obj = response.ResponseObject(resp)
+        with self.assertRaises(exception.ResponseError):
+            resp_obj.extract_response(extract_binds_list)
+
     def test_validate(self):
         url = "http://127.0.0.1:5000/"
         resp = requests.get(url)
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 619bebf3..8b869266 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -154,6 +154,17 @@ class TestUtils(ApiServerUnittest):
         result = utils.query_json(json_content, query)
         self.assertEqual(result, "Leo")
 
+    def test_query_json_content_is_text(self):
+        json_content = ""
+        query = "key"
+        with self.assertRaises(exception.ResponseError):
+            utils.query_json(json_content, query)
+
+        json_content = "<html><body>content</body></html>"
+        query = "key"
+        with self.assertRaises(exception.ParseResponseError):
+            utils.query_json(json_content, query)
+
     def test_match_expected(self):
         self.assertTrue(utils.match_expected(1, 1, "eq"))
         self.assertTrue(utils.match_expected("abc", "abc", "=="))