From ea6496ca24178e87f4ab13c75f125ddc75bc8c81 Mon Sep 17 00:00:00 2001
From: Evgeny Poberezkin <2769109+epoberezkin@users.noreply.github.com>
Date: Sat, 9 Jan 2021 20:35:19 +0000
Subject: [PATCH] better MonadUnliftIO instance for ExceptT; fix tests

---
 src/Simplex/Messaging/Util.hs | 16 +++++++++++++---
 tests/AgentTests/SQLite.hs    |  1 +
 2 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/src/Simplex/Messaging/Util.hs b/src/Simplex/Messaging/Util.hs
index 3b96d0a99..88efef4cd 100644
--- a/src/Simplex/Messaging/Util.hs
+++ b/src/Simplex/Messaging/Util.hs
@@ -1,3 +1,6 @@
+{-# LANGUAGE InstanceSigs #-}
+{-# LANGUAGE RankNTypes #-}
+
 module Simplex.Messaging.Util where
 
 import Control.Monad.Except
@@ -6,10 +9,17 @@ import UnliftIO.Async
 import UnliftIO.Exception (Exception)
 import qualified UnliftIO.Exception as E
 
+newtype InternalException e = InternalException {unInternalException :: e}
+  deriving (Eq, Show)
+
+instance Exception e => Exception (InternalException e)
+
 instance (MonadUnliftIO m, Exception e) => MonadUnliftIO (ExceptT e m) where
-  withRunInIO inner = ExceptT . E.try $
-    withRunInIO $ \run ->
-      inner (run . (either E.throwIO pure <=< runExceptT))
+  withRunInIO :: ((forall a. ExceptT e m a -> IO a) -> IO b) -> ExceptT e m b
+  withRunInIO exceptToIO =
+    withExceptT unInternalException . ExceptT . E.try $
+      withRunInIO $ \run ->
+        exceptToIO $ run . (either (E.throwIO . InternalException) return <=< runExceptT)
 
 raceAny_ :: MonadUnliftIO m => [m a] -> m ()
 raceAny_ = r []
diff --git a/tests/AgentTests/SQLite.hs b/tests/AgentTests/SQLite.hs
index ffafb6c94..63fb06695 100644
--- a/tests/AgentTests/SQLite.hs
+++ b/tests/AgentTests/SQLite.hs
@@ -8,6 +8,7 @@ import Control.Monad.Except
 import qualified Database.SQLite.Simple as DB
 import Simplex.Messaging.Agent.Store
 import Simplex.Messaging.Agent.Store.SQLite
+import Simplex.Messaging.Agent.Store.Types
 import Simplex.Messaging.Agent.Transmission
 import Test.Hspec
 import UnliftIO.Directory