Fix failing tests.

4 files changed, 67 insertions, 35 deletions

diff --git a/pgcli/packages/sqlcompletion.py b/pgcli/packages/sqlcompletion.py
index b0136cbd..f4c58eb6 100644
--- a/pgcli/packages/sqlcompletion.py
+++ b/pgcli/packages/sqlcompletion.py
@@ -207,16 +207,6 @@ def suggest_based_on_last_token(token, text_before_cursor, full_text, identifier
             prev_tok = prev_tok.value.lower()
             if prev_tok == 'exists':
                 return [{'type': 'keyword'}]
-            elif prev_tok in ('any', 'some', 'all'):
-                return column_suggestions + [{'type': 'keyword'}]
-            elif prev_tok == 'in':
-                # Technically, we should suggest columns AND keywords, as
-                # per case 4. However, IN is different from ANY, SOME, ALL
-                # in that it can accept a *list* of columns, or a subquery.
-                # But suggesting keywords for , "SELECT * FROM foo WHERE bar IN
-                # (baz, qux, " would be overwhelming. So we special case 'IN'
-                # to not suggest keywords.
-                return column_suggestions
             else:
                 return column_suggestions
 
diff --git a/tests/test_smart_completion_multiple_schemata.py b/tests/test_smart_completion_multiple_schemata.py
index d6aeacc1..421b25e9 100644
--- a/tests/test_smart_completion_multiple_schemata.py
+++ b/tests/test_smart_completion_multiple_schemata.py
@@ -86,6 +86,7 @@ def test_suggested_column_names_from_shadowed_visible_table(completer, complete_
     result = set(completer.get_completions(
         Document(text=text, cursor_position=position),
         complete_event))
+
     assert set(result) == set([
         Completion(text='*', start_position=0, display_meta='column'),
         Completion(text='id', start_position=0, display_meta='column'),
@@ -94,7 +95,9 @@ def test_suggested_column_names_from_shadowed_visible_table(completer, complete_
         Completion(text='last_name', start_position=0, display_meta='column'),
         Completion(text='func1', start_position=0, display_meta='function'),
         Completion(text='func2', start_position=0, display_meta='function')] +
-        list(map(lambda f: Completion(f, display_meta='function'), completer.functions)))
+        list(map(lambda f: Completion(f, display_meta='function'), completer.functions)) +
+        list(map(lambda x: Completion(x, display_meta='keyword'), completer.keywords))
+        )
 
 def test_suggested_column_names_from_qualified_shadowed_table(completer, complete_event):
     text = 'SELECT  from custom.users'
@@ -108,7 +111,9 @@ def test_suggested_column_names_from_qualified_shadowed_table(completer, complet
         Completion(text='phone_number', start_position=0, display_meta='column'),
         Completion(text='func1', start_position=0, display_meta='function'),
         Completion(text='func2', start_position=0, display_meta='function')] +
-        list(map(lambda f: Completion(f, display_meta='function'), completer.functions)))
+        list(map(lambda f: Completion(f, display_meta='function'), completer.functions)) +
+        list(map(lambda x: Completion(x, display_meta='keyword'), completer.keywords))
+        )
 
 def test_suggested_column_names_from_schema_qualifed_table(completer, complete_event):
     """
@@ -128,7 +133,9 @@ def test_suggested_column_names_from_schema_qualifed_table(completer, complete_e
         Completion(text='price', start_position=0, display_meta='column'),
         Completion(text='func1', start_position=0, display_meta='function'),
         Completion(text='func2', start_position=0, display_meta='function')] +
-        list(map(lambda f: Completion(f, display_meta='function'), completer.functions)))
+        list(map(lambda f: Completion(f, display_meta='function'), completer.functions)) +
+        list(map(lambda x: Completion(x, display_meta='keyword'), completer.keywords))
+        )
 
 def test_suggested_column_names_in_function(completer, complete_event):
     """
@@ -197,7 +204,9 @@ def test_suggested_multiple_column_names(completer, complete_event):
         Completion(text='price', start_position=0, display_meta='column'),
         Completion(text='func1', start_position=0, display_meta='function'),
         Completion(text='func2', start_position=0, display_meta='function')] +
-        list(map(lambda f: Completion(f, display_meta='function'), completer.functions)))
+        map(lambda f: Completion(f, display_meta='function'), completer.functions) +
+        map(lambda x: Completion(x, display_meta='keyword'), completer.keywords)
+        )
 
 def test_suggested_multiple_column_names_with_alias(completer, complete_event):
     """
@@ -278,4 +287,4 @@ def test_schema_qualified_type_name(text, completer, complete_event):
         Completion(text='users', display_meta='table'),
         Completion(text='products', display_meta='table'),
         Completion(text='shipments', display_meta='table'),
-        ])
\ No newline at end of file
+        ])
diff --git a/tests/test_smart_completion_public_schema_only.py b/tests/test_smart_completion_public_schema_only.py
index 128ccdd1..2f75069d 100644
--- a/tests/test_smart_completion_public_schema_only.py
+++ b/tests/test_smart_completion_public_schema_only.py
@@ -96,8 +96,9 @@ def test_builtin_function_name_completion(completer, complete_event):
     position = len('SELECT MA')
     result = completer.get_completions(
         Document(text=text, cursor_position=position), complete_event)
-    assert set(result) == set([Completion(text='MAX', start_position=-2,
-                                          display_meta='function')])
+    assert set(result) == set([Completion(text='MAX', start_position=-2, display_meta='function'),
+                               Completion(text='MAXEXTENTS', start_position=-2, display_meta='keyword'),
+                              ])
 
 
 def test_builtin_function_matches_only_at_start(completer, complete_event):
@@ -118,7 +119,9 @@ def test_user_function_name_completion(completer, complete_event):
         Document(text=text, cursor_position=position), complete_event)
     assert set(result) == set([
         Completion(text='custom_func1', start_position=-2, display_meta='function'),
-        Completion(text='custom_func2', start_position=-2, display_meta='function')])
+        Completion(text='custom_func2', start_position=-2, display_meta='function'),
+        Completion(text='CURRENT', start_position=-2, display_meta='keyword'),
+        ])
 
 
 def test_user_function_name_completion_matches_anywhere(completer,
@@ -152,7 +155,9 @@ def test_suggested_column_names_from_visible_table(completer, complete_event):
         Completion(text='last_name', start_position=0, display_meta='column'),
         Completion(text='custom_func1', start_position=0, display_meta='function'),
         Completion(text='custom_func2', start_position=0, display_meta='function')] +
-        list(map(lambda f: Completion(f, display_meta='function'), completer.functions)))
+        list(map(lambda f: Completion(f, display_meta='function'), completer.functions)) +
+        list(map(lambda x: Completion(x, display_meta='keyword'), completer.keywords))
+        )
 
 
 def test_suggested_column_names_in_function(completer, complete_event):
@@ -234,7 +239,9 @@ def test_suggested_multiple_column_names(completer, complete_event):
         Completion(text='last_name', start_position=0, display_meta='column'),
         Completion(text='custom_func1', start_position=0, display_meta='function'),
         Completion(text='custom_func2', start_position=0, display_meta='function')] +
-        list(map(lambda f: Completion(f, display_meta='function'), completer.functions)))
+        list(map(lambda f: Completion(f, display_meta='function'), completer.functions)) +
+        list(map(lambda x: Completion(x, display_meta='keyword'), completer.keywords))
+        )
 
 def test_suggested_multiple_column_names_with_alias(completer, complete_event):
     """
@@ -363,7 +370,9 @@ def test_auto_escaped_col_names(completer, complete_event):
         Completion(text='"ABC"', start_position=0, display_meta='column'),
         Completion(text='custom_func1', start_position=0, display_meta='function'),
         Completion(text='custom_func2', start_position=0, display_meta='function')] +
-        list(map(lambda f: Completion(f, display_meta='function'), completer.functions)))
+        list(map(lambda f: Completion(f, display_meta='function'), completer.functions)) +
+        list(map(lambda x: Completion(x, display_meta='keyword'), completer.keywords))
+        )
 
 
 @pytest.mark.parametrize('text', [
@@ -396,4 +405,4 @@ def test_suggest_columns_from_escaped_table_alias(completer, complete_event):
         Completion(text='id', start_position=0, display_meta='column'),
         Completion(text='"insert"', start_position=0, display_meta='column'),
         Completion(text='"ABC"', start_position=0, display_meta='column'),
-    ])
\ No newline at end of file
+    ])
diff --git a/tests/test_sqlcompletion.py b/tests/test_sqlcompletion.py
index a1a2ab5f..7ca0b009 100644
--- a/tests/test_sqlcompletion.py
+++ b/tests/test_sqlcompletion.py
@@ -9,13 +9,17 @@ def test_select_suggests_cols_with_visible_table_scope():
     suggestions = suggest_type('SELECT  FROM tabl', 'SELECT ')
     assert sorted_dicts(suggestions) == sorted_dicts([
             {'type': 'column', 'tables': [(None, 'tabl', None)]},
-            {'type': 'function', 'schema': []}])
+            {'type': 'function', 'schema': []},
+            {'type': 'keyword'}
+            ])
 
 def test_select_suggests_cols_with_qualified_table_scope():
     suggestions = suggest_type('SELECT  FROM sch.tabl', 'SELECT ')
     assert sorted_dicts(suggestions) == sorted_dicts([
             {'type': 'column', 'tables': [('sch', 'tabl', None)]},
-            {'type': 'function', 'schema': []}])
+            {'type': 'function', 'schema': []},
+            {'type': 'keyword'}
+            ])
 
 
 @pytest.mark.parametrize('expression', [
@@ -34,7 +38,9 @@ def test_where_suggests_columns_functions(expression):
     suggestions = suggest_type(expression, expression)
     assert sorted_dicts(suggestions) == sorted_dicts([
             {'type': 'column', 'tables': [(None, 'tabl', None)]},
-            {'type': 'function', 'schema': []}])
+            {'type': 'function', 'schema': []},
+            {'type': 'keyword'}
+            ])
 
 @pytest.mark.parametrize('expression', [
     'SELECT * FROM tabl WHERE foo IN (',
@@ -44,7 +50,9 @@ def test_where_in_suggests_columns(expression):
     suggestions = suggest_type(expression, expression)
     assert sorted_dicts(suggestions) == sorted_dicts([
             {'type': 'column', 'tables': [(None, 'tabl', None)]},
-            {'type': 'function', 'schema': []}])
+            {'type': 'function', 'schema': []},
+            {'type': 'keyword'}
+            ])
 
 def test_where_equals_any_suggests_columns_or_keywords():
     text = 'SELECT * FROM tabl WHERE foo = ANY('
@@ -63,7 +71,9 @@ def test_select_suggests_cols_and_funcs():
     suggestions = suggest_type('SELECT ', 'SELECT ')
     assert sorted_dicts(suggestions) == sorted_dicts([
          {'type': 'column', 'tables': []},
-         {'type': 'function', 'schema': []}])
+         {'type': 'function', 'schema': []},
+         {'type': 'keyword'}
+         ])
 
 @pytest.mark.parametrize('expression', [
     'SELECT * FROM ',
@@ -113,7 +123,9 @@ def test_col_comma_suggests_cols():
     suggestions = suggest_type('SELECT a, b, FROM tbl', 'SELECT a, b,')
     assert sorted_dicts(suggestions) == sorted_dicts([
         {'type': 'column', 'tables': [(None, 'tbl', None)]},
-        {'type': 'function', 'schema': []}])
+        {'type': 'function', 'schema': []},
+        {'type': 'keyword'}
+        ])
 
 def test_table_comma_suggests_tables_and_schemas():
     suggestions = suggest_type('SELECT a, b FROM tbl1, ',
@@ -147,7 +159,9 @@ def test_partially_typed_col_name_suggests_col_names():
             'SELECT * FROM tabl WHERE col_n')
     assert sorted_dicts(suggestions) == sorted_dicts([
         {'type': 'column', 'tables': [(None, 'tabl', None)]},
-        {'type': 'function', 'schema': []}])
+        {'type': 'function', 'schema': []},
+        {'type': 'keyword'}
+        ])
 
 def test_dot_suggests_cols_of_a_table_or_schema_qualified_table():
     suggestions = suggest_type('SELECT tabl. FROM tabl', 'SELECT tabl.')
@@ -241,7 +255,9 @@ def test_sub_select_col_name_completion():
             'SELECT * FROM (SELECT ')
     assert sorted_dicts(suggestions) == sorted_dicts([
         {'type': 'column', 'tables': [(None, 'abc', None)]},
-        {'type': 'function', 'schema': []}])
+        {'type': 'function', 'schema': []},
+        {'type': 'keyword'}
+        ])
 
 @pytest.mark.xfail
 def test_sub_select_multiple_col_name_completion():
@@ -249,7 +265,9 @@ def test_sub_select_multiple_col_name_completion():
             'SELECT * FROM (SELECT a, ')
     assert sorted_dicts(suggestions) == sorted_dicts([
         {'type': 'column', 'tables': [(None, 'abc', None)]},
-        {'type': 'function', 'schema': []}])
+        {'type': 'function', 'schema': []},
+        {'type': 'keyword'}
+        ])
 
 def test_sub_select_dot_col_name_completion():
     suggestions = suggest_type('SELECT * FROM (SELECT t. FROM tabl t',
@@ -341,7 +359,9 @@ def test_2_statements_2nd_current():
                                'select * from a; select ')
     assert sorted_dicts(suggestions) == sorted_dicts([
             {'type': 'column', 'tables': [(None, 'b', None)]},
-            {'type': 'function', 'schema': []}])
+            {'type': 'function', 'schema': []},
+            {'type': 'keyword'}
+        ])
 
     # Should work even if first statement is invalid
     suggestions = suggest_type('select * from; select * from ',
@@ -363,7 +383,9 @@ def test_2_statements_1st_current():
                                'select ')
     assert sorted_dicts(suggestions) == sorted_dicts([
             {'type': 'column', 'tables': [(None, 'a', None)]},
-            {'type': 'function', 'schema': []}])
+            {'type': 'function', 'schema': []},
+            {'type': 'keyword'}
+            ])
 
 def test_3_statements_2nd_current():
     suggestions = suggest_type('select * from a; select * from ; select * from c',
@@ -377,7 +399,9 @@ def test_3_statements_2nd_current():
                                'select * from a; select ')
     assert sorted_dicts(suggestions) == sorted_dicts([
             {'type': 'column', 'tables': [(None, 'b', None)]},
-            {'type': 'function', 'schema': []}])
+            {'type': 'function', 'schema': []},
+            {'type': 'keyword'}
+            ])
 
 
 def test_create_db_with_template():
@@ -482,4 +506,4 @@ def test_invalid_sql():
     # issue 317
     text = 'selt *'
     suggestions = suggest_type(text, text)
-    assert suggestions == [{'type': 'keyword'}]
\ No newline at end of file
+    assert suggestions == [{'type': 'keyword'}]