Merge pull request #527 from dbcli/koljonen/suggest_joins

For 'JOIN <cursor>', suggest 'foo on foo.fooid = bar.fooid'

1 files changed, 49 insertions, 1 deletions

diff --git a/tests/test_smart_completion_multiple_schemata.py b/tests/test_smart_completion_multiple_schemata.py
index c107a710..9f74a7d4 100644
--- a/tests/test_smart_completion_multiple_schemata.py
+++ b/tests/test_smart_completion_multiple_schemata.py
@@ -1,8 +1,9 @@
 from __future__ import unicode_literals
 import pytest
+import itertools
 from prompt_toolkit.completion import Completion
 from prompt_toolkit.document import Document
-from pgcli.packages.function_metadata import FunctionMetadata
+from pgcli.packages.function_metadata import FunctionMetadata, ForeignKey
 
 metadata = {
             'tables': {
@@ -37,6 +38,9 @@ metadata = {
                             'public':   ['typ1', 'typ2'],
                             'custom':   ['typ3', 'typ4'],
                          },
+                'foreignkeys': [
+                    ('public', 'users', 'id', 'custom', 'shipments', 'user_id')
+                ],
             }
 
 @pytest.fixture
@@ -63,11 +67,14 @@ def completer():
                     for schema, datatypes in metadata['datatypes'].items()
                     for datatype in datatypes]
 
+    foreignkeys = [ForeignKey(*fk) for fk in metadata['foreignkeys']]
+
     comp.extend_schemata(schemata)
     comp.extend_relations(tables, kind='tables')
     comp.extend_columns(columns, kind='tables')
     comp.extend_functions(functions)
     comp.extend_datatypes(datatypes)
+    comp.extend_foreignkeys(foreignkeys)
     comp.set_search_path(['public'])
 
     return comp
@@ -135,6 +142,47 @@ def test_suggested_column_names_from_qualified_shadowed_table(completer, complet
         list(map(lambda x: Completion(x, display_meta='keyword'), completer.keywords))
         )
 
+@pytest.mark.parametrize('text', [
+    'SELECT * FROM users JOIN custom.shipments ON ',
+    '''SELECT *
+    FROM public.users
+    JOIN custom.shipments ON '''
+])
+def test_suggested_join_conditions(completer, complete_event, text):
+    position = len(text)
+    result = set(completer.get_completions(
+        Document(text=text, cursor_position=position),
+        complete_event))
+    assert set(result) == set([
+        Completion(text='users', start_position=0, display_meta='table alias'),
+        Completion(text='shipments', start_position=0, display_meta='table alias'),
+        Completion(text='shipments.id = users.id', start_position=0, display_meta='name join'),
+        Completion(text='shipments.user_id = users.id', start_position=0, display_meta='fk join')])
+
+@pytest.mark.parametrize(('query', 'tbl'), itertools.product((
+    'SELECT * FROM public.{0} RIGHT OUTER JOIN ',
+    '''SELECT *
+    FROM {0}
+    JOIN '''
+), ('users', '"users"', 'Users')))
+def test_suggested_joins(completer, complete_event, query, tbl):
+    text = query.format(tbl)
+    position = len(text)
+    result = set(completer.get_completions(
+        Document(text=text, cursor_position=position),
+        complete_event))
+    join = 'custom.shipments ON shipments.user_id = {0}.id'.format(tbl)
+    assert set(result) == set([
+        Completion(text=join, start_position=0, display_meta='join'),
+        Completion(text='public', start_position=0, display_meta='schema'),
+        Completion(text='custom', start_position=0, display_meta='schema'),
+        Completion(text='"Custom"', start_position=0, display_meta='schema'),
+        Completion(text='orders', start_position=0, display_meta='table'),
+        Completion(text='users', start_position=0, display_meta='table'),
+        Completion(text='"select"', start_position=0, display_meta='table'),
+        Completion(text='func1', start_position=0, display_meta='function'),
+        Completion(text='func2', start_position=0, display_meta='function')])
+
 def test_suggested_column_names_from_schema_qualifed_table(completer, complete_event):
     """
     Suggest column and function names when selecting from a qualified-table