Fix pgexecute tests

1 files changed, 12 insertions, 7 deletions

diff --git a/tests/utils.py b/tests/utils.py
index 93e9f258..d7d5993d 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -58,12 +58,17 @@ def drop_tables(conn):
             DROP SCHEMA IF EXISTS schema2 CASCADE''')
 
 
-def run(executor, sql, join=False, expanded=False, pgspecial=None):
+def run(executor, sql, join=False, expanded=False, pgspecial=None,
+        exception_formatter=None):
     " Return string output for the sql to be run "
-    result = []
-    for title, rows, headers, status in executor.run(sql, pgspecial):
-        result.extend(format_output(title, rows, headers, status, 'psql',
-                                    expanded=expanded))
+
+    results = executor.run(sql, pgspecial, exception_formatter)
+    formatted = []
+
+    for title, rows, headers, status in results:
+        formatted.extend(format_output(title, rows, headers, status, 'psql',
+                                       expanded=expanded))
     if join:
-        result = '\n'.join(result)
-    return result
+        formatted = '\n'.join(formatted)
+
+    return formatted