Move assertBytesEqual to base test class, improve it, and hook into assertEqual

2 files changed, 12 insertions, 6 deletions

diff --git a/Lib/test/test_email/__init__.py b/Lib/test/test_email/__init__.py
index 2ef21ba..54d1c9a 100644
--- a/Lib/test/test_email/__init__.py
+++ b/Lib/test/test_email/__init__.py
@@ -25,6 +25,10 @@ def openfile(filename, *args, **kws):
 # Base test class
 class TestEmailBase(unittest.TestCase):
 
+    def __init__(self, *args, **kw):
+        super().__init__(*args, **kw)
+        self.addTypeEqualityFunc(bytes, self.assertBytesEqual)
+
     def ndiffAssertEqual(self, first, second):
         """Like assertEqual except use ndiff for readable output."""
         if first != second:
@@ -38,3 +42,10 @@ class TestEmailBase(unittest.TestCase):
     def _msgobj(self, filename):
         with openfile(filename) as fp:
             return email.message_from_file(fp)
+
+    def _bytes_repr(self, b):
+        return [repr(x) for x in b.splitlines(True)]
+
+    def assertBytesEqual(self, first, second, msg):
+        """Our byte strings are really encoded strings; improve diff output"""
+        self.assertEqual(self._bytes_repr(first), self._bytes_repr(second))
diff --git a/Lib/test/test_email/test_email.py b/Lib/test/test_email/test_email.py
index 2519cc0..3d86b6a 100644
--- a/Lib/test/test_email/test_email.py
+++ b/Lib/test/test_email/test_email.py
@@ -3161,12 +3161,7 @@ class BaseTestBytesGeneratorIdempotent:
         b = BytesIO()
         g = email.generator.BytesGenerator(b, maxheaderlen=0)
         g.flatten(msg, unixfrom=unixfrom, linesep=self.linesep)
-        self.assertByteStringsEqual(data, b.getvalue())
-
-    def assertByteStringsEqual(self, str1, str2):
-        # Not using self.blinesep here is intentional.  This way the output
-        # is more useful when the failure results in mixed line endings.
-        self.assertListEqual(str1.split(b'\n'), str2.split(b'\n'))
+        self.assertEqual(data, b.getvalue())
 
 
 class TestBytesGeneratorIdempotentNL(BaseTestBytesGeneratorIdempotent,