1 files changed, 555 insertions, 0 deletions
diff --git a/test/test_completion_engine.py b/test/test_completion_engine.py
new file mode 100644
index 0000000..318b632
--- /dev/null
+++ b/test/test_completion_engine.py
@@ -0,0 +1,555 @@
+from mycli.packages.completion_engine import suggest_type
+import pytest
+
+
+def sorted_dicts(dicts):
+    """input is a list of dicts."""
+    return sorted(tuple(x.items()) for x in dicts)
+
+
+def test_select_suggests_cols_with_visible_table_scope():
+    suggestions = suggest_type('SELECT  FROM tabl', 'SELECT ')
+    assert sorted_dicts(suggestions) == sorted_dicts([
+        {'type': 'alias', 'aliases': ['tabl']},
+        {'type': 'column', 'tables': [(None, 'tabl', None)]},
+        {'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': 'alias', 'aliases': ['tabl']},
+        {'type': 'column', 'tables': [('sch', 'tabl', None)]},
+        {'type': 'function', 'schema': []},
+        {'type': 'keyword'},
+    ])
+
+
+@pytest.mark.parametrize('expression', [
+    'SELECT * FROM tabl WHERE ',
+    'SELECT * FROM tabl WHERE (',
+    'SELECT * FROM tabl WHERE foo = ',
+    'SELECT * FROM tabl WHERE bar OR ',
+    'SELECT * FROM tabl WHERE foo = 1 AND ',
+    'SELECT * FROM tabl WHERE (bar > 10 AND ',
+    'SELECT * FROM tabl WHERE (bar AND (baz OR (qux AND (',
+    'SELECT * FROM tabl WHERE 10 < ',
+    'SELECT * FROM tabl WHERE foo BETWEEN ',
+    'SELECT * FROM tabl WHERE foo BETWEEN foo AND ',
+])
+def test_where_suggests_columns_functions(expression):
+    suggestions = suggest_type(expression, expression)
+    assert sorted_dicts(suggestions) == sorted_dicts([
+        {'type': 'alias', 'aliases': ['tabl']},
+        {'type': 'column', 'tables': [(None, 'tabl', None)]},
+        {'type': 'function', 'schema': []},
+        {'type': 'keyword'},
+    ])
+
+
+@pytest.mark.parametrize('expression', [
+    'SELECT * FROM tabl WHERE foo IN (',
+    'SELECT * FROM tabl WHERE foo IN (bar, ',
+])
+def test_where_in_suggests_columns(expression):
+    suggestions = suggest_type(expression, expression)
+    assert sorted_dicts(suggestions) == sorted_dicts([
+        {'type': 'alias', 'aliases': ['tabl']},
+        {'type': 'column', 'tables': [(None, 'tabl', None)]},
+        {'type': 'function', 'schema': []},
+        {'type': 'keyword'},
+    ])
+
+
+def test_where_equals_any_suggests_columns_or_keywords():
+    text = 'SELECT * FROM tabl WHERE foo = ANY('
+    suggestions = suggest_type(text, text)
+    assert sorted_dicts(suggestions) == sorted_dicts([
+        {'type': 'alias', 'aliases': ['tabl']},
+        {'type': 'column', 'tables': [(None, 'tabl', None)]},
+        {'type': 'function', 'schema': []},
+        {'type': 'keyword'}])
+
+
+def test_lparen_suggests_cols():
+    suggestion = suggest_type('SELECT MAX( FROM tbl', 'SELECT MAX(')
+    assert suggestion == [
+        {'type': 'column', 'tables': [(None, 'tbl', None)]}]
+
+
+def test_operand_inside_function_suggests_cols1():
+    suggestion = suggest_type(
+        'SELECT MAX(col1 +  FROM tbl', 'SELECT MAX(col1 + ')
+    assert suggestion == [
+        {'type': 'column', 'tables': [(None, 'tbl', None)]}]
+
+
+def test_operand_inside_function_suggests_cols2():
+    suggestion = suggest_type(
+        'SELECT MAX(col1 + col2 +  FROM tbl', 'SELECT MAX(col1 + col2 + ')
+    assert suggestion == [
+        {'type': 'column', 'tables': [(None, 'tbl', None)]}]
+
+
+def test_select_suggests_cols_and_funcs():
+    suggestions = suggest_type('SELECT ', 'SELECT ')
+    assert sorted_dicts(suggestions) == sorted_dicts([
+        {'type': 'alias', 'aliases': []},
+        {'type': 'column', 'tables': []},
+        {'type': 'function', 'schema': []},
+        {'type': 'keyword'},
+    ])
+
+
+@pytest.mark.parametrize('expression', [
+    'SELECT * FROM ',
+    'INSERT INTO ',
+    'COPY ',
+    'UPDATE ',
+    'DESCRIBE ',
+    'DESC ',
+    'EXPLAIN ',
+    'SELECT * FROM foo JOIN ',
+])
+def test_expression_suggests_tables_views_and_schemas(expression):
+    suggestions = suggest_type(expression, expression)
+    assert sorted_dicts(suggestions) == sorted_dicts([
+        {'type': 'table', 'schema': []},
+        {'type': 'view', 'schema': []},
+        {'type': 'schema'}])
+
+
+@pytest.mark.parametrize('expression', [
+    'SELECT * FROM sch.',
+    'INSERT INTO sch.',
+    'COPY sch.',
+    'UPDATE sch.',
+    'DESCRIBE sch.',
+    'DESC sch.',
+    'EXPLAIN sch.',
+    'SELECT * FROM foo JOIN sch.',
+])
+def test_expression_suggests_qualified_tables_views_and_schemas(expression):
+    suggestions = suggest_type(expression, expression)
+    assert sorted_dicts(suggestions) == sorted_dicts([
+        {'type': 'table', 'schema': 'sch'},
+        {'type': 'view', 'schema': 'sch'}])
+
+
+def test_truncate_suggests_tables_and_schemas():
+    suggestions = suggest_type('TRUNCATE ', 'TRUNCATE ')
+    assert sorted_dicts(suggestions) == sorted_dicts([
+        {'type': 'table', 'schema': []},
+        {'type': 'schema'}])
+
+
+def test_truncate_suggests_qualified_tables():
+    suggestions = suggest_type('TRUNCATE sch.', 'TRUNCATE sch.')
+    assert sorted_dicts(suggestions) == sorted_dicts([
+        {'type': 'table', 'schema': 'sch'}])
+
+
+def test_distinct_suggests_cols():
+    suggestions = suggest_type('SELECT DISTINCT ', 'SELECT DISTINCT ')
+    assert suggestions == [{'type': 'column', 'tables': []}]
+
+
+def test_col_comma_suggests_cols():
+    suggestions = suggest_type('SELECT a, b, FROM tbl', 'SELECT a, b,')
+    assert sorted_dicts(suggestions) == sorted_dicts([
+        {'type': 'alias', 'aliases': ['tbl']},
+        {'type': 'column', 'tables': [(None, 'tbl', None)]},
+        {'type': 'function', 'schema': []},
+        {'type': 'keyword'},
+    ])
+
+
+def test_table_comma_suggests_tables_and_schemas():
+    suggestions = suggest_type('SELECT a, b FROM tbl1, ',
+                               'SELECT a, b FROM tbl1, ')
+    assert sorted_dicts(suggestions) == sorted_dicts([
+        {'type': 'table', 'schema': []},
+        {'type': 'view', 'schema': []},
+        {'type': 'schema'}])
+
+
+def test_into_suggests_tables_and_schemas():
+    suggestion = suggest_type('INSERT INTO ', 'INSERT INTO ')
+    assert sorted_dicts(suggestion) == sorted_dicts([
+        {'type': 'table', 'schema': []},
+        {'type': 'view', 'schema': []},
+        {'type': 'schema'}])
+
+
+def test_insert_into_lparen_suggests_cols():
+    suggestions = suggest_type('INSERT INTO abc (', 'INSERT INTO abc (')
+    assert suggestions == [{'type': 'column', 'tables': [(None, 'abc', None)]}]
+
+
+def test_insert_into_lparen_partial_text_suggests_cols():
+    suggestions = suggest_type('INSERT INTO abc (i', 'INSERT INTO abc (i')
+    assert suggestions == [{'type': 'column', 'tables': [(None, 'abc', None)]}]
+
+
+def test_insert_into_lparen_comma_suggests_cols():
+    suggestions = suggest_type('INSERT INTO abc (id,', 'INSERT INTO abc (id,')
+    assert suggestions == [{'type': 'column', 'tables': [(None, 'abc', None)]}]
+
+
+def test_partially_typed_col_name_suggests_col_names():
+    suggestions = suggest_type('SELECT * FROM tabl WHERE col_n',
+                               'SELECT * FROM tabl WHERE col_n')
+    assert sorted_dicts(suggestions) == sorted_dicts([
+        {'type': 'alias', 'aliases': ['tabl']},
+        {'type': 'column', 'tables': [(None, 'tabl', None)]},
+        {'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.')
+    assert sorted_dicts(suggestions) == sorted_dicts([
+        {'type': 'column', 'tables': [(None, 'tabl', None)]},
+        {'type': 'table', 'schema': 'tabl'},
+        {'type': 'view', 'schema': 'tabl'},
+        {'type': 'function', 'schema': 'tabl'}])
+
+
+def test_dot_suggests_cols_of_an_alias():
+    suggestions = suggest_type('SELECT t1. FROM tabl1 t1, tabl2 t2',
+                               'SELECT t1.')
+    assert sorted_dicts(suggestions) == sorted_dicts([
+        {'type': 'table', 'schema': 't1'},
+        {'type': 'view', 'schema': 't1'},
+        {'type': 'column', 'tables': [(None, 'tabl1', 't1')]},
+        {'type': 'function', 'schema': 't1'}])
+
+
+def test_dot_col_comma_suggests_cols_or_schema_qualified_table():
+    suggestions = suggest_type('SELECT t1.a, t2. FROM tabl1 t1, tabl2 t2',
+                               'SELECT t1.a, t2.')
+    assert sorted_dicts(suggestions) == sorted_dicts([
+        {'type': 'column', 'tables': [(None, 'tabl2', 't2')]},
+        {'type': 'table', 'schema': 't2'},
+        {'type': 'view', 'schema': 't2'},
+        {'type': 'function', 'schema': 't2'}])
+
+
+@pytest.mark.parametrize('expression', [
+    'SELECT * FROM (',
+    'SELECT * FROM foo WHERE EXISTS (',
+    'SELECT * FROM foo WHERE bar AND NOT EXISTS (',
+    'SELECT 1 AS',
+])
+def test_sub_select_suggests_keyword(expression):
+    suggestion = suggest_type(expression, expression)
+    assert suggestion == [{'type': 'keyword'}]
+
+
+@pytest.mark.parametrize('expression', [
+    'SELECT * FROM (S',
+    'SELECT * FROM foo WHERE EXISTS (S',
+    'SELECT * FROM foo WHERE bar AND NOT EXISTS (S',
+])
+def test_sub_select_partial_text_suggests_keyword(expression):
+    suggestion = suggest_type(expression, expression)
+    assert suggestion == [{'type': 'keyword'}]
+
+
+def test_outer_table_reference_in_exists_subquery_suggests_columns():
+    q = 'SELECT * FROM foo f WHERE EXISTS (SELECT 1 FROM bar WHERE f.'
+    suggestions = suggest_type(q, q)
+    assert suggestions == [
+        {'type': 'column', 'tables': [(None, 'foo', 'f')]},
+        {'type': 'table', 'schema': 'f'},
+        {'type': 'view', 'schema': 'f'},
+        {'type': 'function', 'schema': 'f'}]
+
+
+@pytest.mark.parametrize('expression', [
+    'SELECT * FROM (SELECT * FROM ',
+    'SELECT * FROM foo WHERE EXISTS (SELECT * FROM ',
+    'SELECT * FROM foo WHERE bar AND NOT EXISTS (SELECT * FROM ',
+])
+def test_sub_select_table_name_completion(expression):
+    suggestion = suggest_type(expression, expression)
+    assert sorted_dicts(suggestion) == sorted_dicts([
+        {'type': 'table', 'schema': []},
+        {'type': 'view', 'schema': []},
+        {'type': 'schema'}])
+
+
+def test_sub_select_col_name_completion():
+    suggestions = suggest_type('SELECT * FROM (SELECT  FROM abc',
+                               'SELECT * FROM (SELECT ')
+    assert sorted_dicts(suggestions) == sorted_dicts([
+        {'type': 'alias', 'aliases': ['abc']},
+        {'type': 'column', 'tables': [(None, 'abc', None)]},
+        {'type': 'function', 'schema': []},
+        {'type': 'keyword'},
+    ])
+
+
+@pytest.mark.xfail
+def test_sub_select_multiple_col_name_completion():
+    suggestions = suggest_type('SELECT * FROM (SELECT a, FROM abc',
+                               'SELECT * FROM (SELECT a, ')
+    assert sorted_dicts(suggestions) == sorted_dicts([
+        {'type': 'column', 'tables': [(None, 'abc', None)]},
+        {'type': 'function', 'schema': []}])
+
+
+def test_sub_select_dot_col_name_completion():
+    suggestions = suggest_type('SELECT * FROM (SELECT t. FROM tabl t',
+                               'SELECT * FROM (SELECT t.')
+    assert sorted_dicts(suggestions) == sorted_dicts([
+        {'type': 'column', 'tables': [(None, 'tabl', 't')]},
+        {'type': 'table', 'schema': 't'},
+        {'type': 'view', 'schema': 't'},
+        {'type': 'function', 'schema': 't'}])
+
+
+@pytest.mark.parametrize('join_type', ['', 'INNER', 'LEFT', 'RIGHT OUTER'])
+@pytest.mark.parametrize('tbl_alias', ['', 'foo'])
+def test_join_suggests_tables_and_schemas(tbl_alias, join_type):
+    text = 'SELECT * FROM abc {0} {1} JOIN '.format(tbl_alias, join_type)
+    suggestion = suggest_type(text, text)
+    assert sorted_dicts(suggestion) == sorted_dicts([
+        {'type': 'table', 'schema': []},
+        {'type': 'view', 'schema': []},
+        {'type': 'schema'}])
+
+
+@pytest.mark.parametrize('sql', [
+    'SELECT * FROM abc a JOIN def d ON a.',
+    'SELECT * FROM abc a JOIN def d ON a.id = d.id AND a.',
+])
+def test_join_alias_dot_suggests_cols1(sql):
+    suggestions = suggest_type(sql, sql)
+    assert sorted_dicts(suggestions) == sorted_dicts([
+        {'type': 'column', 'tables': [(None, 'abc', 'a')]},
+        {'type': 'table', 'schema': 'a'},
+        {'type': 'view', 'schema': 'a'},
+        {'type': 'function', 'schema': 'a'}])
+
+
+@pytest.mark.parametrize('sql', [
+    'SELECT * FROM abc a JOIN def d ON a.id = d.',
+    'SELECT * FROM abc a JOIN def d ON a.id = d.id AND a.id2 = d.',
+])
+def test_join_alias_dot_suggests_cols2(sql):
+    suggestions = suggest_type(sql, sql)
+    assert sorted_dicts(suggestions) == sorted_dicts([
+        {'type': 'column', 'tables': [(None, 'def', 'd')]},
+        {'type': 'table', 'schema': 'd'},
+        {'type': 'view', 'schema': 'd'},
+        {'type': 'function', 'schema': 'd'}])
+
+
+@pytest.mark.parametrize('sql', [
+    'select a.x, b.y from abc a join bcd b on ',
+    'select a.x, b.y from abc a join bcd b on a.id = b.id OR ',
+])
+def test_on_suggests_aliases(sql):
+    suggestions = suggest_type(sql, sql)
+    assert suggestions == [{'type': 'alias', 'aliases': ['a', 'b']}]
+
+
+@pytest.mark.parametrize('sql', [
+    'select abc.x, bcd.y from abc join bcd on ',
+    'select abc.x, bcd.y from abc join bcd on abc.id = bcd.id AND ',
+])
+def test_on_suggests_tables(sql):
+    suggestions = suggest_type(sql, sql)
+    assert suggestions == [{'type': 'alias', 'aliases': ['abc', 'bcd']}]
+
+
+@pytest.mark.parametrize('sql', [
+    'select a.x, b.y from abc a join bcd b on a.id = ',
+    'select a.x, b.y from abc a join bcd b on a.id = b.id AND a.id2 = ',
+])
+def test_on_suggests_aliases_right_side(sql):
+    suggestions = suggest_type(sql, sql)
+    assert suggestions == [{'type': 'alias', 'aliases': ['a', 'b']}]
+
+
+@pytest.mark.parametrize('sql', [
+    'select abc.x, bcd.y from abc join bcd on ',
+    'select abc.x, bcd.y from abc join bcd on abc.id = bcd.id and ',
+])
+def test_on_suggests_tables_right_side(sql):
+    suggestions = suggest_type(sql, sql)
+    assert suggestions == [{'type': 'alias', 'aliases': ['abc', 'bcd']}]
+
+
+@pytest.mark.parametrize('col_list', ['', 'col1, '])
+def test_join_using_suggests_common_columns(col_list):
+    text = 'select * from abc inner join def using (' + col_list
+    assert suggest_type(text, text) == [
+        {'type': 'column',
+         'tables': [(None, 'abc', None), (None, 'def', None)],
+         'drop_unique': True}]
+
+@pytest.mark.parametrize('sql', [
+    'SELECT * FROM abc a JOIN def d ON a.id = d.id JOIN ghi g ON g.',
+    'SELECT * FROM abc a JOIN def d ON a.id = d.id AND a.id2 = d.id2 JOIN ghi g ON d.id = g.id AND g.',
+])
+def test_two_join_alias_dot_suggests_cols1(sql):
+    suggestions = suggest_type(sql, sql)
+    assert sorted_dicts(suggestions) == sorted_dicts([
+        {'type': 'column', 'tables': [(None, 'ghi', 'g')]},
+        {'type': 'table', 'schema': 'g'},
+        {'type': 'view', 'schema': 'g'},
+        {'type': 'function', 'schema': 'g'}])
+
+def test_2_statements_2nd_current():
+    suggestions = suggest_type('select * from a; select * from ',
+                               'select * from a; select * from ')
+    assert sorted_dicts(suggestions) == sorted_dicts([
+        {'type': 'table', 'schema': []},
+        {'type': 'view', 'schema': []},
+        {'type': 'schema'}])
+
+    suggestions = suggest_type('select * from a; select  from b',
+                               'select * from a; select ')
+    assert sorted_dicts(suggestions) == sorted_dicts([
+        {'type': 'alias', 'aliases': ['b']},
+        {'type': 'column', 'tables': [(None, 'b', None)]},
+        {'type': 'function', 'schema': []},
+        {'type': 'keyword'},
+    ])
+
+    # Should work even if first statement is invalid
+    suggestions = suggest_type('select * from; select * from ',
+                               'select * from; select * from ')
+    assert sorted_dicts(suggestions) == sorted_dicts([
+        {'type': 'table', 'schema': []},
+        {'type': 'view', 'schema': []},
+        {'type': 'schema'}])
+
+
+def test_2_statements_1st_current():
+    suggestions = suggest_type('select * from ; select * from b',
+                               'select * from ')
+    assert sorted_dicts(suggestions) == sorted_dicts([
+        {'type': 'table', 'schema': []},
+        {'type': 'view', 'schema': []},
+        {'type': 'schema'}])
+
+    suggestions = suggest_type('select  from a; select * from b',
+                               'select ')
+    assert sorted_dicts(suggestions) == sorted_dicts([
+        {'type': 'alias', 'aliases': ['a']},
+        {'type': 'column', 'tables': [(None, 'a', None)]},
+        {'type': 'function', 'schema': []},
+        {'type': 'keyword'},
+    ])
+
+
+def test_3_statements_2nd_current():
+    suggestions = suggest_type('select * from a; select * from ; select * from c',
+                               'select * from a; select * from ')
+    assert sorted_dicts(suggestions) == sorted_dicts([
+        {'type': 'table', 'schema': []},
+        {'type': 'view', 'schema': []},
+        {'type': 'schema'}])
+
+    suggestions = suggest_type('select * from a; select  from b; select * from c',
+                               'select * from a; select ')
+    assert sorted_dicts(suggestions) == sorted_dicts([
+        {'type': 'alias', 'aliases': ['b']},
+        {'type': 'column', 'tables': [(None, 'b', None)]},
+        {'type': 'function', 'schema': []},
+        {'type': 'keyword'},
+    ])
+
+
+def test_create_db_with_template():
+    suggestions = suggest_type('create database foo with template ',
+                               'create database foo with template ')
+
+    assert sorted_dicts(suggestions) == sorted_dicts([{'type': 'database'}])
+
+
+@pytest.mark.parametrize('initial_text', ['', '    ', '\t \t'])
+def test_specials_included_for_initial_completion(initial_text):
+    suggestions = suggest_type(initial_text, initial_text)
+
+    assert sorted_dicts(suggestions) == \
+        sorted_dicts([{'type': 'keyword'}, {'type': 'special'}])
+
+
+def test_specials_not_included_after_initial_token():
+    suggestions = suggest_type('create table foo (dt d',
+                               'create table foo (dt d')
+
+    assert sorted_dicts(suggestions) == sorted_dicts([{'type': 'keyword'}])
+
+
+def test_drop_schema_qualified_table_suggests_only_tables():
+    text = 'DROP TABLE schema_name.table_name'
+    suggestions = suggest_type(text, text)
+    assert suggestions == [{'type': 'table', 'schema': 'schema_name'}]
+
+
+@pytest.mark.parametrize('text', [',', '  ,', 'sel ,'])
+def test_handle_pre_completion_comma_gracefully(text):
+    suggestions = suggest_type(text, text)
+
+    assert iter(suggestions)
+
+
+def test_cross_join():
+    text = 'select * from v1 cross join v2 JOIN v1.id, '
+    suggestions = suggest_type(text, text)
+    assert sorted_dicts(suggestions) == sorted_dicts([
+        {'type': 'table', 'schema': []},
+        {'type': 'view', 'schema': []},
+        {'type': 'schema'}])
+
+
+@pytest.mark.parametrize('expression', [
+    'SELECT 1 AS ',
+    'SELECT 1 FROM tabl AS ',
+])
+def test_after_as(expression):
+    suggestions = suggest_type(expression, expression)
+    assert set(suggestions) == set()
+
+
+@pytest.mark.parametrize('expression', [
+    '\\. ',
+    'select 1; \\. ',
+    'select 1;\\. ',
+    'select 1 ; \\. ',
+    'source ',
+    'truncate table test; source ',
+    'truncate table test ; source ',
+    'truncate table test;source ',
+])
+def test_source_is_file(expression):
+    suggestions = suggest_type(expression, expression)
+    assert suggestions == [{'type': 'file_name'}]
+
+
+@pytest.mark.parametrize("expression", [
+    "\\f ",
+])
+def test_favorite_name_suggestion(expression):
+    suggestions = suggest_type(expression, expression)
+    assert suggestions == [{'type': 'favoritequery'}]
+
+
+def test_order_by():
+    text = 'select * from foo order by '
+    suggestions = suggest_type(text, text)
+    assert suggestions == [{'tables': [(None, 'foo', None)], 'type': 'column'}]
+
+
+def test_quoted_where():
+    text = "'where i=';"
+    suggestions = suggest_type(text, text)
+    assert suggestions == [{'type': 'keyword'}]