pgexecute returns additional function metadata

3 files changed, 23 insertions, 3 deletions

diff --git a/tests/test_pgexecute.py b/tests/test_pgexecute.py
index 3f6d4e86..dc289220 100644
--- a/tests/test_pgexecute.py
+++ b/tests/test_pgexecute.py
@@ -4,6 +4,7 @@ import pytest
 from pgcli.packages.pgspecial import PGSpecial
 from textwrap import dedent
 from utils import run, dbtest, requires_json, requires_jsonb
+from pgcli.pgexecute import FunctionMetadata
 
 @dbtest
 def test_conn(executor):
@@ -68,8 +69,24 @@ def test_functions_query(executor):
     run(executor, '''create function schema1.func2() returns int
                      language sql as $$select 2$$''')
 
+    run(executor, '''create function func3()
+                     returns table(x int, y int) language sql
+                     as $$select 1, 2 from generate_series(1,5)$$;''')
+
+    run(executor, '''create function func4(x int) returns setof int language sql
+                     as $$select generate_series(1,5)$$;''')
+
     funcs = set(executor.functions())
-    assert funcs >= set([('public', 'func1'), ('schema1', 'func2')])
+    assert funcs >= set([
+        FunctionMetadata('public', 'func1', '',
+                         'integer', False, False, False),
+        FunctionMetadata('public', 'func3', '',
+                         'TABLE(x integer, y integer)', False, False, True),
+        FunctionMetadata('public', 'func4', 'x integer',
+                         'SETOF integer', False, False, True),
+        FunctionMetadata('schema1', 'func2', '',
+                         'integer', False, False, False),
+      ])
 
 
 @dbtest
diff --git a/tests/test_smart_completion_multiple_schemata.py b/tests/test_smart_completion_multiple_schemata.py
index fd8854d6..3a80a42f 100644
--- a/tests/test_smart_completion_multiple_schemata.py
+++ b/tests/test_smart_completion_multiple_schemata.py
@@ -2,6 +2,7 @@ from __future__ import unicode_literals
 import pytest
 from prompt_toolkit.completion import Completion
 from prompt_toolkit.document import Document
+from pgcli.pgexecute import FunctionMetadata
 
 metadata = {
             'tables': {
@@ -40,7 +41,7 @@ def completer():
             tables.append((schema, table))
             columns.extend([(schema, table, col) for col in cols])
 
-    functions = [(schema, func)
+    functions = [FunctionMetadata(schema, func, '', '', False, False, False)
                     for schema, funcs in metadata['functions'].items()
                     for func in funcs]
 
diff --git a/tests/test_smart_completion_public_schema_only.py b/tests/test_smart_completion_public_schema_only.py
index 2f75069d..e1e2ac17 100644
--- a/tests/test_smart_completion_public_schema_only.py
+++ b/tests/test_smart_completion_public_schema_only.py
@@ -2,6 +2,7 @@ from __future__ import unicode_literals
 import pytest
 from prompt_toolkit.completion import Completion
 from prompt_toolkit.document import Document
+from pgcli.pgexecute import FunctionMetadata
 
 metadata = {
                 'tables': {
@@ -42,7 +43,8 @@ def completer():
     comp.extend_columns(columns, kind='views')
 
     # functions
-    functions = [('public', func) for func in metadata['functions']]
+    functions = [FunctionMetadata('public', func, '', '', False, False, False)
+                 for func in metadata['functions']]
     comp.extend_functions(functions)
 
     # types