From 83ca01c0b5fcc844ca364aa0b93038d817dbbae8 Mon Sep 17 00:00:00 2001
From: debugtalk <mail@debugtalk.com>
Date: Wed, 28 Jun 2017 21:25:28 +0800
Subject: [PATCH] query_json: Do an xpath-like query with json_content.

---
 ate/utils.py       | 33 +++++++++++++++++++++++++++++++++
 test/test_utils.py | 40 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 73 insertions(+)

diff --git a/ate/utils.py b/ate/utils.py
index 221ad47a..18a5a660 100644
--- a/ate/utils.py
+++ b/ate/utils.py
@@ -159,3 +159,36 @@ def parse_content_with_variables(content, variables_binds):
         return value
 
     return content
+
+def query_json(json_content, query, delimiter='.'):
+    """ Do an xpath-like query with json_content.
+    @param (json_content) json_content
+        json_content = {
+            "ids": [1, 2, 3, 4],
+            "person": {
+                "name": {
+                    "first_name": "Leo",
+                    "last_name": "Lee",
+                },
+                "age": 29,
+                "cities": ["Guangzhou", "Shenzhen"]
+            }
+        }
+    @param (str) query
+        "person.name.first_name"  =>  "Leo"
+        "person.cities.0"         =>  "Guangzhou"
+    @return queried result
+    """
+    stripped_query = query.strip(delimiter)
+    if not stripped_query:
+        return None
+
+    try:
+        for key in stripped_query.split(delimiter):
+            if isinstance(json_content, list):
+                key = int(key)
+            json_content = json_content[key]
+    except (KeyError, ValueError, IndexError):
+        raise ParamsError("invalid query string in extract_binds!")
+
+    return json_content
diff --git a/test/test_utils.py b/test/test_utils.py
index e33ad439..dd323e58 100644
--- a/test/test_utils.py
+++ b/test/test_utils.py
@@ -139,3 +139,43 @@ class TestUtils(ApiServerUnittest):
         }
         with self.assertRaises(exception.ParamsError):
             utils.parse_content_with_variables(content, variables_binds)
+
+    def test_query_json(self):
+        json_content = {
+            "ids": [1, 2, 3, 4],
+            "person": {
+                "name": {
+                    "first_name": "Leo",
+                    "last_name": "Lee",
+                },
+                "age": 29,
+                "cities": ["Guangzhou", "Shenzhen"]
+            }
+        }
+        query = "ids.2"
+        result = utils.query_json(json_content, query)
+        self.assertEqual(result, 3)
+
+        query = "ids.str_key"
+        with self.assertRaises(exception.ParamsError):
+            utils.query_json(json_content, query)
+
+        query = "ids.5"
+        with self.assertRaises(exception.ParamsError):
+            utils.query_json(json_content, query)
+
+        query = "person.age"
+        result = utils.query_json(json_content, query)
+        self.assertEqual(result, 29)
+
+        query = "person.not_exist_key"
+        with self.assertRaises(exception.ParamsError):
+            utils.query_json(json_content, query)
+
+        query = "person.cities.0"
+        result = utils.query_json(json_content, query)
+        self.assertEqual(result, "Guangzhou")
+
+        query = "person.name.first_name"
+        result = utils.query_json(json_content, query)
+        self.assertEqual(result, "Leo")