Fix Python 3 unit tests

2 files changed, 20 insertions, 7 deletions

diff --git a/suite/generator.py b/suite/generator.py
index 78925e97..3d15a813 100755
--- a/suite/generator.py
+++ b/suite/generator.py
@@ -7,6 +7,14 @@ from optparse import OptionParser
 import testcommon
 import time
 
+if sys.version_info.major == 3:
+    # The differences between the behavior of str() on Python 2 and Python 3 make it very difficult to make
+    # test results created on Python 3 work when testing on Python 2. Test results generated with Python 2
+    # will work when testing on both Python 2 and Python 3.
+    # FIXME: When Python 2 is dropped fully, use Python 3 to generate the tests.
+    print("Generate unit tests on Python 2. Python 3 generated test results are NOT COMPATIBLE with Python 2.")
+    sys.exit(1)
+
 unit_test_template = """#!/usr/bin/env python
 # This is an auto generated unit test file do not edit directly
 import os
@@ -84,7 +92,7 @@ class TestBinaryNinjaAPI(unittest.TestCase):
         pickle_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "oracle.pkl")
         try:
             # Python 2 does not have the encodings option
-            self.oracle_test_data = pickle.load(open(pickle_path, "rb"), encoding='charmap')
+            self.oracle_test_data = pickle.load(open(pickle_path, "rb"), encoding='utf8')
         except TypeError:
             self.oracle_test_data = pickle.load(open(pickle_path, "rb"))
         self.verifybuilder = testcommon.VerifyBuilder("{3}")
@@ -99,7 +107,7 @@ class TestBinaryNinjaAPI(unittest.TestCase):
         self.assertTrue(pickle_path, "Test pickle doesn't exist")
         try:
             # Python 2 does not have the encodings option
-            binary_oracle = pickle.load(open(pickle_path, "rb"), encoding='charmap')
+            binary_oracle = pickle.load(open(pickle_path, "rb"), encoding='utf8')
         except TypeError:
             binary_oracle = pickle.load(open(pickle_path, "rb"))
 
@@ -185,7 +193,7 @@ class UnitTestFile:
         api_path = map(lambda x: '"{0}"'.format(x), api_path.split(os.sep))
         api_path = '{0}'.format(', '.join(api_path))
         test_store = self.test_store.replace(os.sep, '/') if os.name == 'nt' else self.test_store
-        self.f.write(self.template.format(self.outdir, self.tests, self.binary_tests, test_store, api_path).encode('charmap'))
+        self.f.write(self.template.format(self.outdir, self.tests, self.binary_tests, test_store, api_path).encode('utf8'))
         self.f.close()
 
     def add_verify(self, test_name):
diff --git a/suite/testcommon.py b/suite/testcommon.py
index ce1039f1..25a8b4a5 100644
--- a/suite/testcommon.py
+++ b/suite/testcommon.py
@@ -54,6 +54,11 @@ def fixSet(string):
         return string
 
 
+def fixStrRepr(string):
+    # Python 2 and Python 3 represent Unicode character reprs differently
+    return string.replace(b"\xe2\x80\xa6".decode("utf8"), "\\xe2\\x80\\xa6")
+
+
 def get_file_list(test_store_rel):
     test_store = os.path.join(os.path.dirname(__file__), test_store_rel)
     all_files = []
@@ -199,7 +204,7 @@ class BinaryViewTestBuilder(Builder):
                             prefixList.append(str(contents))
                         else:
                             prefixList.append(i)
-                    retinfo.append("Function: {:x} Instruction: {:x} Prefix operands: {}".format(func.start, ins.address, str(prefixList)))
+                    retinfo.append("Function: {:x} Instruction: {:x} Prefix operands: {}".format(func.start, ins.address, fixStrRepr(str(prefixList))))
 
                     postfixList = []
                     for i in ins.postfix_operands:
@@ -210,7 +215,7 @@ class BinaryViewTestBuilder(Builder):
                             postfixList.append(str(contents))
                         else:
                             postfixList.append(i)
-                    retinfo.append("Function: {:x} Instruction: {:x} Postfix operands: {}".format(func.start, ins.address, str(postfixList)))
+                    retinfo.append("Function: {:x} Instruction: {:x} Postfix operands: {}".format(func.start, ins.address, fixStrRepr(str(postfixList))))
 
                     retinfo.append("Function: {:x} Instruction: {:x} SSA form: {}".format(func.start, ins.address, str(ins.ssa_form)))
                     retinfo.append("Function: {:x} Instruction: {:x} Non-SSA form: {}".format(func.start, ins.address, str(ins.non_ssa_form)))
@@ -264,7 +269,7 @@ class BinaryViewTestBuilder(Builder):
                             prefixList.append(str(contents))
                         else:
                             prefixList.append(str(i))
-                    retinfo.append("Function: {:x} Instruction: {:x} Prefix operands:  {}".format(func.start, ins.address, str(sorted(prefixList))))
+                    retinfo.append("Function: {:x} Instruction: {:x} Prefix operands:  {}".format(func.start, ins.address, fixStrRepr(str(sorted(prefixList)))))
                     postfixList = []
                     for i in ins.postfix_operands:
                         if isinstance(i, float) and 'e' in str(i):
@@ -279,7 +284,7 @@ class BinaryViewTestBuilder(Builder):
                         else:
                             postfixList.append(str(i))
 
-                    retinfo.append("Function: {:x} Instruction: {:x} Postfix operands:  {}".format(func.start, ins.address, str(sorted(postfixList))))
+                    retinfo.append("Function: {:x} Instruction: {:x} Postfix operands:  {}".format(func.start, ins.address, fixStrRepr(str(sorted(postfixList)))))
                     retinfo.append("Function: {:x} Instruction: {:x} SSA form:  {}".format(func.start, ins.address, str(ins.ssa_form)))
                     retinfo.append("Function: {:x} Instruction: {:x} Non-SSA form: {}".format(func.start, ins.address, str(ins.non_ssa_form)))
         return fixOutput(retinfo)