From f2b284c407ed5afd2061fbe7481aa617f39277b8 Mon Sep 17 00:00:00 2001
From: Edwin Eefting <edwin@datux.nl>
Date: Wed, 8 Jul 2020 17:53:54 +0200
Subject: [PATCH] make tests python2 compat

---
 basetest.py              | 32 ++++++++++++++++++++++++++++++++
 test_externalfailures.py |  8 ++++----
 test_zfsautobackup.py    |  8 +++-----
 3 files changed, 39 insertions(+), 9 deletions(-)

diff --git a/basetest.py b/basetest.py
index 10a954f..404fb95 100644
--- a/basetest.py
+++ b/basetest.py
@@ -10,10 +10,42 @@ import time
 from pprint import *
 from bin.zfs_autobackup import *
 from mock import *
+import contextlib 
+import sys
+import io
 
 TEST_POOLS="test_source1 test_source2 test_target1"
 
 
+# for python2 compatibility
+if sys.version_info.major==2:
+    OutputIO=io.BytesIO
+else:
+    OutputIO=io.StringIO
+
+
+# for python2 compatibility (python 3 has this already)
+@contextlib.contextmanager
+def redirect_stdout(target):
+    original = sys.stdout
+    try:
+        sys.stdout = target
+        yield
+    finally:
+        sys.stdout = original
+
+# for python2 compatibility (python 3 has this already)
+@contextlib.contextmanager
+def redirect_stderr(target):
+    original = sys.stderr
+    try:
+        sys.stderr = target
+        yield
+    finally:
+        sys.stderr = original
+
+
+
 def shelltest(cmd):
     """execute and print result as nice copypastable string for unit tests (adds extra newlines on top/bottom)"""
     ret=(subprocess.check_output(cmd , shell=True).decode('utf-8'))
diff --git a/test_externalfailures.py b/test_externalfailures.py
index e084a53..14e5481 100644
--- a/test_externalfailures.py
+++ b/test_externalfailures.py
@@ -1,7 +1,6 @@
 
 from basetest import *
-import contextlib
-import io
+
 
 class TestZfsNode(unittest2.TestCase):
 
@@ -32,9 +31,10 @@ class TestZfsNode(unittest2.TestCase):
         r=shelltest("rm /test_target1/waste")
         r=shelltest("zfs umount test_target1")
 
+
         #should resume and succeed
-        with io.StringIO() as buf:
-            with contextlib.redirect_stdout(buf):
+        with OutputIO() as buf:
+            with redirect_stdout(buf):
                 with patch('time.strftime', return_value="20101111000002"):
                     self.assertFalse(ZfsAutobackup("test test_target1 --verbose --allow-empty".split(" ")).run())
 
diff --git a/test_zfsautobackup.py b/test_zfsautobackup.py
index f2963e7..a965791 100644
--- a/test_zfsautobackup.py
+++ b/test_zfsautobackup.py
@@ -1,7 +1,5 @@
 from basetest import *
 import time
-import contextlib
-import io
 
 
 
@@ -16,13 +14,13 @@ class TestZfsAutobackup(unittest2.TestCase):
         with self.subTest("no datasets selected"):
             #should resume and succeed
             
-            with io.StringIO() as buf:
-                with contextlib.redirect_stderr(buf):
+            with OutputIO() as buf:
+                with redirect_stderr(buf):
                     with patch('time.strftime', return_value="20101111000000"):
                         self.assertTrue(ZfsAutobackup("nonexisting test_target1 --verbose --debug".split(" ")).run())
 
                 print(buf.getvalue())
-                #did we really resume?
+                #correct message?
                 self.assertIn("No source filesystems selected", buf.getvalue())