From 415f45adf48e859750fadc4785bbac07235bf158 Mon Sep 17 00:00:00 2001
From: Honza Kral <honza.kral@gmail.com>
Date: Wed, 10 Jul 2013 16:43:20 +0200
Subject: [PATCH] Support for set and is_true yaml commands

---
 test_elasticsearch/test_server/test_common.py | 58 ++++++++++++-------
 1 file changed, 37 insertions(+), 21 deletions(-)

diff --git a/test_elasticsearch/test_server/test_common.py b/test_elasticsearch/test_server/test_common.py
index b9757e53..33b6444c 100644
--- a/test_elasticsearch/test_server/test_common.py
+++ b/test_elasticsearch/test_server/test_common.py
@@ -26,6 +26,7 @@ class YamlTestCase(TestCase):
     def setUp(self):
         self.client = Elasticsearch(['localhost:9900'])
         self.last_response = None
+        self._state = {}
 
     def run_code(self, test):
         """ Execute an instruction based on it's type. """
@@ -58,6 +59,10 @@ class YamlTestCase(TestCase):
             if k in args:
                 args[PARAMS_RENAMES[k]] = args.pop(k)
 
+        # resolve vars
+        for k in args:
+            args[k] = self._resolve(args[k])
+
         try:
             self.last_response = api(**args)
         except:
@@ -71,37 +76,49 @@ class YamlTestCase(TestCase):
     def run_catch(self, catch):
         pass
 
+    def set_state(self, key, value):
+        self._state[key] = value
+
+    def run_set(self, action):
+        for key, value in action.items():
+            self.run_match({key: None}, lambda x: self.set_state(value, x))
+
+    def run_is_true(self, action):
+        self.run_match({action: True}, bool)
+
     def run_length(self, action):
-        self.run_is(action, len)
+        self.run_match(action, len)
+
+    def _resolve(self, value):
+        # resolve variables
+        if isinstance(value, (type(u''), type(''))) and value.startswith('$'):
+            value = value[1:]
+            self.assertIn(value, self._state)
+            value = self._state[value]
+        return value
 
     def run_match(self, action, transform=None):
         """ Match part of last response to test data. """
 
-        # matching part of the reponse dict
-        if isinstance(action, dict):
-            self.assertEquals(1, len(action))
-            path, expected = list(action.items())[0]
+        self.assertEquals(1, len(action))
+        path, expected = list(action.items())[0]
 
-            # fetch the possibly nested value from last_response
-            value = self.last_response
-            for step in path.split('.'):
-                if step.isdigit():
-                    step = int(step)
-                    self.assertIsInstance(value, list)
-                    self.assertGreater(len(value), step)
-                else:
-                    self.assertIn(step, value)
-                value = value[step]
-
-        # matching the entire response
-        else:
-            value = self.last_response
-            expected = action
+        # fetch the possibly nested value from last_response
+        value = self.last_response
+        for step in path.split('.'):
+            if step.isdigit():
+                step = int(step)
+                self.assertIsInstance(value, list)
+                self.assertGreater(len(value), step)
+            else:
+                self.assertIn(step, value)
+            value = value[step]
 
         # sometimes we need to transform the json value before comparing
         if transform:
             value = transform(value)
 
+        expected = self._resolve(expected)
         # compare target value
         self.assertEquals(expected, value)
 
@@ -112,7 +129,6 @@ class YamlTestCase(TestCase):
     def test_from_yaml(self):
         for test in self._definition:
             for name, definition in test.items():
-                print name
                 self.run_code(definition)