In completion tests, get functions from metadata instead of manually listing them

3 files changed, 24 insertions, 53 deletions

diff --git a/tests/metadata.py b/tests/metadata.py
index 6b9cc5a4..fff424a4 100644
--- a/tests/metadata.py
+++ b/tests/metadata.py
@@ -48,6 +48,10 @@ class MetaData(object):
         return [view(escape(x), pos)
             for x in self.metadata.get('views', {}).get(schema, [])]
 
+    def functions(self, schema='public', pos=0):
+        return [function(escape(x[0]), pos)
+            for x in self.metadata.get('functions', {}).get(schema, [])]
+
     def schemas(self, pos=0):
         schemas = set(sch for schs in self.metadata.values() for sch in schs)
         return [schema(escape(s), pos=pos) for s in schemas]
diff --git a/tests/test_smart_completion_multiple_schemata.py b/tests/test_smart_completion_multiple_schemata.py
index ff176305..86486f63 100644
--- a/tests/test_smart_completion_multiple_schemata.py
+++ b/tests/test_smart_completion_multiple_schemata.py
@@ -63,10 +63,7 @@ def test_schema_or_visible_table_completion(completer, complete_event):
     position = len(text)
     result = completer.get_completions(
         Document(text=text, cursor_position=position), complete_event)
-    assert set(result) == set(testdata.schemas() + [
-       function('func1'),
-       function('func2'),
-       ] + testdata.tables())
+    assert set(result) == set(testdata.schemas() + testdata.functions() + testdata.tables())
 
 
 @pytest.mark.parametrize('table', [
@@ -91,8 +88,7 @@ def test_suggested_column_names_from_shadowed_visible_table(completer, complete_
         column('email'),
         column('first_name'),
         column('last_name'),
-        function('func1'),
-        function('func2')] +
+        ] + testdata.functions() +
         list(testdata.builtin_functions() +
         testdata.keywords())
         )
@@ -106,8 +102,7 @@ def test_suggested_column_names_from_qualified_shadowed_table(completer, complet
     assert set(result) == set([
         column('id'),
         column('phone_number'),
-        function('func1'),
-        function('func2')] +
+        ] + testdata.functions() +
         list(testdata.builtin_functions() +
         testdata.keywords())
         )
@@ -143,8 +138,7 @@ def test_suggested_joins(completer, complete_event, query, tbl):
         complete_event))
     assert set(result) == set(testdata.schemas() + testdata.tables() + [
         join('custom.shipments ON shipments.user_id = {0}.id'.format(tbl)),
-        function('func1'),
-        function('func2')])
+        ] + testdata.functions())
 
 def test_suggested_column_names_from_schema_qualifed_table(completer, complete_event):
     """
@@ -161,8 +155,7 @@ def test_suggested_column_names_from_schema_qualifed_table(completer, complete_e
         column('id'),
         column('product_name'),
         column('price'),
-        function('func1'),
-        function('func2')] +
+        ] + testdata.functions() +
         list(testdata.builtin_functions() +
         testdata.keywords())
         )
@@ -203,10 +196,8 @@ def test_suggested_table_names_with_schema_dot(completer, complete_event,
     position = len(text)
     result = completer.get_completions(
         Document(text=text, cursor_position=position), complete_event)
-    assert set(result) == set(testdata.tables('custom', start_pos) + [
-        function('func3', start_pos),
-        function('set_returning_func', start_pos),
-    ])
+    assert set(result) == set(testdata.tables('custom', start_pos)
+        + testdata.functions('custom', start_pos))
 
 @pytest.mark.parametrize('text', [
     'SELECT * FROM "Custom".',
@@ -223,10 +214,8 @@ def test_suggested_table_names_with_schema_dot2(completer, complete_event,
     position = len(text)
     result = completer.get_completions(
         Document(text=text, cursor_position=position), complete_event)
-    assert set(result) == set([
-        function('func4', start_pos)]
-        + testdata.tables('Custom', start_pos)
-    )
+    assert set(result) == set(testdata.functions('Custom', start_pos) +
+        testdata.tables('Custom', start_pos))
 
 def test_suggested_column_names_with_qualified_alias(completer, complete_event):
     """
@@ -262,8 +251,7 @@ def test_suggested_multiple_column_names(completer, complete_event):
         column('id'),
         column('product_name'),
         column('price'),
-        function('func1'),
-        function('func2')] +
+        ] + testdata.functions() +
         list(testdata.builtin_functions() +
         testdata.keywords())
         )
@@ -318,10 +306,8 @@ def test_table_names_after_from(completer, complete_event):
     result = set(completer.get_completions(
         Document(text=text, cursor_position=position),
         complete_event))
-    assert set(result) == set(testdata.schemas() + testdata.tables() + [
-        function('func1'),
-        function('func2'),
-        ])
+    assert set(result) == set(testdata.schemas() + testdata.tables()
+        + testdata.functions())
 
 def test_schema_qualified_function_name(completer, complete_event):
     text = 'SELECT custom.func'
diff --git a/tests/test_smart_completion_public_schema_only.py b/tests/test_smart_completion_public_schema_only.py
index 580f567b..9794c2e7 100644
--- a/tests/test_smart_completion_public_schema_only.py
+++ b/tests/test_smart_completion_public_schema_only.py
@@ -67,10 +67,7 @@ def test_schema_or_visible_table_completion(completer, complete_event):
     result = completer.get_completions(
         Document(text=text, cursor_position=position), complete_event)
     assert set(result) == set(testdata.schemas()
-        + testdata.views() + testdata.tables() + [
-        function('custom_func1'),
-        function('custom_func2'),
-        function('set_returning_func')])
+        + testdata.views() + testdata.tables() + testdata.functions())
 
 
 def test_builtin_function_name_completion(completer, complete_event):
@@ -135,9 +132,7 @@ def test_suggested_column_names_from_visible_table(completer, complete_event):
         column('email'),
         column('first_name'),
         column('last_name'),
-        function('custom_func1'),
-        function('custom_func2'),
-    function('set_returning_func')] +
+        ] + testdata.functions() +
         list(testdata.builtin_functions() +
         testdata.keywords())
         )
@@ -220,9 +215,7 @@ def test_suggested_multiple_column_names(completer, complete_event):
         column('email'),
         column('first_name'),
         column('last_name'),
-        function('custom_func1'),
-        function('custom_func2'),
-        function('set_returning_func')] +
+        ] + testdata.functions() +
         list(testdata.builtin_functions() +
         testdata.keywords())
         )
@@ -377,9 +370,7 @@ def test_suggested_joins(completer, complete_event, text):
         join('"Users" ON "Users".userid = users.id'),
         join('users users2 ON users2.id = users.parentid'),
         join('users users2 ON users2.parentid = users.id'),
-        function('custom_func2'),
-        function('set_returning_func'),
-        function('custom_func1')])
+        ] + testdata.functions())
 
 @pytest.mark.parametrize('text', [
     'SELECT * FROM public."Users" JOIN ',
@@ -396,9 +387,7 @@ def test_suggested_joins_quoted_schema_qualified_table(completer, complete_event
     assert set(result) == set(testdata.schemas() + testdata.tables()
         + testdata.views() + [
         join('public.users ON users.id = "Users".userid'),
-        function('custom_func2'),
-        function('set_returning_func'),
-        function('custom_func1')])
+        ] + testdata.functions())
 
 @pytest.mark.parametrize('text', [
     'SELECT u.name, o.id FROM users u JOIN orders o ON ',
@@ -508,11 +497,7 @@ def test_table_names_after_from(completer, complete_event, text):
         Document(text=text, cursor_position=position),
         complete_event)
     assert set(result) == set(testdata.schemas() + testdata.tables()
-        + testdata.views() + [
-        function('custom_func1'),
-        function('custom_func2'),
-        function('set_returning_func')
-        ])
+        + testdata.views() + testdata.functions())
     assert [c.text for c in result] == [
         '"Users"',
         'custom_func1',
@@ -535,9 +520,7 @@ def test_auto_escaped_col_names(completer, complete_event):
         column('id'),
         column('"insert"'),
         column('"ABC"'),
-        function('custom_func1'),
-        function('custom_func2'),
-        function('set_returning_func')] +
+        ] + testdata.functions() +
         list(testdata.builtin_functions() +
         testdata.keywords())
         )
@@ -588,9 +571,7 @@ def test_suggest_columns_from_set_returning_function(completer, complete_event):
     assert set(result) == set([
         column('x'),
         column('y'),
-        function('custom_func1'),
-        function('custom_func2'),
-        function('set_returning_func')]
+        ] + testdata.functions()
          + list(testdata.builtin_functions()
          + testdata.keywords()))