unittest assertEqual difference output foiled by newlines

Author: nanjekyejoannah

Lib/unittest/case.py

@@ -1238,6 +1238,11 @@ def assertCountEqual(self, first, second, msg=None):
             msg = self._formatMessage(msg, standardMsg)
             self.fail(msg)
 
+    def addTrailingNewLine(self, line):
+        if line != '' and line[-1] != '\n':
+            line = line + '\n'
+        return line
+
     def assertMultiLineEqual(self, first, second, msg=None):
         """Assert that two multi-line strings are equal."""
         self.assertIsInstance(first, str, 'First argument is not a string')
@@ -1248,13 +1253,16 @@ def assertMultiLineEqual(self, first, second, msg=None):
             if (len(first) > self._diffThreshold or
                 len(second) > self._diffThreshold):
                 self._baseAssertEqual(first, second, msg)
-            firstlines = first.splitlines(keepends=True)
+            firstlines = first.splitlines(keepends=True) 
             secondlines = second.splitlines(keepends=True)
             if len(firstlines) == 1 and first.strip('\r\n') == first:
                 firstlines = [first + '\n']
                 secondlines = [second + '\n']
             standardMsg = '%s != %s' % _common_shorten_repr(first, second)
-            diff = '\n' + ''.join(difflib.ndiff(firstlines, secondlines))
+            difflines = list(difflib.ndiff(firstlines, secondlines))
+            if len(difflines) > 1:
+                difflines = [self.addTrailingNewLine(line) for line in difflines]
+            diff = '\n' + ''.join(difflines)
             standardMsg = self._truncateMessage(standardMsg, diff)
             self.fail(self._formatMessage(msg, standardMsg))
 

Lib/unittest/test/test_case.py

@@ -1839,5 +1839,20 @@ def test2(self):
             self.assertEqual(MyException.ninstance, 0)
 
 
+class testAssertEqualSingleLine(unittest.TestCase):
+
+    def test_trailing_new_line_at_end(self):
+        self.assertEqual("abc\n", "abc\n")
+    def test_trailing_space_at_end(self):
+        self.assertEqual("abc ", "abc ")
+    def test_no_trailing_new_line(self):
+        self.assertEqual("abc", "abc")
+    def test_new_line_at_beginning(self):
+        self.assertEqual("\nabc", "\nabc")
+    def test_new_line_at_start_and_end(self):
+        self.assertEqual("\nabc\n", "\nabc\n")
+    def test_with_space_at_start_and_end(self):
+        self.assertEqual(" abc ", " abc ")
+
 if __name__ == "__main__":
     unittest.main()