Fix formatting of output from select steps

2 files changed, 28 insertions, 4 deletions

diff --git a/yoyo/migrations.py b/yoyo/migrations.py
index 1ae942c..bb1fddc 100755
--- a/yoyo/migrations.py
+++ b/yoyo/migrations.py
@@ -181,11 +181,13 @@ class MigrationStep(StepBase):
         self._rollback = rollback
         self._apply = apply
 
-    def _execute(self, cursor, stmt, out=stdout):
+    def _execute(self, cursor, stmt, out=None):
         """
         Execute the given statement. If rows are returned, output these in a
         tabulated format.
         """
+        if out is None:
+            out = stdout
         if isinstance(stmt, ustr):
             logger.debug(" - executing %r", stmt.encode('ascii', 'replace'))
         else:
@@ -202,9 +204,10 @@ class MigrationStep(StepBase):
                     if len(value) > column_sizes[ix]:
                         column_sizes[ix] = len(value)
             format = '|'.join(' %%- %ds ' % size for size in column_sizes)
-            out.write(format % tuple(column_names) + "\n")
+            format += '\n'
+            out.write(format % tuple(column_names))
             out.write('+'.join('-' * (size + 2) for size in column_sizes)
-                      + "\n")
+                      + '\n')
             for row in result:
                 out.write(format % tuple(row))
             out.write(plural(len(result), '(%d row)', '(%d rows)') + "\n")
diff --git a/yoyo/tests/test_migrations.py b/yoyo/tests/test_migrations.py
index 90c900d..dad6aae 100644
--- a/yoyo/tests/test_migrations.py
+++ b/yoyo/tests/test_migrations.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import pytest
-from mock import Mock
+from mock import Mock, patch
 
 from yoyo.connections import get_backend
 from yoyo import read_migrations
@@ -212,6 +212,27 @@ def test_migrations_can_import_step_and_group(tmpdir):
     assert cursor.fetchall() == [(1,)]
 
 
+@with_migrations(
+    '''
+    step("CREATE TABLE _yoyo_test (id INT, c VARCHAR(1))")
+    step("INSERT INTO _yoyo_test VALUES (1, 'a')")
+    step("INSERT INTO _yoyo_test VALUES (2, 'b')")
+    step("SELECT * FROM _yoyo_test")
+    '''
+)
+def test_migrations_display_selected_data(tmpdir):
+    backend = get_backend(dburi)
+    migrations = read_migrations(tmpdir)
+    with patch('yoyo.migrations.stdout') as stdout:
+        backend.apply_migrations(migrations)
+        written = ''.join(a[0] for a, kw in stdout.write.call_args_list)
+        assert written == (' id | c \n'
+                           '----+---\n'
+                           ' 1  | a \n'
+                           ' 2  | b \n'
+                           '(2 rows)\n')
+
+
 class TestTopologicalSort(object):
 
     def get_mock_migrations(self):