2 files changed, 12 insertions, 15 deletions
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 4ae8da3..7ce57eb 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -15,25 +15,23 @@ from io import StringIO, BytesIO
 import rsa
 import rsa.cli
 import rsa.util
+from rsa._compat import PY2
 
-if sys.version_info[0] < 3:
-    def make_buffer():
+
+def make_buffer():
+    if PY2:
         return BytesIO()
+    buf = StringIO()
+    buf.buffer = BytesIO()
+    return buf
 
 
-    def get_bytes_out(out):
-        # Python 2.x writes 'str' to stdout:
+def get_bytes_out(out):
+    if PY2:
+        # Python 2.x writes 'str' to stdout
         return out.getvalue()
-else:
-    def make_buffer():
-        buf = StringIO()
-        buf.buffer = BytesIO()
-        return buf
-
-
-    def get_bytes_out(out):
-        # Python 3.x writes 'bytes' to stdout.buffer:
-        return out.buffer.getvalue()
+    # Python 3.x writes 'bytes' to stdout.buffer
+    return out.buffer.getvalue()
 
 
 @contextmanager
diff --git a/tests/test_compat.py b/tests/test_compat.py
index 0013155..a47f890 100644
--- a/tests/test_compat.py
+++ b/tests/test_compat.py
@@ -16,7 +16,6 @@
 
 import unittest
 import struct
-import sys
 
 from rsa._compat import byte, is_bytes, range