Add tests for the format_output.

1 files changed, 20 insertions, 1 deletions

diff --git a/tests/test_main.py b/tests/test_main.py
index 0d026994..1642d1f5 100644
--- a/tests/test_main.py
+++ b/tests/test_main.py
@@ -4,7 +4,7 @@ try:
     import setproctitle
 except ImportError:
     setproctitle = None
-from pgcli.main import obfuscate_process_password
+from pgcli.main import obfuscate_process_password, format_output
 
 
 @pytest.mark.skipif(platform.system() == 'Windows',
@@ -39,3 +39,22 @@ def test_obfuscate_process_password():
     assert title == expected
 
     setproctitle.setproctitle(original_title)
+
+def test_format_output():
+    results = format_output('Title', [('abc', 'def')], ['head1', 'head2'],
+                            'test status', 'psql')
+    expected = ['Title', '+---------+---------+\n| head1   | head2   |\n|---------+---------|\n| abc     | def     |\n+---------+---------+', 'test status']
+    assert results == expected
+
+def test_format_output_auto_expand():
+    table_results = format_output('Title', [('abc', 'def')],
+                                  ['head1', 'head2'], 'test status', 'psql',
+                                  max_width=100)
+    table = ['Title', '+---------+---------+\n| head1   | head2   |\n|---------+---------|\n| abc     | def     |\n+---------+---------+', 'test status']
+    assert table_results == table
+
+    expanded_results = format_output('Title', [('abc', 'def')],
+                                     ['head1', 'head2'], 'test status', 'psql',
+                                     max_width=1)
+    expanded = ['Title', u'-[ RECORD 0 ]-------------------------\nhead1 | abc\nhead2 | def\n', 'test status']
+    assert expanded_results == expanded