diff options
Diffstat (limited to 'tests/test_smart_completion.py')
| -rw-r--r-- | tests/test_smart_completion.py | 85 |
1 files changed, 74 insertions, 11 deletions
diff --git a/tests/test_smart_completion.py b/tests/test_smart_completion.py index 5344d13d..8b938894 100644 --- a/tests/test_smart_completion.py +++ b/tests/test_smart_completion.py @@ -1,20 +1,45 @@ import pytest +from pandas import DataFrame from prompt_toolkit.completion import Completion from prompt_toolkit.document import Document -tables = { - 'users': ['id', 'email', 'first_name', 'last_name'], - 'orders': ['id', 'user_id', 'ordered_date', 'status'] -} +schemata = { + 'public': { + 'users': ['id', 'email', 'first_name', 'last_name'], + 'orders': ['id', 'ordered_date', 'status'] + }, + 'custom': { + '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) - comp.extend_table_names(tables.keys()) - for t in tables: - comp.extend_column_names(t, tables[t]) + + # Table metadata is a dataframe with columns [schema, table, is_visible] + tables = DataFrame.from_records( + ((schema, table, schema=='public') + for schema, tables in schemata.items() + for table, columns in tables.items()), + columns=['schema', 'table', 'is_visible']) + + # Column metadata is a dataframe with columns [schema, table, column] + columns = DataFrame.from_records( + ((schema, table, column) + for schema, tables in schemata.items() + for table, columns in tables.items() + for column in columns), + columns=['schema', 'table', 'column']) + + + comp.extend_tables(tables) + comp.extend_columns(columns) + return comp @pytest.fixture @@ -39,15 +64,26 @@ def test_select_keyword_completion(completer, complete_event): complete_event) assert set(result) == set([Completion(text='SELECT', start_position=-3)]) + +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='orders', start_position=0)]) + + def test_function_name_completion(completer, complete_event): text = 'SELECT MA' position = len('SELECT MA') result = completer.get_completions( - Document(text=text, cursor_position=position), - complete_event) + Document(text=text, cursor_position=position), complete_event) assert set(result) == set([Completion(text='MAX', start_position=-2)]) -def test_suggested_column_names(completer, complete_event): +def test_suggested_column_names_from_visible_table(completer, complete_event): """ Suggest column and function names when selecting from table :param completer: @@ -67,6 +103,24 @@ def test_suggested_column_names(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 @@ -87,7 +141,7 @@ def test_suggested_column_names_in_function(completer, complete_event): Completion(text='first_name', start_position=0), Completion(text='last_name', start_position=0)]) -def test_suggested_column_names_with_dot(completer, complete_event): +def test_suggested_column_names_with_table_dot(completer, complete_event): """ Suggest column names on table name and dot :param completer: @@ -106,6 +160,15 @@ def test_suggested_column_names_with_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 |