webssh

commit 90e7ea032756f8aca62b84daee235c21541cc071
parent fdcf1718c110462ee682157e64e7b45c3d9211af
Author: Sheng <webmaster0115@gmail.com>
Date:   Wed,  3 Oct 2018 19:22:20 +0800

Fixed test_print_version

Diffstat:
M
tests/test_settings.py
 | 
15
+++++++++++++--

1 file changed, 13 insertions(+), 2 deletions(-)
diff --git a/tests/test_settings.py b/tests/test_settings.py
@@ -1,3 +1,5 @@
+import io
+import sys
 import os.path
 import unittest
 import paramiko
@@ -8,15 +10,24 @@ from webssh.policy import load_host_keys
 from webssh.settings import (
     get_host_keys_settings, get_policy_setting, base_dir, print_version
 )
+from webssh.utils import UnicodeType
 from webssh._version import __version__
 
 
 class TestSettings(unittest.TestCase):
 
     def test_print_version(self):
-        self.assertNotEqual(print_version(False), 2, msg=__version__)
+        sys_stdout = sys.stdout
+        sys.stdout = io.StringIO() if UnicodeType == str else io.BytesIO()
+
+        self.assertEqual(print_version(False), None)
+        self.assertEqual(sys.stdout.getvalue(), '')
+
         with self.assertRaises(SystemExit):
-            self.assertEqual(print_version(True), 2, msg=__version__)
+            self.assertEqual(print_version(True), None)
+        self.assertEqual(sys.stdout.getvalue(), __version__ + '\n')
+
+        sys.stdout = sys_stdout
 
     def test_get_host_keys_settings(self):
         options.hostFile = ''