From 4771dd272eecef0ed5ecd240be409125c8b9dad1 Mon Sep 17 00:00:00 2001
From: debugtalk <mail@debugtalk.com>
Date: Thu, 26 Jul 2018 22:33:35 +0800
Subject: [PATCH] bugfix: set repsonse attribute

---
 httprunner/response.py | 16 +++++++++++
 tests/debugtalk.py     |  1 +
 tests/test_runner.py   | 65 ++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 82 insertions(+)

diff --git a/httprunner/response.py b/httprunner/response.py
index d562769a..6b7435f5 100644
--- a/httprunner/response.py
+++ b/httprunner/response.py
@@ -18,6 +18,7 @@ class ResponseObject(object):
         @param (requests.Response instance) resp_obj
         """
         self.resp_obj = resp_obj
+        self.attributes = {}
 
     def __getattr__(self, key):
         try:
@@ -152,6 +153,21 @@ class ResponseObject(object):
                 logger.log_error(err_msg)
                 raise exceptions.ExtractFailure(err_msg)
 
+        # new set response attributes
+        elif top_query == "attributes":
+            if not sub_query:
+                # extract response attributes
+                return self.attributes
+
+            if sub_query in self.attributes:
+                return self.attributes[sub_query]
+            else:
+                # content = "attributes.new_attribute_not_exist"
+                err_msg = u"Failed to extract cumstom set attribute! => {}\n".format(field)
+                err_msg += u"response set attributes: {}\n".format(self.attributes)
+                logger.log_error(err_msg)
+                raise exceptions.ExtractFailure(err_msg)
+
         # others
         else:
             err_msg = u"Failed to extract attribute from response! => {}\n".format(field)
diff --git a/tests/debugtalk.py b/tests/debugtalk.py
index 5cd762ce..1d10b20a 100644
--- a/tests/debugtalk.py
+++ b/tests/debugtalk.py
@@ -102,3 +102,4 @@ def alter_response(response):
     response.status_code = 500
     response.headers["Content-Type"] = "html/text"
     response.json["headers"]["Host"] = "127.0.0.1:8888"
+    response.attributes["new_attribute"] = "new_attribute"
diff --git a/tests/test_runner.py b/tests/test_runner.py
index 5aa1b83f..029bea7e 100644
--- a/tests/test_runner.py
+++ b/tests/test_runner.py
@@ -163,6 +163,71 @@ class TestRunner(ApiServerUnittest):
         self.assertTrue(summary["success"])
         self.assertLess(end_time - start_time, 1)
 
+    def test_run_httprunner_with_teardown_hooks_alter_response(self):
+        testsets = [
+            {
+                "name": "test teardown hooks",
+                "config": {
+                    'path': 'tests/httpbin/hooks.yml',
+                },
+                "testcases": [
+                    {
+                        "name": "test teardown hooks",
+                        "request": {
+                            "url": "http://127.0.0.1:3458/headers",
+                            "method": "GET",
+                            "data": "abc"
+                        },
+                        "teardown_hooks": [
+                            "${alter_response($response)}"
+                        ],
+                        "validate": [
+                            {"eq": ["status_code", 500]},
+                            {"eq": ["headers.content-type", "html/text"]},
+                            {"eq": ["json.headers.Host", "127.0.0.1:8888"]},
+                            {"eq": ["content.headers.Host", "127.0.0.1:8888"]},
+                            {"eq": ["text.headers.Host", "127.0.0.1:8888"]},
+                            {"eq": ["attributes.new_attribute", "new_attribute"]}
+                        ]
+                    }
+                ]
+            }
+        ]
+        runner = HttpRunner().run(testsets)
+        summary = runner.summary
+        self.assertTrue(summary["success"])
+
+    def test_run_httprunner_with_teardown_hooks_not_exist_attribute(self):
+        testsets = [
+            {
+                "name": "test teardown hooks",
+                "config": {
+                    'path': 'tests/httpbin/hooks.yml',
+                },
+                "testcases": [
+                    {
+                        "name": "test teardown hooks",
+                        "request": {
+                            "url": "http://127.0.0.1:3458/headers",
+                            "method": "GET",
+                            "data": "abc"
+                        },
+                        "teardown_hooks": [
+                            "${alter_response($response)}"
+                        ],
+                        "validate": [
+                            {"eq": ["attributes.attribute_not_exist", "new_attribute"]}
+                        ]
+                    }
+                ]
+            }
+        ]
+        # with self.assertRaises(exceptions.AssertionError):
+        runner = HttpRunner().run(testsets)
+
+        summary = runner.summary
+        self.assertFalse(summary["success"])
+
     def test_run_testset_with_teardown_hooks_success(self):
         test = {
             "name": "get token",