Query token ancestry.

1 files changed, 24 insertions, 0 deletions

diff --git a/tests/test_parse.py b/tests/test_parse.py
index 3ee2203..d14e329 100644
--- a/tests/test_parse.py
+++ b/tests/test_parse.py
@@ -5,6 +5,7 @@
 from tests.utils import TestCaseBase
 
 import sqlparse
+import sqlparse.sql
 
 
 class SQLParseTest(TestCaseBase):
@@ -37,3 +38,26 @@ class SQLParseTest(TestCaseBase):
         sql = u'select\r\n*from foo\n'
         p = sqlparse.parse(sql)[0]
         self.assertEqual(unicode(p), sql)
+
+    def test_within(self):
+        sql = 'foo(col1, col2)'
+        p = sqlparse.parse(sql)[0]
+        col1 = p.tokens[0].tokens[1].tokens[1].tokens[0]
+        self.assert_(col1.within(sqlparse.sql.Function))
+
+    def test_child_of(self):
+        sql = '(col1, col2)'
+        p = sqlparse.parse(sql)[0]
+        self.assert_(p.tokens[0].tokens[1].is_child_of(p.tokens[0]))
+        sql = 'select foo'
+        p = sqlparse.parse(sql)[0]
+        self.assert_(not p.tokens[2].is_child_of(p.tokens[0]))
+        self.assert_(p.tokens[2].is_child_of(p))
+
+    def test_has_ancestor(self):
+        sql = 'foo or (bar, baz)'
+        p = sqlparse.parse(sql)[0]
+        baz = p.tokens[-1].tokens[1].tokens[-1]
+        self.assert_(baz.has_ancestor(p.tokens[-1].tokens[1]))
+        self.assert_(baz.has_ancestor(p.tokens[-1]))
+        self.assert_(baz.has_ancestor(p))