Added tests for multiple column suggestion, dot and alias.

1 files changed, 100 insertions, 0 deletions

diff --git a/tests/test_smart_completion.py b/tests/test_smart_completion.py
index 9ea98975..8ec0925e 100644
--- a/tests/test_smart_completion.py
+++ b/tests/test_smart_completion.py
@@ -47,7 +47,34 @@ def test_function_name_completion(completer, complete_event):
         complete_event)
     assert set(result) == set([Completion(text='MAX', start_position=-2)])
 
+def test_suggested_column_names(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)] +
+        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 users'
     position = len('SELECT MAX(')
     result = completer.get_completions(
@@ -61,6 +88,12 @@ def test_suggested_column_names_in_function(completer, complete_event):
         Completion(text='last_name', start_position=0)])
 
 def test_suggested_column_names_with_dot(completer, complete_event):
+    """
+    Suggest column names on table name and dot
+    :param completer:
+    :param complete_event:
+    :return:
+    """
     text = 'SELECT users. from users'
     position = len('SELECT users.')
     result = set(completer.get_completions(
@@ -74,6 +107,12 @@ def test_suggested_column_names_with_dot(completer, complete_event):
         Completion(text='last_name', start_position=0)])
 
 def test_suggested_column_names_with_alias(completer, complete_event):
+    """
+    Suggest column names on table alias and dot
+    :param completer:
+    :param complete_event:
+    :return:
+    """
     text = 'SELECT u. from users u'
     position = len('SELECT u.')
     result = set(completer.get_completions(
@@ -85,3 +124,64 @@ def test_suggested_column_names_with_alias(completer, complete_event):
         Completion(text='email', start_position=0),
         Completion(text='first_name', start_position=0),
         Completion(text='last_name', 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 users u'
+    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='email', start_position=0),
+        Completion(text='first_name', start_position=0),
+        Completion(text='last_name', start_position=0)] +
+        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 u.id, u. from users u'
+    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='email', start_position=0),
+        Completion(text='first_name', start_position=0),
+        Completion(text='last_name', start_position=0)])
+
+def test_suggested_multiple_column_names_with_dot(completer, complete_event):
+    """
+    Suggest column names on table names and dot
+    when selecting multiple columns from table
+    :param completer:
+    :param complete_event:
+    :return:
+    """
+    text = 'SELECT users.id, users. from users u'
+    position = len('SELECT users.id, users.')
+    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)])