Cleanup sql_completion tests.

2 files changed, 31 insertions, 16 deletions

diff --git a/pgcli/packages/sqlcompletion.py b/pgcli/packages/sqlcompletion.py
index 3d46e85c..81195e67 100644
--- a/pgcli/packages/sqlcompletion.py
+++ b/pgcli/packages/sqlcompletion.py
@@ -1,6 +1,8 @@
+from __future__ import print_function
 import sqlparse
 from parseutils import last_word, extract_tables
 
+
 def suggest_type(full_text, text_before_cursor):
     """Takes the full_text that is typed so far and also the text before the
     cursor to suggest completion type and scope.
diff --git a/tests/test_sqlcompletion.py b/tests/test_sqlcompletion.py
index d31042f0..ff8e8e45 100644
--- a/tests/test_sqlcompletion.py
+++ b/tests/test_sqlcompletion.py
@@ -1,18 +1,31 @@
 from pgcli.packages.sqlcompletion import suggest_type
 
-class TestSqlCompletion:
-    def test_lparen_suggest_cols(self):
-        suggestion = suggest_type('SELECT MAX( FROM table_name', 'SELECT MAX(')
-        assert (suggestion and suggestion[0] == 'columns')
-
-    def test_select_suggest_cols_and_funcs(self):
-        suggestion = suggest_type('SELECT ', 'SELECT ')
-        assert (suggestion and suggestion[0] == 'columns-and-functions')
-
-    def test_from_suggest_tables(self):
-        suggestion = suggest_type('SELECT * FROM ', 'SELECT * FROM ')
-        assert (suggestion and suggestion[0] == 'tables')
-
-    def test_distinct_suggest_cols(self):
-        suggestion = suggest_type('SELECT DISTINCT ', 'SELECT DISTINCT ')
-        assert (suggestion and suggestion[0] == 'columns')
+
+def test_select_suggests_cols_with_table_scope():
+    suggestion = suggest_type('SELECT  FROM tabl', 'SELECT ')
+    assert suggestion == ('columns-and-functions', ['tabl'])
+
+def test_lparen_suggest_cols():
+    suggestion = suggest_type('SELECT MAX( FROM tbl', 'SELECT MAX(')
+    assert suggestion == ('columns', ['tbl'])
+
+def test_select_suggest_cols_and_funcs():
+    suggestion = suggest_type('SELECT ', 'SELECT ')
+    assert suggestion == ('columns-and-functions', [])
+
+def test_from_suggest_tables():
+    suggestion = suggest_type('SELECT * FROM ', 'SELECT * FROM ')
+    assert suggestion == ('tables', [])
+
+def test_distinct_suggest_cols():
+    suggestion = suggest_type('SELECT DISTINCT ', 'SELECT DISTINCT ')
+    assert suggestion == ('columns', [])
+
+def test_multiple_cols_suggest_cols():
+    suggestion = suggest_type('SELECT a, b, FROM tbl', 'SELECT a, b,')
+    assert suggestion == ('columns-and-functions', ['tbl'])
+
+def test_multiple_tables_suggest_tables():
+    suggestion = suggest_type('SELECT a, b FROM tbl1, ',
+            'SELECT a, b FROM tbl1, ')
+    assert suggestion == ('tables', [])