Split up smart completion tests into two files

One file for simple public schema only tests, and another more focused on schema-aware autocompletion

2 files changed, 235 insertions, 45 deletions

diff --git a/tests/test_smart_completion_multiple_schemata.py b/tests/test_smart_completion_multiple_schemata.py
new file mode 100644
index 00000000..d36f4c5a
--- /dev/null
+++ b/tests/test_smart_completion_multiple_schemata.py
@@ -0,0 +1,227 @@
+import pytest
+from prompt_toolkit.completion import Completion
+from prompt_toolkit.document import Document
+
+metadata = {
+            'public':   {
+                            'users': ['id', 'email', 'first_name', 'last_name'],
+                            'orders': ['id', 'ordered_date', 'status'],
+                            'select': ['id', 'insert', 'ABC']
+                        },
+            'custom':   {
+                            'users': ['id', 'phone_number'],
+                            'products': ['id', 'product_name', 'price'],
+                            'shipments': ['id', 'address', 'user_id']
+                        }
+            }
+
+@pytest.fixture
+def completer():
+
+    import pgcli.pgcompleter as pgcompleter
+    comp = pgcompleter.PGCompleter(smart_completion=True)
+
+    schemata, tables, columns = [], [], []
+
+    for schema, tbls in metadata.items():
+        schemata.append(schema)
+
+        for table, cols in tbls.items():
+            tables.append((schema, table))
+            columns.extend([(schema, table, col) for col in cols])
+
+    comp.extend_schemata(schemata)
+    comp.extend_tables(tables)
+    comp.extend_columns(columns)
+    comp.set_search_path(['public'])
+
+    return comp
+
+@pytest.fixture
+def complete_event():
+    from mock import Mock
+    return Mock()
+
+def test_schema_or_visible_table_completion(completer, complete_event):
+    text = 'SELECT * FROM '
+    position = len(text)
+    result = completer.get_completions(
+        Document(text=text, cursor_position=position), complete_event)
+    assert set(result) == set([Completion(text='public', start_position=0),
+                               Completion(text='custom', start_position=0),
+                               Completion(text='users', start_position=0),
+                               Completion(text='"select"', start_position=0),
+                               Completion(text='orders', start_position=0)])
+
+def test_suggested_column_names_from_shadowed_visible_table(completer, complete_event):
+    """
+    Suggest column and function names when selecting from table
+    :param completer:
+    :param complete_event:
+    :return:
+    """
+    text = 'SELECT  from users'
+    position = len('SELECT ')
+    result = set(completer.get_completions(
+        Document(text=text, cursor_position=position),
+        complete_event))
+    assert set(result) == set([
+        Completion(text='*', start_position=0),
+        Completion(text='id', start_position=0),
+        Completion(text='email', start_position=0),
+        Completion(text='first_name', start_position=0),
+        Completion(text='last_name', start_position=0)] +
+        list(map(Completion, completer.functions)))
+
+def test_suggested_column_names_from_qualified_shadowed_table(completer, complete_event):
+    text = 'SELECT  from custom.users'
+    position = len('SELECT ')
+    result = set(completer.get_completions(
+        Document(text=text, cursor_position=position),
+        complete_event))
+    assert set(result) == set([
+        Completion(text='*', start_position=0),
+        Completion(text='id', start_position=0),
+        Completion(text='phone_number', start_position=0)] +
+        list(map(Completion, completer.functions)))
+
+def test_suggested_column_names_from_schema_qualifed_table(completer, complete_event):
+    """
+    Suggest column and function names when selecting from a qualified-table
+    :param completer:
+    :param complete_event:
+    :return:
+    """
+    text = 'SELECT  from custom.products'
+    position = len('SELECT ')
+    result = set(completer.get_completions(
+        Document(text=text, cursor_position=position), complete_event))
+    assert set(result) == set([
+        Completion(text='*', start_position=0),
+        Completion(text='id', start_position=0),
+        Completion(text='product_name', start_position=0),
+        Completion(text='price', start_position=0)] +
+        list(map(Completion, completer.functions)))
+
+def test_suggested_column_names_in_function(completer, complete_event):
+    """
+    Suggest column and function names when selecting multiple
+    columns from table
+    :param completer:
+    :param complete_event:
+    :return:
+    """
+    text = 'SELECT MAX( from custom.products'
+    position = len('SELECT MAX(')
+    result = completer.get_completions(
+        Document(text=text, cursor_position=position),
+        complete_event)
+    assert set(result) == set([
+        Completion(text='*', start_position=0),
+        Completion(text='id', start_position=0),
+        Completion(text='product_name', start_position=0),
+        Completion(text='price', start_position=0)])
+
+def test_suggested_table_names_with_schema_dot(completer, complete_event):
+    text = 'SELECT * FROM custom.'
+    position = len(text)
+    result = completer.get_completions(
+        Document(text=text, cursor_position=position), complete_event)
+    assert set(result) == set([
+        Completion(text='users', start_position=0),
+        Completion(text='products', start_position=0),
+        Completion(text='shipments', start_position=0)])
+
+def test_suggested_column_names_with_qualified_alias(completer, complete_event):
+    """
+    Suggest column names on table alias and dot
+    :param completer:
+    :param complete_event:
+    :return:
+    """
+    text = 'SELECT p. from custom.products p'
+    position = len('SELECT p.')
+    result = set(completer.get_completions(
+        Document(text=text, cursor_position=position),
+        complete_event))
+    assert set(result) == set([
+        Completion(text='*', start_position=0),
+        Completion(text='id', start_position=0),
+        Completion(text='product_name', start_position=0),
+        Completion(text='price', start_position=0)])
+
+def test_suggested_multiple_column_names(completer, complete_event):
+    """
+    Suggest column and function names when selecting multiple
+    columns from table
+    :param completer:
+    :param complete_event:
+    :return:
+    """
+    text = 'SELECT id,  from custom.products'
+    position = len('SELECT id, ')
+    result = set(completer.get_completions(
+        Document(text=text, cursor_position=position),
+        complete_event))
+    assert set(result) == set([
+        Completion(text='*', start_position=0),
+        Completion(text='id', start_position=0),
+        Completion(text='product_name', start_position=0),
+        Completion(text='price', start_position=0)] +
+        list(map(Completion, completer.functions)))
+
+def test_suggested_multiple_column_names_with_alias(completer, complete_event):
+    """
+    Suggest column names on table alias and dot
+    when selecting multiple columns from table
+    :param completer:
+    :param complete_event:
+    :return:
+    """
+    text = 'SELECT p.id, p. from custom.products p'
+    position = len('SELECT u.id, u.')
+    result = set(completer.get_completions(
+        Document(text=text, cursor_position=position),
+        complete_event))
+    assert set(result) == set([
+        Completion(text='*', start_position=0),
+        Completion(text='id', start_position=0),
+        Completion(text='product_name', start_position=0),
+        Completion(text='price', start_position=0)])
+
+def test_suggested_aliases_after_on(completer, complete_event):
+    text = 'SELECT x.id, y.product_name FROM custom.products x JOIN custom.products y ON '
+    position = len(text)
+    result = set(completer.get_completions(
+        Document(text=text, cursor_position=position),
+        complete_event))
+    assert set(result) == set([
+        Completion(text='x', start_position=0),
+        Completion(text='y', start_position=0)])
+
+def test_suggested_aliases_after_on_right_side(completer, complete_event):
+    text = 'SELECT x.id, y.product_name FROM custom.products x JOIN custom.products y ON x.id = '
+    position = len(text)
+    result = set(completer.get_completions(
+        Document(text=text, cursor_position=position),
+        complete_event))
+    assert set(result) == set([
+        Completion(text='x', start_position=0),
+        Completion(text='y', start_position=0)])
+
+def test_table_names_after_from(completer, complete_event):
+    text = 'SELECT * FROM '
+    position = len('SELECT * FROM ')
+    result = set(completer.get_completions(
+        Document(text=text, cursor_position=position),
+        complete_event))
+    assert set(result) == set([
+        Completion(text='public', start_position=0),
+        Completion(text='custom', start_position=0),
+        Completion(text='users', start_position=0),
+        Completion(text='orders', start_position=0),
+        Completion(text='"select"', start_position=0),
+        ])
+
+
+
diff --git a/tests/test_smart_completion.py b/tests/test_smart_completion_public_schema_only.py
index 2d34a218..4c64c825 100644
--- a/tests/test_smart_completion.py
+++ b/tests/test_smart_completion_public_schema_only.py
@@ -3,15 +3,9 @@ from prompt_toolkit.completion import Completion
 from prompt_toolkit.document import Document
 
 metadata = {
-            'public':   {
-                            'users': ['id', 'email', 'first_name', 'last_name'],
-                            'orders': ['id', 'ordered_date', 'status'],
-                            'select': ['id', 'insert', 'ABC'],
-                        },
-            'custom':   {
-                            'products': ['id', 'product_name', 'price'],
-                            'shipments': ['id', 'address', 'user_id']
-                        }
+                'users': ['id', 'email', 'first_name', 'last_name'],
+                'orders': ['id', 'ordered_date', 'status'],
+                'select': ['id', 'insert', 'ABC']
             }
 
 @pytest.fixture
@@ -20,14 +14,12 @@ def completer():
     import pgcli.pgcompleter as pgcompleter
     comp = pgcompleter.PGCompleter(smart_completion=True)
 
-    schemata, tables, columns = [], [], []
+    schemata = ['public']
+    tables, columns = [], []
 
-    for schema, tbls in metadata.items():
-        schemata.append(schema)
-
-        for table, cols in tbls.items():
-            tables.append((schema, table))
-            columns.extend([(schema, table, col) for col in cols])
+    for table, cols in metadata.items():
+        tables.append(('public', table))
+        columns.extend([('public', table, col) for col in cols])
 
     comp.extend_schemata(schemata)
     comp.extend_tables(tables)
@@ -65,7 +57,6 @@ 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([Completion(text='public', start_position=0),
-                               Completion(text='custom', start_position=0),
                                Completion(text='users', start_position=0),
                                Completion(text='"select"', start_position=0),
                                Completion(text='orders', start_position=0)])
@@ -98,24 +89,6 @@ def test_suggested_column_names_from_visible_table(completer, complete_event):
         Completion(text='last_name', start_position=0)] +
         list(map(Completion, completer.functions)))
 
-def test_suggested_column_names_from_schema_qualifed_table(completer, complete_event):
-    """
-    Suggest column and function names when selecting from a qualified-table
-    :param completer:
-    :param complete_event:
-    :return:
-    """
-    text = 'SELECT  from custom.products'
-    position = len('SELECT ')
-    result = set(completer.get_completions(
-        Document(text=text, cursor_position=position), complete_event))
-    assert set(result) == set([
-        Completion(text='*', start_position=0),
-        Completion(text='id', start_position=0),
-        Completion(text='product_name', start_position=0),
-        Completion(text='price', start_position=0)] +
-        list(map(Completion, completer.functions)))
-
 def test_suggested_column_names_in_function(completer, complete_event):
     """
     Suggest column and function names when selecting multiple
@@ -155,15 +128,6 @@ def test_suggested_column_names_with_table_dot(completer, complete_event):
         Completion(text='first_name', start_position=0),
         Completion(text='last_name', start_position=0)])
 
-def test_suggested_table_names_with_schema_dot(completer, complete_event):
-    text = 'SELECT * FROM custom.'
-    position = len(text)
-    result = completer.get_completions(
-        Document(text=text, cursor_position=position), complete_event)
-    assert set(result) == set([
-        Completion(text='products', start_position=0),
-        Completion(text='shipments', start_position=0)])
-
 def test_suggested_column_names_with_alias(completer, complete_event):
     """
     Suggest column names on table alias and dot
@@ -292,7 +256,6 @@ def test_table_names_after_from(completer, complete_event):
         complete_event))
     assert set(result) == set([
         Completion(text='public', start_position=0),
-        Completion(text='custom', start_position=0),
         Completion(text='users', start_position=0),
         Completion(text='orders', start_position=0),
         Completion(text='"select"', start_position=0),