1 files changed, 21 insertions, 1 deletions
diff --git a/test/test_main.py b/test/test_main.py
index 7731603..64cba0a 100644
--- a/test/test_main.py
+++ b/test/test_main.py
@@ -25,7 +25,7 @@ os.environ['MYSQL_TEST_LOGIN_FILE'] = login_path_file
 CLI_ARGS = ['--user', USER, '--host', HOST, '--port', PORT,
             '--password', PASSWORD, '--myclirc', default_config_file,
             '--defaults-file', default_config_file,
-            '_test_db']
+            'mycli_test_db']
 
 
 @dbtest
@@ -283,6 +283,20 @@ def test_list_dsn():
         assert result.output == "test : mysql://test/test\n"
 
 
+def test_prettify_statement():
+    statement = 'SELECT 1'
+    m = MyCli()
+    pretty_statement = m.handle_prettify_binding(statement)
+    assert pretty_statement == 'SELECT\n    1;'
+
+
+def test_unprettify_statement():
+    statement = 'SELECT\n    1'
+    m = MyCli()
+    unpretty_statement = m.handle_unprettify_binding(statement)
+    assert unpretty_statement == 'SELECT 1;'
+
+
 def test_list_ssh_config():
     runner = CliRunner()
     with NamedTemporaryFile(mode="w") as ssh_config:
@@ -305,19 +319,25 @@ def test_dsn(monkeypatch):
     # Setup classes to mock mycli.main.MyCli
     class Formatter:
         format_name = None
+
     class Logger:
         def debug(self, *args, **args_dict):
             pass
+
         def warning(self, *args, **args_dict):
             pass
+
     class MockMyCli:
         config = {'alias_dsn': {}}
+
         def __init__(self, **args):
             self.logger = Logger()
             self.destructive_warning = False
             self.formatter = Formatter()
+
         def connect(self, **args):
             MockMyCli.connect_args = args
+
         def run_query(self, query, new_line=True):
             pass