From 97df5c1db305b626ffa0b80055361b7b28e69cec Mon Sep 17 00:00:00 2001
From: Stefan Matting <stefan@wire.com>
Date: Fri, 16 Sep 2022 14:04:43 +0200
Subject: [PATCH] Show callstack HUnit failures, not only source loc

---
 hunit/Test/Tasty/HUnit.hs       |  2 +-
 hunit/Test/Tasty/HUnit/Orig.hs  | 42 ++++++++++++++++++++++-----------
 hunit/Test/Tasty/HUnit/Steps.hs |  2 +-
 3 files changed, 30 insertions(+), 16 deletions(-)

diff --git a/hunit/Test/Tasty/HUnit.hs b/hunit/Test/Tasty/HUnit.hs
index 6ccfdec9..f5d14c35 100644
--- a/hunit/Test/Tasty/HUnit.hs
+++ b/hunit/Test/Tasty/HUnit.hs
@@ -103,6 +103,6 @@ instance IsTest TestCase where
     return $
       case hunitResult of
         Right info -> testPassed info
-        Left (HUnitFailure mbloc message) -> testFailed $ prependLocation mbloc message
+        Left (HUnitFailure cs message) -> testFailed $ prependCallStack cs message
 
   testOptions = return []
diff --git a/hunit/Test/Tasty/HUnit/Orig.hs b/hunit/Test/Tasty/HUnit/Orig.hs
index 55f674ad..885a10db 100644
--- a/hunit/Test/Tasty/HUnit/Orig.hs
+++ b/hunit/Test/Tasty/HUnit/Orig.hs
@@ -1,4 +1,4 @@
-{-# LANGUAGE DeriveDataTypeable, FlexibleInstances, TypeSynonymInstances #-}
+{-# LANGUAGE DeriveDataTypeable, FlexibleInstances, TypeSynonymInstances, RecordWildCards #-}
 
 -- required for HasCallStack by different versions of GHC
 {-# LANGUAGE ConstraintKinds, FlexibleContexts #-}
@@ -11,6 +11,7 @@ import qualified Control.Exception as E
 import Control.Monad
 import Data.Typeable (Typeable)
 import Data.CallStack
+import Data.List
 
 -- Interfaces
 -- ----------
@@ -38,12 +39,7 @@ assertFailure
   :: HasCallStack
   => String -- ^ A message that is displayed with the assertion failure
   -> IO a
-assertFailure msg = E.throwIO (HUnitFailure location msg)
-  where
-    location :: Maybe SrcLoc
-    location = case reverse callStack of
-      (_, loc) : _ -> Just loc
-      [] -> Nothing
+assertFailure msg = E.throwIO (HUnitFailure callStack msg)
 
 -- Conditional Assertion Functions
 -- -------------------------------
@@ -133,16 +129,34 @@ instance (AssertionPredicable t) => AssertionPredicable (IO t)
 
 
 -- | Exception thrown by 'assertFailure' etc.
-data HUnitFailure = HUnitFailure (Maybe SrcLoc) String
+data HUnitFailure = HUnitFailure CallStack String
     deriving (Eq, Show, Typeable)
 instance E.Exception HUnitFailure where
-  displayException (HUnitFailure mbloc s) = prependLocation mbloc s
+  displayException (HUnitFailure mbloc s) = prependCallStack mbloc s
+
+prependCallStack :: CallStack -> String -> String
+prependCallStack cs s =
+  "Error message: " <> s <> "\n\n" <> prettyCallStack cs
 
-prependLocation :: Maybe SrcLoc -> String -> String
-prependLocation mbloc s =
-  case mbloc of
-    Nothing -> s
-    Just loc -> srcLocFile loc ++ ":" ++ show (srcLocStartLine loc) ++ ":\n" ++ s
+prettyCallStack :: CallStack -> String
+prettyCallStack = intercalate "\n" . prettyCallStackLines
+
+prettyCallStackLines :: CallStack -> [String]
+prettyCallStackLines cs = case cs of
+  []  -> []
+  stk -> "CallStack (from HasCallStack):"
+       : map (("  " ++) . prettyCallSite) stk
+  where
+    prettyCallSite (f, loc) = f ++ ", called at " ++ prettySrcLoc loc
+
+prettySrcLoc :: SrcLoc -> String
+prettySrcLoc SrcLoc {..}
+  = foldr (++) ""
+      [ srcLocFile, ":"
+      , show srcLocStartLine, ":"
+      , show srcLocStartCol, " in "
+      , srcLocPackage, ":", srcLocModule
+      ]
 
 ----------------------------------------------------------------------
 --                          DEPRECATED CODE
diff --git a/hunit/Test/Tasty/HUnit/Steps.hs b/hunit/Test/Tasty/HUnit/Steps.hs
index 80595f13..d7e75674 100644
--- a/hunit/Test/Tasty/HUnit/Steps.hs
+++ b/hunit/Test/Tasty/HUnit/Steps.hs
@@ -31,7 +31,7 @@ instance IsTest TestCaseSteps where
         atomicModifyIORef ref (\l -> ((tme,msg):l, ()))
 
     hunitResult <- (Right <$> assertionFn stepFn) `catch`
-        \(SomeException ex) -> return $ Left (displayException ex)
+      \(SomeException ex) -> return $ Left (displayException ex)
 
     endTime <- getTime