JetBrains#236: store diff data it tls, rethraw real exception When cu…

…stom exception is thrown, stack trace is broken

teamcity/diff_tools.py

@@ -31,12 +31,13 @@ def _patched_equals(self, first, second, msg=None):
         try:
             old(self, first, second, msg)
             return
-        except AssertionError as native_error:
+        except AssertionError as e:
             if not test_filter or test_filter(self):
-                error = EqualsAssertionError(first, second, msg)
+                error = EqualsAssertionError(first, second, msg, real_exception=e)
                 if error.can_be_serialized():
-                    raise error
-            raise native_error
+                    from .jb_local_exc_store import store_exception
+                    store_exception(error)
+            raise
 
     unittest.TestCase.assertEqual = _patched_equals
 
@@ -50,7 +51,9 @@ class EqualsAssertionError(AssertionError):
     MESSAGE_SEP = " :: "
     NOT_EQ_SEP = " != "
 
-    def __init__(self, expected, actual, msg=None, preformated=False):
+    # Real exception could be provided, but not serialized
+    def __init__(self, expected, actual, msg=None, preformated=False, real_exception=None):
+        self.real_exception = real_exception
         self.expected = expected
         self.actual = actual
         self.msg = text_type(msg)

teamcity/jb_local_exc_store.py

@@ -0,0 +1,16 @@
+import threading
+
+_JB_PYTEST_LOCAL = threading.local()
+
+
+def store_exception(exc):
+    _JB_PYTEST_LOCAL.exception = exc
+
+
+def get_exception():
+    try:
+        exception = _JB_PYTEST_LOCAL.exception
+        _JB_PYTEST_LOCAL.exception = None
+        return exception
+    except AttributeError:
+        return None

teamcity/nose_report.py

@@ -169,9 +169,10 @@ def report_fail(self, test, fail_type, err):
             details = details[:start_index] + details[end_index + len(_captured_output_end_marker):]
 
         try:
-            error = err[1]
-            if isinstance(error, EqualsAssertionError):
-                details = convert_error_to_string(err, 2)
+            from .jb_local_exc_store import get_exception
+            error = get_exception()
+            if isinstance(err[1], AssertionError) and isinstance(error, EqualsAssertionError):
+                details = convert_error_to_string(err, 1)
                 self.messages.testFailed(test_id, message=error.msg, details=details, flowId=test_id, comparison_failure=error)
                 return
         except Exception:

teamcity/pytest_plugin.py

@@ -286,12 +286,21 @@ def report_test_failure(self, test_id, report, message=None, report_output=True)
         except Exception:
             pass
 
+        if not diff_error:
+            from .jb_local_exc_store import get_exception
+            diff_error = get_exception()
+
         if diff_error:
             # Cut everything after postfix: it is internal view of DiffError
             strace = str(report.longrepr)
             data_postfix = "_ _ _ _ _"
+            # Error message in pytest must be in "file.py:22 AssertionError" format
+            # This message goes to strace
+            # With custom error we must add real exception class explicitly
             if data_postfix in strace:
-                strace = strace[0:strace.index(data_postfix)]
+                strace = strace[0:strace.index(data_postfix)].strip()
+                if strace.endswith(":") and diff_error.real_exception:
+                    strace += " " + type(diff_error.real_exception).__name__
             self.teamcity.testFailed(test_id, diff_error.msg if diff_error.msg else message, strace,
                                      flowId=test_id,
                                      comparison_failure=diff_error

teamcity/unittestpy.py

@@ -195,7 +195,8 @@ def report_fail(self, test, fail_type, err):
 
         diff_failed = None
         try:
-            error = err[1]
+            from .jb_local_exc_store import get_exception
+            error = get_exception()
             if isinstance(error, EqualsAssertionError):
                 diff_failed = error
         except Exception:

tests/guinea-pigs/pytest/unittest_error_test.py

@@ -8,3 +8,6 @@ def test_error(self):
 
     def test_fail(self):
         self.assertTrue(False)
+
+    def test_fail_diff(self):
+        self.assertEqual("A", "B")

tests/integration-tests/pytest_integration_test.py

@@ -211,17 +211,20 @@ def test_unittest_error(venv):
     ms = assert_service_messages(
         output,
         [
-            ServiceMessage('testCount', {'count': "2"}),
+            ServiceMessage('testCount', {'count': "3"}),
             ServiceMessage('testStarted', {'name': 'tests.guinea-pigs.pytest.unittest_error_test.TestErrorFail.test_error'}),
             ServiceMessage('testFailed', {}),
             ServiceMessage('testFinished', {'name': 'tests.guinea-pigs.pytest.unittest_error_test.TestErrorFail.test_error'}),
             ServiceMessage('testStarted', {'name': 'tests.guinea-pigs.pytest.unittest_error_test.TestErrorFail.test_fail'}),
             ServiceMessage('testFailed', {}),
             ServiceMessage('testFinished', {'name': 'tests.guinea-pigs.pytest.unittest_error_test.TestErrorFail.test_fail'}),
+            ServiceMessage('testStarted', {'name': 'tests.guinea-pigs.pytest.unittest_error_test.TestErrorFail.test_fail_diff'}),
+            ServiceMessage('testFailed', {}),
+            ServiceMessage('testFinished', {'name': 'tests.guinea-pigs.pytest.unittest_error_test.TestErrorFail.test_fail_diff'}),
         ])
     assert ms[2].params["details"].find("raise Exception") > 0
     assert ms[2].params["details"].find("oops") > 0
-    assert ms[5].params["details"].find("AssertionError") > 0
+    assert ms[8].params["details"].find("unittest_error_test.py:13: AssertionError") > 0
 
 
 def test_fixture_error(venv):