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

3 files changed, 44 insertions, 123 deletions

diff --git a/tests/metadata.py b/tests/metadata.py
index fff424a4..05f90fa9 100644
--- a/tests/metadata.py
+++ b/tests/metadata.py
@@ -36,6 +36,14 @@ class MetaData(object):
     def keywords(self, pos=0):
         return [keyword(kw, pos) for kw in self.completer.keywords]
 
+    def columns(self, parent, schema='public', typ='tables', pos=0):
+        if typ == 'functions':
+            fun = [x for x in self.metadata[typ][schema] if x[0] == parent][0]
+            cols = fun[1]
+        else:
+            cols = self.metadata[typ][schema][parent]
+        return [column(escape(col), pos) for col in cols]
+
     def datatypes(self, schema='public', pos=0):
         return [datatype(escape(x), pos)
             for x in self.metadata.get('datatypes', {}).get(schema, [])]
diff --git a/tests/test_smart_completion_multiple_schemata.py b/tests/test_smart_completion_multiple_schemata.py
index 86486f63..e87c9f1f 100644
--- a/tests/test_smart_completion_multiple_schemata.py
+++ b/tests/test_smart_completion_multiple_schemata.py
@@ -2,9 +2,7 @@ from __future__ import unicode_literals
 import pytest
 import itertools
 from metadata import (MetaData, alias, name_join, fk_join, join,
-    function,
-    column,
-    wildcard_expansion)
+    function, wildcard_expansion)
 from prompt_toolkit.document import Document
 from pgcli.packages.function_metadata import FunctionMetadata, ForeignKey
 
@@ -83,12 +81,8 @@ def test_suggested_column_names_from_shadowed_visible_table(completer, complete_
         Document(text=text, cursor_position=position),
         complete_event))
 
-    assert set(result) == set([
-        column('id'),
-        column('email'),
-        column('first_name'),
-        column('last_name'),
-        ] + testdata.functions() +
+    assert set(result) == set(testdata.columns('users') +
+        testdata.functions() +
         list(testdata.builtin_functions() +
         testdata.keywords())
         )
@@ -99,10 +93,8 @@ def test_suggested_column_names_from_qualified_shadowed_table(completer, complet
     result = set(completer.get_completions(
         Document(text=text, cursor_position=position),
         complete_event))
-    assert set(result) == set([
-        column('id'),
-        column('phone_number'),
-        ] + testdata.functions() +
+    assert set(result) == set(testdata.columns('users', 'custom') +
+        testdata.functions() +
         list(testdata.builtin_functions() +
         testdata.keywords())
         )
@@ -151,11 +143,7 @@ def test_suggested_column_names_from_schema_qualifed_table(completer, complete_e
     position = len('SELECT ')
     result = set(completer.get_completions(
         Document(text=text, cursor_position=position), complete_event))
-    assert set(result) == set([
-        column('id'),
-        column('product_name'),
-        column('price'),
-        ] + testdata.functions() +
+    assert set(result) == set(testdata.columns('products', 'custom') + testdata.functions() +
         list(testdata.builtin_functions() +
         testdata.keywords())
         )
@@ -173,10 +161,7 @@ def test_suggested_column_names_in_function(completer, complete_event):
     result = completer.get_completions(
         Document(text=text, cursor_position=position),
         complete_event)
-    assert set(result) == set([
-        column('id'),
-        column('product_name'),
-        column('price')])
+    assert set(result) == set(testdata.columns('products', 'custom'))
 
 
 @pytest.mark.parametrize('text', [
@@ -229,10 +214,7 @@ def test_suggested_column_names_with_qualified_alias(completer, complete_event):
     result = set(completer.get_completions(
         Document(text=text, cursor_position=position),
         complete_event))
-    assert set(result) == set([
-        column('id'),
-        column('product_name'),
-        column('price')])
+    assert set(result) == set(testdata.columns('products', 'custom'))
 
 def test_suggested_multiple_column_names(completer, complete_event):
     """
@@ -247,11 +229,8 @@ def test_suggested_multiple_column_names(completer, complete_event):
     result = set(completer.get_completions(
         Document(text=text, cursor_position=position),
         complete_event))
-    assert set(result) == set([
-        column('id'),
-        column('product_name'),
-        column('price'),
-        ] + testdata.functions() +
+    assert set(result) == set(testdata.columns('products', 'custom') +
+        testdata.functions() +
         list(testdata.builtin_functions() +
         testdata.keywords())
         )
@@ -269,10 +248,7 @@ def test_suggested_multiple_column_names_with_alias(completer, complete_event):
     result = set(completer.get_completions(
         Document(text=text, cursor_position=position),
         complete_event))
-    assert set(result) == set([
-        column('id'),
-        column('product_name'),
-        column('price')])
+    assert set(result) == set(testdata.columns('products', 'custom'))
 
 @pytest.mark.parametrize('text', [
     'SELECT x.id, y.product_name FROM custom.products x JOIN custom.products y ON ',
@@ -338,8 +314,8 @@ def test_suggest_columns_from_aliased_set_returning_function(completer, complete
     pos = len('select f.')
     result = completer.get_completions(Document(text=sql, cursor_position=pos),
                                        complete_event)
-    assert set(result) == set([
-        column('x')])
+    assert set(result) == set(
+        testdata.columns('set_returning_func', 'custom', 'functions'))
 
 @pytest.mark.parametrize('text', [
     'SELECT * FROM custom.set_returning_func()',
@@ -445,9 +421,7 @@ def test_suggest_columns_from_unquoted_table(completer, complete_event, text):
     pos = len('SELECT U.')
     result = completer.get_completions(Document(text=text, cursor_position=pos),
                                        complete_event)
-    assert set(result) == set([
-        column('id'),
-        column('phone_number')])
+    assert set(result) == set(testdata.columns('users', 'custom'))
 
 @pytest.mark.parametrize('text', [
     'SELECT U. FROM custom."Users" U',
@@ -457,6 +431,4 @@ def test_suggest_columns_from_quoted_table(completer, complete_event, text):
     pos = len('SELECT U.')
     result = completer.get_completions(Document(text=text, cursor_position=pos),
                                        complete_event)
-    assert set(result) == set([
-        column('userid'),
-        column('username')])
+    assert set(result) == set(testdata.columns('Users', 'custom'))
diff --git a/tests/test_smart_completion_public_schema_only.py b/tests/test_smart_completion_public_schema_only.py
index 9794c2e7..2b722e51 100644
--- a/tests/test_smart_completion_public_schema_only.py
+++ b/tests/test_smart_completion_public_schema_only.py
@@ -1,10 +1,7 @@
 from __future__ import unicode_literals
 import pytest
 from metadata import (MetaData, alias, name_join, fk_join, join, keyword,
-    table,
-    function,
-    column,
-    wildcard_expansion)
+    table, function, column, wildcard_expansion)
 from prompt_toolkit.document import Document
 from pgcli.packages.function_metadata import FunctionMetadata, ForeignKey
 
@@ -126,13 +123,7 @@ def test_suggested_column_names_from_visible_table(completer, complete_event):
     result = set(completer.get_completions(
         Document(text=text, cursor_position=position),
         complete_event))
-    assert set(result) == set([
-        column('id'),
-        column('parentid'),
-        column('email'),
-        column('first_name'),
-        column('last_name'),
-        ] + testdata.functions() +
+    assert set(result) == set(testdata.columns('users') + testdata.functions() +
         list(testdata.builtin_functions() +
         testdata.keywords())
         )
@@ -151,12 +142,7 @@ def test_suggested_column_names_in_function(completer, complete_event):
     result = completer.get_completions(
         Document(text=text, cursor_position=position),
         complete_event)
-    assert set(result) == set([
-        column('id'),
-        column('parentid'),
-        column('email'),
-        column('first_name'),
-        column('last_name')])
+    assert set(result) == set(testdata.columns('users'))
 
 def test_suggested_column_names_with_table_dot(completer, complete_event):
     """
@@ -170,12 +156,7 @@ def test_suggested_column_names_with_table_dot(completer, complete_event):
     result = set(completer.get_completions(
         Document(text=text, cursor_position=position),
         complete_event))
-    assert set(result) == set([
-        column('id'),
-        column('parentid'),
-        column('email'),
-        column('first_name'),
-        column('last_name')])
+    assert set(result) == set(testdata.columns('users'))
 
 def test_suggested_column_names_with_alias(completer, complete_event):
     """
@@ -189,12 +170,7 @@ def test_suggested_column_names_with_alias(completer, complete_event):
     result = set(completer.get_completions(
         Document(text=text, cursor_position=position),
         complete_event))
-    assert set(result) == set([
-        column('id'),
-        column('parentid'),
-        column('email'),
-        column('first_name'),
-        column('last_name')])
+    assert set(result) == set(testdata.columns('users'))
 
 def test_suggested_multiple_column_names(completer, complete_event):
     """
@@ -209,13 +185,7 @@ def test_suggested_multiple_column_names(completer, complete_event):
     result = set(completer.get_completions(
         Document(text=text, cursor_position=position),
         complete_event))
-    assert set(result) == set([
-        column('id'),
-        column('parentid'),
-        column('email'),
-        column('first_name'),
-        column('last_name'),
-        ] + testdata.functions() +
+    assert set(result) == set(testdata.columns('users') + testdata.functions() +
         list(testdata.builtin_functions() +
         testdata.keywords())
         )
@@ -233,12 +203,7 @@ def test_suggested_multiple_column_names_with_alias(completer, complete_event):
     result = set(completer.get_completions(
         Document(text=text, cursor_position=position),
         complete_event))
-    assert set(result) == set([
-        column('id'),
-        column('parentid'),
-        column('email'),
-        column('first_name'),
-        column('last_name')])
+    assert set(result) == set(testdata.columns('users'))
 
 def test_suggested_multiple_column_names_with_dot(completer, complete_event):
     """
@@ -253,12 +218,7 @@ def test_suggested_multiple_column_names_with_dot(completer, complete_event):
     result = set(completer.get_completions(
         Document(text=text, cursor_position=position),
         complete_event))
-    assert set(result) == set([
-        column('id'),
-        column('parentid'),
-        column('email'),
-        column('first_name'),
-        column('last_name')])
+    assert set(result) == set(testdata.columns('users'))
 
 
 def test_suggest_columns_after_three_way_join(completer, complete_event):
@@ -516,10 +476,7 @@ def test_auto_escaped_col_names(completer, complete_event):
     result = set(completer.get_completions(
         Document(text=text, cursor_position=position),
         complete_event))
-    assert set(result) == set([
-        column('id'),
-        column('"insert"'),
-        column('"ABC"'),
+    assert set(result) == set(testdata.columns('select') + [
         ] + testdata.functions() +
         list(testdata.builtin_functions() +
         testdata.keywords())
@@ -556,11 +513,7 @@ def test_suggest_columns_from_escaped_table_alias(completer, complete_event):
     pos = len(sql)
     result = completer.get_completions(Document(text=sql, cursor_position=pos),
                                        complete_event)
-    assert set(result) == set([
-        column('id'),
-        column('"insert"'),
-        column('"ABC"'),
-    ])
+    assert set(result) == set(testdata.columns('select'))
 
 
 def test_suggest_columns_from_set_returning_function(completer, complete_event):
@@ -568,12 +521,11 @@ def test_suggest_columns_from_set_returning_function(completer, complete_event):
     pos = len('select ')
     result = completer.get_completions(Document(text=sql, cursor_position=pos),
                                        complete_event)
-    assert set(result) == set([
-        column('x'),
-        column('y'),
-        ] + testdata.functions()
-         + list(testdata.builtin_functions()
-         + testdata.keywords()))
+    assert set(result) == set(
+        testdata.columns('set_returning_func', typ='functions')
+        + testdata.functions()
+        + list(testdata.builtin_functions()
+        + testdata.keywords()))
 
 
 def test_suggest_columns_from_aliased_set_returning_function(completer, complete_event):
@@ -581,9 +533,7 @@ def test_suggest_columns_from_aliased_set_returning_function(completer, complete
     pos = len('select f.')
     result = completer.get_completions(Document(text=sql, cursor_position=pos),
                                        complete_event)
-    assert set(result) == set([
-        column('x'),
-        column('y')])
+    assert set(result) == set(testdata.columns('set_returning_func', typ='functions'))
 
 
 def test_join_functions_using_suggests_common_columns(completer, complete_event):
@@ -592,9 +542,8 @@ def test_join_functions_using_suggests_common_columns(completer, complete_event)
     pos = len(text)
     result = set(completer.get_completions(
         Document(text=text, cursor_position=pos), complete_event))
-    assert set(result) == set([
-         column('x'),
-         column('y')])
+    assert set(result) == set(
+        testdata.columns('set_returning_func', typ='functions'))
 
 
 def test_join_functions_on_suggests_columns_and_join_conditions(completer, complete_event):
@@ -606,8 +555,7 @@ def test_join_functions_on_suggests_columns_and_join_conditions(completer, compl
     assert set(result) == set([
          name_join('y = f2.y'),
          name_join('x = f2.x'),
-         column('x'),
-         column('y')])
+         ] + testdata.columns('set_returning_func', typ='functions'))
 
 
 def test_learn_keywords(completer, complete_event):
@@ -724,12 +672,7 @@ def test_suggest_columns_from_unquoted_table(completer, complete_event, text):
     pos = len('SELECT U.')
     result = completer.get_completions(Document(text=text, cursor_position=pos),
                                        complete_event)
-    assert set(result) == set([
-        column('id'),
-        column('parentid'),
-        column('email'),
-        column('first_name'),
-        column('last_name')])
+    assert set(result) == set(testdata.columns('users'))
 
 
 def test_suggest_columns_from_quoted_table(completer, complete_event):
@@ -737,6 +680,4 @@ def test_suggest_columns_from_quoted_table(completer, complete_event):
     pos = len('SELECT U.')
     result = completer.get_completions(Document(text=text, cursor_position=pos),
                                        complete_event)
-    assert set(result) == set([
-        column('userid'),
-        column('username')])
+    assert set(result) == set(testdata.columns('Users'))