extract_response: extract content from requests.Response, and bind extracted value to context.extractors

ate/response.py

@@ -1,16 +1,17 @@
-from ate import utils
+from ate import utils, exception
 
 
+def parse_response_body(resp_obj):
+    try:
+        return resp_obj.json()
+    except ValueError:
+        return resp_obj.text
+
 def parse_response_object(resp_obj):
-    try:
-        resp_body = resp_obj.json()
-    except ValueError:
-        resp_body = resp_obj.text
-
     return {
         'status_code': resp_obj.status_code,
         'headers': resp_obj.headers,
-        'body': resp_body
+        'body': parse_response_body(resp_obj)
     }
 
 def diff_response(resp_obj, expected_resp_json):
@@ -51,3 +52,40 @@ def diff_response(resp_obj, expected_resp_json):
         diff_content['body'] = body_diff
 
     return diff_content
+
+def extract_response(resp_obj, context, delimiter='.'):
+    """ extract content from requests.Response, and bind extracted value to context.extractors
+    @param (requests.Response instance) resp_obj
+    @param (ate.context.Context instance) context
+        context.extractors:
+        {
+            "resp_status_code": "status_code",
+            "resp_headers_content_type": "headers.content-type",
+            "resp_content": "content",
+            "resp_content_person_first_name": "content.person.name.first_name"
+        }
+    """
+    for key, value in context.extractors.items():
+        try:
+            if isinstance(value, str):
+                value += "."
+                top_query, sub_query = value.split(delimiter, maxsplit=1)
+
+                if top_query in ["body", "content", "text"]:
+                    json_content = parse_response_body(resp_obj)
+                else:
+                    json_content = getattr(resp_obj, top_query)
+
+                if sub_query:
+                    # e.g. key: resp_headers_content_type, sub_query = "content-type"
+                    answer = utils.query_json(json_content, sub_query)
+                    context.extractors[key] = answer
+                else:
+                    # e.g. key: resp_status_code, resp_content
+                    context.extractors[key] = json_content
+
+            else:
+                raise NotImplementedError("TODO: support template.")
+
+        except AttributeError:
+            raise exception.ParamsError("invalid extract_binds!")

test/test_response.py

@@ -1,9 +1,9 @@
 import random
 import requests
-from ate import response
+from ate import response, context, exception
 from test.base import ApiServerUnittest
 
-class TestUtils(ApiServerUnittest):
+class TestResponse(ApiServerUnittest):
 
     def test_parse_response_object_json(self):
         url = "http://127.0.0.1:5000/api/users"
@@ -210,3 +210,131 @@ class TestUtils(ApiServerUnittest):
                 }
             }
         )
+
+    def test_extract_response_json(self):
+        resp_obj = requests.post(
+            url="http://127.0.0.1:5000/customize-response",
+            json={
+                'headers': {
+                    'Content-Type': "application/json"
+                },
+                'body': {
+                    'success': False,
+                    "person": {
+                        "name": {
+                            "first_name": "Leo",
+                            "last_name": "Lee",
+                        },
+                        "age": 29,
+                        "cities": ["Guangzhou", "Shenzhen"]
+                    }
+                }
+            }
+        )
+
+        extract_binds = {
+            "resp_status_code": "status_code",
+            "resp_headers_content_type": "headers.content-type",
+            "resp_content_body_success": "body.success",
+            "resp_content_content_success": "content.success",
+            "resp_content_text_success": "text.success",
+            "resp_content_person_first_name": "content.person.name.first_name",
+            "resp_content_cities_1": "content.person.cities.1"
+        }
+
+        test_context = context.Context()
+        test_context.bind_extractors(extract_binds)
+        response.extract_response(resp_obj, test_context)
+
+        extract_binds_dict = test_context.extractors
+        self.assertEqual(
+            extract_binds_dict["resp_status_code"],
+            200
+        )
+        self.assertEqual(
+            extract_binds_dict["resp_headers_content_type"],
+            "application/json"
+        )
+        self.assertEqual(
+            extract_binds_dict["resp_content_content_success"],
+            False
+        )
+        self.assertEqual(
+            extract_binds_dict["resp_content_text_success"],
+            False
+        )
+        self.assertEqual(
+            extract_binds_dict["resp_content_person_first_name"],
+            "Leo"
+        )
+        self.assertEqual(
+            extract_binds_dict["resp_content_cities_1"],
+            "Shenzhen"
+        )
+
+
+    def test_extract_response_fail(self):
+        resp_obj = requests.post(
+            url="http://127.0.0.1:5000/customize-response",
+            json={
+                'headers': {
+                    'Content-Type': "application/json"
+                },
+                'body': {
+                    'success': False,
+                    "person": {
+                        "name": {
+                            "first_name": "Leo",
+                            "last_name": "Lee",
+                        },
+                        "age": 29,
+                        "cities": ["Guangzhou", "Shenzhen"]
+                    }
+                }
+            }
+        )
+
+        extract_binds = {
+            "resp_content_dict_key_error": "content.not_exist"
+        }
+
+        test_context = context.Context()
+        test_context.bind_extractors(extract_binds)
+
+        with self.assertRaises(exception.ParamsError):
+            response.extract_response(resp_obj, test_context)
+
+        extract_binds = {
+            "resp_content_list_index_error": "content.person.cities.3"
+        }
+
+        test_context = context.Context()
+        test_context.bind_extractors(extract_binds)
+
+        with self.assertRaises(exception.ParamsError):
+            response.extract_response(resp_obj, test_context)
+
+    def test_extract_response_json_string(self):
+        resp_obj = requests.post(
+            url="http://127.0.0.1:5000/customize-response",
+            json={
+                'headers': {
+                    'Content-Type': "application/json"
+                },
+                'body': "abc"
+            }
+        )
+
+        extract_binds = {
+            "resp_content_body": "content"
+        }
+
+        test_context = context.Context()
+        test_context.bind_extractors(extract_binds)
+        response.extract_response(resp_obj, test_context)
+
+        extract_binds_dict = test_context.extractors
+        self.assertEqual(
+            extract_binds_dict["resp_content_body"],
+            "abc"
+        )