From 7b074c1ddf21daf0d8320252d34b253fa7a3fc93 Mon Sep 17 00:00:00 2001
From: Pepe Iborra <pepeiborra@gmail.com>
Date: Fri, 12 Apr 2019 15:15:54 +0100
Subject: [PATCH 1/7] Stateless version of MockT

---
 library/Control/Monad/Mock.hs                 | 36 ++++---
 library/Control/Monad/Mock/Stateless.hs       | 95 +++++++++++++++++++
 library/Control/Monad/Mock/TH.hs              | 17 ++--
 .../Control/Monad/Mock/StatelessSpec.hs       | 71 ++++++++++++++
 4 files changed, 198 insertions(+), 21 deletions(-)
 create mode 100644 library/Control/Monad/Mock/Stateless.hs
 create mode 100644 test-suite/Control/Monad/Mock/StatelessSpec.hs

diff --git a/library/Control/Monad/Mock.hs b/library/Control/Monad/Mock.hs
index b5d01d3..8d8bf25 100644
--- a/library/Control/Monad/Mock.hs
+++ b/library/Control/Monad/Mock.hs
@@ -63,12 +63,15 @@ boilerplate, look at 'Control.Monad.Mock.TH.makeAction' from
 "Control.Monad.Mock.TH".
 -}
 module Control.Monad.Mock
-  ( -- * The MockT monad transformer
-    MockT
+  (
+  -- * The MonadMock class
+    MonadMock(..)
+
+  -- * The MockT monad transformer
+  , MockT
   , Mock
   , runMockT
   , runMock
-  , mockAction
 
   -- * Actions and actions with results
   , Action(..)
@@ -164,15 +167,18 @@ runMockT actions (MockT x) = runStateT x actions >>= \case
 runMock :: forall f a. Action f => [WithResult f] -> Mock f a -> a
 runMock actions x = runIdentity $ runMockT actions x
 
--- | Logs a method call within a mock.
-mockAction :: (Action f, Monad m) => String -> f r -> MockT f m r
-mockAction fnName action = MockT $ get >>= \case
-  [] -> error'
-     $ "runMockT: expected end of program, called " ++ fnName ++ "\n"
-    ++ "  given action: " ++ showAction action ++ "\n"
-  (action' :-> r) : actions
-    | Just Refl <- action `eqAction` action' -> put actions >> return r
-    | otherwise -> error'
-         $ "runMockT: argument mismatch in " ++ fnName ++ "\n"
-        ++ "  given: " ++ showAction action ++ "\n"
-        ++ "  expected: " ++ showAction action' ++ "\n"
+class MonadMock f m where
+  -- | Logs a method call within a mock.
+  mockAction :: Action f => String -> f r -> m r
+
+instance Monad m => MonadMock f (MockT f m) where
+  mockAction fnName action = MockT $ get >>= \case
+    [] -> error'
+      $ "runMockT: expected end of program, called " ++ fnName ++ "\n"
+      ++ "  given action: " ++ showAction action ++ "\n"
+    (action' :-> r) : actions
+      | Just Refl <- action `eqAction` action' -> put actions >> return r
+      | otherwise -> error'
+          $ "runMockT: argument mismatch in " ++ fnName ++ "\n"
+          ++ "  given: " ++ showAction action ++ "\n"
+          ++ "  expected: " ++ showAction action' ++ "\n"
diff --git a/library/Control/Monad/Mock/Stateless.hs b/library/Control/Monad/Mock/Stateless.hs
new file mode 100644
index 0000000..bb3ab4f
--- /dev/null
+++ b/library/Control/Monad/Mock/Stateless.hs
@@ -0,0 +1,95 @@
+{-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE CPP                  #-}
+{-# LANGUAGE UndecidableInstances #-}
+
+-- | A version of 'MockT' with a stateless 'MonadTransControl' instance
+module Control.Monad.Mock.Stateless
+  (
+  -- * The MonadMock class
+    MonadMock(..)
+
+  -- * The MockT monad transformer
+  , MockT
+  , Mock
+  , runMockT
+  , runMock
+  , MockT_
+
+  -- * Actions and actions with results
+  , Action(..)
+  , WithResult(..)
+  ) where
+
+import Control.Monad.Base (MonadBase)
+import Control.Monad.Catch (MonadCatch, MonadThrow, MonadMask)
+import Control.Monad.Cont (MonadCont)
+import Control.Monad.Except (MonadError)
+import Control.Monad.IO.Class (MonadIO)
+import Control.Monad.Primitive (PrimMonad(..))
+import Control.Monad.Reader (ReaderT(..), MonadReader(..))
+import Control.Monad.State (MonadState)
+import Control.Monad.ST (ST, runST)
+import Control.Monad.Trans (MonadTrans(..))
+import Control.Monad.Trans.Control (MonadBaseControl, MonadTransControl)
+import Control.Monad.Writer (MonadWriter)
+import Data.Primitive.MutVar (MutVar, newMutVar, readMutVar, writeMutVar)
+import Data.Type.Equality ((:~:)(..))
+
+import Control.Monad.Mock (Action(..), MonadMock(..), WithResult(..))
+
+type MockT f m = MockT_ (PrimState m) f m
+
+type Mock s f = MockT f (ST s)
+
+newtype MockT_ s f m a = MockT (ReaderT (MutVar s [WithResult f]) m a)
+  deriving ( Functor, Applicative, Monad, MonadTrans, MonadIO, MonadBase b
+           , MonadState st, MonadCont, MonadError e, MonadWriter w
+           , MonadCatch, MonadThrow, MonadMask
+           , MonadBaseControl b, MonadTransControl )
+
+instance MonadReader r m => MonadReader r (MockT_ s f m) where
+  ask = lift ask
+  local f (MockT act) = MockT $ do
+    env <- ask
+    lift $ local f $ runReaderT act env
+
+runMockT :: forall f m a .
+            (Action f, PrimMonad m) =>
+            [WithResult f] -> MockT f m a -> m a
+runMockT actions (MockT x) = do
+  ref <- newMutVar actions
+  r <- runReaderT x ref
+  leftovers <- readMutVar ref
+  case leftovers of
+    [] -> return r
+    remainingActions -> error'
+      $ "runMockT: expected the following unexecuted actions to be run:\n"
+      ++ unlines (map (\(action :-> _) -> "  " ++ showAction action) remainingActions)
+
+runMock :: forall f a. Action f => [WithResult f] -> (forall s. Mock s f a) -> a
+runMock actions x = runST $ runMockT actions x
+
+instance (PrimMonad m, PrimState m ~ s) => MonadMock f (MockT_ s f m) where
+  mockAction fnName action = MockT $ do
+    ref <- ask
+    results <- lift $ readMutVar ref
+    case results of
+      [] -> error'
+        $ "runMockT: expected end of program, called " ++ fnName ++ "\n"
+        ++ "  given action: " ++ showAction action ++ "\n"
+      (action' :-> r) : actions
+        | Just Refl <- action `eqAction` action' -> do
+            lift $ writeMutVar ref actions
+            return r
+        | otherwise -> error'
+            $ "runMockT: argument mismatch in " ++ fnName ++ "\n"
+            ++ "  given: " ++ showAction action ++ "\n"
+            ++ "  expected: " ++ showAction action' ++ "\n"
+
+
+error' :: String -> a
+#if MIN_VERSION_base(4,9,0)
+error' = errorWithoutStackTrace
+#else
+error' = error
+#endif
diff --git a/library/Control/Monad/Mock/TH.hs b/library/Control/Monad/Mock/TH.hs
index 7093d22..cab7599 100644
--- a/library/Control/Monad/Mock/TH.hs
+++ b/library/Control/Monad/Mock/TH.hs
@@ -81,6 +81,7 @@ spec = describe "copyFile" '$'
 module Control.Monad.Mock.TH (makeAction, deriveAction, ts) where
 
 import Control.Monad (replicateM, when, zipWithM)
+import Control.Monad.Primitive (PrimMonad, PrimState)
 import Data.Char (toUpper)
 import Data.Foldable (traverse_)
 import Data.List (foldl', nub, partition)
@@ -89,6 +90,7 @@ import GHC.Exts (Constraint)
 import Language.Haskell.TH
 
 import Control.Monad.Mock (Action(..), MockT, mockAction)
+import qualified Control.Monad.Mock.Stateless as Stateless
 import Control.Monad.Mock.TH.Internal.TypesQuasi (ts)
 
 -- | Given a list of monadic typeclass constraints of kind @* -> 'Constraint'@,
@@ -119,9 +121,12 @@ makeAction actionNameStr classTs = do
         mkStandaloneDec derivT = standaloneDeriveD' [] (derivT `AppT` (actionTypeCon `AppT` VarT actionParamName))
         standaloneDecs = [mkStandaloneDec (ConT ''Eq), mkStandaloneDec (ConT ''Show)]
     actionInstanceDec <- deriveAction' actionTypeCon actionCons
-    classInstanceDecs <- zipWithM (mkInstance actionTypeCon) classTs methods
+    classInstanceDecs1 <- zipWithM (mkInstance (ConT ''MockT) (const []) actionTypeCon) classTs methods
+    primStateVar <- newName "s"
+    let primStateConstraint baseM = [ConT ''PrimMonad `AppT` baseM, EqualityT `AppT` VarT primStateVar `AppT` (ConT ''PrimState `AppT` baseM)]
+    classInstanceDecs2 <- zipWithM (mkInstance (ConT ''Stateless.MockT_ `AppT` VarT primStateVar) primStateConstraint actionTypeCon) classTs methods
 
-    return $ [actionDec] ++ standaloneDecs ++ [actionInstanceDec] ++ classInstanceDecs
+    return $ [actionDec] ++ standaloneDecs ++ [actionInstanceDec] ++ classInstanceDecs1 ++ classInstanceDecs2
   where
     -- | Ensures that a provided constraint is something monad-mock can actually
     -- derive an instance for. Specifically, it must be a constraint of kind
@@ -203,8 +208,8 @@ makeAction actionNameStr classTs = do
     methodNameToConstructorName name = mkName (toUpper c : cs)
       where (c:cs) = nameBase name
 
-    mkInstance :: Type -> Type -> [Dec] -> Q Dec
-    mkInstance actionT classT methodSigs = do
+    mkInstance :: Type -> (Type -> [Pred]) -> Type -> Type -> [Dec] -> Q Dec
+    mkInstance mockT mkExtraConstraints actionT classT methodSigs = do
       mVar <- newName "m"
 
       -- In order to calculate the constraints on the instance, we need to look
@@ -229,10 +234,10 @@ makeAction actionNameStr classTs = do
           contextSubFns = map (uncurry substituteTypeVar) classBindsToInstanceBinds
           instanceContext = foldr map classContext contextSubFns
 
-      let instanceHead = classT `AppT` (ConT ''MockT `AppT` actionT `AppT` VarT mVar)
+      let instanceHead = classT `AppT` (mockT `AppT` actionT `AppT` VarT mVar)
       methodImpls <- traverse mkInstanceMethod methodSigs
 
-      return $ instanceD' instanceContext instanceHead methodImpls
+      return $ instanceD' (instanceContext ++ mkExtraConstraints (VarT mVar)) instanceHead methodImpls
 
     mkInstanceMethod :: Dec -> Q Dec
     mkInstanceMethod (SigD name typ) = do
diff --git a/test-suite/Control/Monad/Mock/StatelessSpec.hs b/test-suite/Control/Monad/Mock/StatelessSpec.hs
new file mode 100644
index 0000000..166db0a
--- /dev/null
+++ b/test-suite/Control/Monad/Mock/StatelessSpec.hs
@@ -0,0 +1,71 @@
+{-# LANGUAGE QuasiQuotes #-}
+{-# LANGUAGE TemplateHaskell #-}
+{-# LANGUAGE FunctionalDependencies #-}
+
+module Control.Monad.Mock.StatelessSpec (spec) where
+
+import Prelude hiding (readFile, writeFile)
+
+import Control.Exception (evaluate)
+import Control.Monad.Except (MonadError, runExceptT)
+import Control.Monad.ST (runST)
+import Data.Function ((&))
+import Test.Hspec
+
+import Control.Monad.Mock.Stateless
+import Control.Monad.Mock.TH
+
+class MonadError e m => MonadFileSystem e m | m -> e where
+  readFile :: FilePath -> m String
+  writeFile :: FilePath -> String -> m ()
+makeAction "FileSystemAction" [ts| MonadFileSystem String |]
+
+copyFileAndReturn :: MonadFileSystem e m => FilePath -> FilePath -> m String
+copyFileAndReturn a b = do
+  x <- readFile a
+  writeFile b x
+  return x
+
+spec :: Spec
+spec = describe "MockT" $ do
+  it "runs computations with mocked method implementations" $ do
+    let result = runST
+          $ copyFileAndReturn "foo.txt" "bar.txt"
+          & runMockT [ ReadFile "foo.txt" :-> "file contents"
+                     , WriteFile "bar.txt" "file contents" :-> () ]
+          & runExceptT
+    result `shouldBe` Right "file contents"
+
+  it "raises an exception if calls are not in the right order" $ do
+    let result = runST
+          $ copyFileAndReturn "foo.txt" "bar.txt"
+          & runMockT [ WriteFile "bar.txt" "file contents" :-> ()
+                     , ReadFile "foo.txt" :-> "file contents" ]
+          & runExceptT
+        exnMessage =
+          "runMockT: argument mismatch in readFile\n\
+          \  given: ReadFile \"foo.txt\"\n\
+          \  expected: WriteFile \"bar.txt\" \"file contents\"\n"
+    evaluate result `shouldThrow` errorCall exnMessage
+
+  it "raises an exception if calls are missing" $ do
+    let result = -- running on top of IO
+            copyFileAndReturn "foo.txt" "bar.txt"
+          & runMockT [ ReadFile "foo.txt" :-> "file contents"
+                     , WriteFile "bar.txt" "file contents" :-> ()
+                     , ReadFile "qux.txt" :-> "file contents 2" ]
+          & runExceptT
+    let exnMessage =
+          "runMockT: expected the following unexecuted actions to be run:\n\
+          \  ReadFile \"qux.txt\"\n"
+    result `shouldThrow` errorCall exnMessage
+
+  it "raises an exception if there are too many calls" $ do
+    let result = runST
+          $ copyFileAndReturn "foo.txt" "bar.txt"
+          & runMockT [ ReadFile "foo.txt" :-> "file contents" ]
+          & runExceptT
+        exnMessage =
+          "runMockT: expected end of program, called writeFile\n\
+          \  given action: WriteFile \"bar.txt\" \"file contents\"\n"
+    evaluate result `shouldThrow` errorCall exnMessage

From e226ce52f2f6def2b0fd6928e767119b6ef23d46 Mon Sep 17 00:00:00 2001
From: Pepe Iborra <pepeiborra@gmail.com>
Date: Mon, 15 Apr 2019 16:21:37 +0100
Subject: [PATCH 2/7] Add missing dependency on primitive

---
 package.yaml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/package.yaml b/package.yaml
index 69e1786..6331c98 100644
--- a/package.yaml
+++ b/package.yaml
@@ -48,6 +48,7 @@ library:
   - th-orphans
   - monad-control >= 1.0.0.0 && < 2
   - mtl
+  - primitive
   - template-haskell >= 2.10.0.0 && < 2.13
   - transformers-base
   when:

From c28909efc1d502f7b48e3e4c481361055f4638e1 Mon Sep 17 00:00:00 2001
From: Pepe Iborra <pepeiborra@gmail.com>
Date: Mon, 15 Apr 2019 16:31:50 +0100
Subject: [PATCH 3/7] Relax upper bound for template Haskell

---
 package.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/package.yaml b/package.yaml
index 6331c98..a9402ff 100644
--- a/package.yaml
+++ b/package.yaml
@@ -49,7 +49,7 @@ library:
   - monad-control >= 1.0.0.0 && < 2
   - mtl
   - primitive
-  - template-haskell >= 2.10.0.0 && < 2.13
+  - template-haskell >= 2.10.0.0 && < 2.15
   - transformers-base
   when:
   - condition: impl(ghc < 8)

From 2d4e6170516fc391e84cc4e1dfa5cd01a00ef900 Mon Sep 17 00:00:00 2001
From: Pepe Iborra <pepeiborra@gmail.com>
Date: Wed, 17 Apr 2019 15:35:39 +0100
Subject: [PATCH 4/7] Derive also PrimMonad

---
 library/Control/Monad/Mock/Stateless.hs | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/library/Control/Monad/Mock/Stateless.hs b/library/Control/Monad/Mock/Stateless.hs
index bb3ab4f..44bfcaa 100644
--- a/library/Control/Monad/Mock/Stateless.hs
+++ b/library/Control/Monad/Mock/Stateless.hs
@@ -1,5 +1,8 @@
+{-# LANGUAGE DataKinds  #-}
+{-# LANGUAGE PolyKinds  #-}
 {-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE CPP                  #-}
+{-# LANGUAGE UnboxedTuples #-}
 {-# LANGUAGE UndecidableInstances #-}
 
 -- | A version of 'MockT' with a stateless 'MonadTransControl' instance
@@ -45,7 +48,8 @@ newtype MockT_ s f m a = MockT (ReaderT (MutVar s [WithResult f]) m a)
   deriving ( Functor, Applicative, Monad, MonadTrans, MonadIO, MonadBase b
            , MonadState st, MonadCont, MonadError e, MonadWriter w
            , MonadCatch, MonadThrow, MonadMask
-           , MonadBaseControl b, MonadTransControl )
+           , MonadBaseControl b, MonadTransControl
+           , PrimMonad)
 
 instance MonadReader r m => MonadReader r (MockT_ s f m) where
   ask = lift ask

From 87a3f223e84f1717ba11f0f677e21c64d5025264 Mon Sep 17 00:00:00 2001
From: Pepe Iborra <pepeiborra@gmail.com>
Date: Fri, 19 Apr 2019 00:02:09 +0200
Subject: [PATCH 5/7] Derive also MonadFix

---
 library/Control/Monad/Mock.hs           | 3 ++-
 library/Control/Monad/Mock/Stateless.hs | 5 +++--
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/library/Control/Monad/Mock.hs b/library/Control/Monad/Mock.hs
index 8d8bf25..257a67a 100644
--- a/library/Control/Monad/Mock.hs
+++ b/library/Control/Monad/Mock.hs
@@ -82,6 +82,7 @@ import Control.Monad.Base (MonadBase)
 import Control.Monad.Catch (MonadCatch, MonadThrow, MonadMask)
 import Control.Monad.Cont (MonadCont)
 import Control.Monad.Except (MonadError)
+import Control.Monad.Fix
 import Control.Monad.IO.Class (MonadIO)
 import Control.Monad.Reader (MonadReader)
 import Control.Monad.State (StateT, MonadState(..), runStateT)
@@ -132,7 +133,7 @@ data WithResult f where
 -- f m a@, @f@ should be an 'Action', which should be a GADT that represents a
 -- reified version of typeclass method calls.
 newtype MockT f m a = MockT (StateT [WithResult f] m a)
-  deriving ( Functor, Applicative, Monad, MonadTrans, MonadIO, MonadBase b
+  deriving ( Functor, Applicative, Monad, MonadTrans, MonadIO, MonadFix, MonadBase b
            , MonadReader r, MonadCont, MonadError e, MonadWriter w
            , MonadCatch, MonadThrow, MonadMask )
 
diff --git a/library/Control/Monad/Mock/Stateless.hs b/library/Control/Monad/Mock/Stateless.hs
index 44bfcaa..06fe5d5 100644
--- a/library/Control/Monad/Mock/Stateless.hs
+++ b/library/Control/Monad/Mock/Stateless.hs
@@ -27,6 +27,7 @@ import Control.Monad.Base (MonadBase)
 import Control.Monad.Catch (MonadCatch, MonadThrow, MonadMask)
 import Control.Monad.Cont (MonadCont)
 import Control.Monad.Except (MonadError)
+import Control.Monad.Fix
 import Control.Monad.IO.Class (MonadIO)
 import Control.Monad.Primitive (PrimMonad(..))
 import Control.Monad.Reader (ReaderT(..), MonadReader(..))
@@ -45,10 +46,10 @@ type MockT f m = MockT_ (PrimState m) f m
 type Mock s f = MockT f (ST s)
 
 newtype MockT_ s f m a = MockT (ReaderT (MutVar s [WithResult f]) m a)
-  deriving ( Functor, Applicative, Monad, MonadTrans, MonadIO, MonadBase b
+  deriving ( Functor, Applicative, Monad, MonadTrans, MonadIO, MonadFix
            , MonadState st, MonadCont, MonadError e, MonadWriter w
            , MonadCatch, MonadThrow, MonadMask
-           , MonadBaseControl b, MonadTransControl
+           , MonadBase b, MonadBaseControl b, MonadTransControl
            , PrimMonad)
 
 instance MonadReader r m => MonadReader r (MockT_ s f m) where

From 52bac815962edbbf929b163d852fdf46028fecf7 Mon Sep 17 00:00:00 2001
From: Pepe Iborra <pepeiborra@gmail.com>
Date: Fri, 19 Apr 2019 09:08:59 +0200
Subject: [PATCH 6/7] Generalize WithResult to computations

---
 library/Control/Monad/Mock/Stateless.hs | 30 +++++++++++++++----------
 1 file changed, 18 insertions(+), 12 deletions(-)

diff --git a/library/Control/Monad/Mock/Stateless.hs b/library/Control/Monad/Mock/Stateless.hs
index 06fe5d5..7eb12bf 100644
--- a/library/Control/Monad/Mock/Stateless.hs
+++ b/library/Control/Monad/Mock/Stateless.hs
@@ -28,31 +28,37 @@ import Control.Monad.Catch (MonadCatch, MonadThrow, MonadMask)
 import Control.Monad.Cont (MonadCont)
 import Control.Monad.Except (MonadError)
 import Control.Monad.Fix
+import Control.Monad.Identity
 import Control.Monad.IO.Class (MonadIO)
 import Control.Monad.Primitive (PrimMonad(..))
 import Control.Monad.Reader (ReaderT(..), MonadReader(..))
 import Control.Monad.State (MonadState)
 import Control.Monad.ST (ST, runST)
 import Control.Monad.Trans (MonadTrans(..))
-import Control.Monad.Trans.Control (MonadBaseControl, MonadTransControl)
+import Control.Monad.Trans.Control
 import Control.Monad.Writer (MonadWriter)
 import Data.Primitive.MutVar (MutVar, newMutVar, readMutVar, writeMutVar)
 import Data.Type.Equality ((:~:)(..))
 
-import Control.Monad.Mock (Action(..), MonadMock(..), WithResult(..))
+import Control.Monad.Mock (Action(..), MonadMock(..))
 
-type MockT f m = MockT_ (PrimState m) f m
+type MockT f m = MockT_ (PrimState m) f m m
 
 type Mock s f = MockT f (ST s)
 
-newtype MockT_ s f m a = MockT (ReaderT (MutVar s [WithResult f]) m a)
-  deriving ( Functor, Applicative, Monad, MonadTrans, MonadIO, MonadFix
+-- | Represents both an expected call (an 'Action') and its expected result.
+data WithResult m f where
+  (:->) :: f r -> m r -> WithResult m f
+
+newtype MockT_ s f n m a = MockT (ReaderT (MutVar s [WithResult n f]) m a)
+  deriving ( Functor, Applicative, Monad, MonadIO, MonadFix
            , MonadState st, MonadCont, MonadError e, MonadWriter w
            , MonadCatch, MonadThrow, MonadMask
-           , MonadBase b, MonadBaseControl b, MonadTransControl
+           , MonadTrans, MonadTransControl
+           , MonadBase b, MonadBaseControl b
            , PrimMonad)
 
-instance MonadReader r m => MonadReader r (MockT_ s f m) where
+instance MonadReader r m => MonadReader r (MockT_ s f n m) where
   ask = lift ask
   local f (MockT act) = MockT $ do
     env <- ask
@@ -60,7 +66,7 @@ instance MonadReader r m => MonadReader r (MockT_ s f m) where
 
 runMockT :: forall f m a .
             (Action f, PrimMonad m) =>
-            [WithResult f] -> MockT f m a -> m a
+            [WithResult m f] -> MockT f m a -> m a
 runMockT actions (MockT x) = do
   ref <- newMutVar actions
   r <- runReaderT x ref
@@ -71,10 +77,10 @@ runMockT actions (MockT x) = do
       $ "runMockT: expected the following unexecuted actions to be run:\n"
       ++ unlines (map (\(action :-> _) -> "  " ++ showAction action) remainingActions)
 
-runMock :: forall f a. Action f => [WithResult f] -> (forall s. Mock s f a) -> a
-runMock actions x = runST $ runMockT actions x
+runMock :: forall f a. Action f => [WithResult Identity f] -> (forall s. Mock s f a) -> a
+runMock actions x = runST $ runMockT (map (\(a :-> b) -> a :-> return(runIdentity b)) actions) x
 
-instance (PrimMonad m, PrimState m ~ s) => MonadMock f (MockT_ s f m) where
+instance (PrimMonad m, PrimState m ~ s) => MonadMock f (MockT_ s f m m) where
   mockAction fnName action = MockT $ do
     ref <- ask
     results <- lift $ readMutVar ref
@@ -85,7 +91,7 @@ instance (PrimMonad m, PrimState m ~ s) => MonadMock f (MockT_ s f m) where
       (action' :-> r) : actions
         | Just Refl <- action `eqAction` action' -> do
             lift $ writeMutVar ref actions
-            return r
+            lift r
         | otherwise -> error'
             $ "runMockT: argument mismatch in " ++ fnName ++ "\n"
             ++ "  given: " ++ showAction action ++ "\n"

From 64ada5860e28e22720a1452d4e0f2f2e8ea3803d Mon Sep 17 00:00:00 2001
From: Pepe Iborra <pepeiborra@gmail.com>
Date: Fri, 19 Apr 2019 09:18:17 +0200
Subject: [PATCH 7/7] Extend the WithResult language with a skipping primitive

---
 library/Control/Monad/Mock/Stateless.hs | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/library/Control/Monad/Mock/Stateless.hs b/library/Control/Monad/Mock/Stateless.hs
index 7eb12bf..0b09c45 100644
--- a/library/Control/Monad/Mock/Stateless.hs
+++ b/library/Control/Monad/Mock/Stateless.hs
@@ -48,7 +48,10 @@ type Mock s f = MockT f (ST s)
 
 -- | Represents both an expected call (an 'Action') and its expected result.
 data WithResult m f where
-  (:->) :: f r -> m r -> WithResult m f
+  -- | Matches a specific command
+  (:->)     :: f r -> m r -> WithResult m f
+  -- | Skips commands as long as the predicate returns something
+  SkipWhile :: (forall r. f r -> Maybe (m r)) -> WithResult m f
 
 newtype MockT_ s f n m a = MockT (ReaderT (MutVar s [WithResult n f]) m a)
   deriving ( Functor, Applicative, Monad, MonadIO, MonadFix
@@ -81,13 +84,19 @@ runMock :: forall f a. Action f => [WithResult Identity f] -> (forall s. Mock s
 runMock actions x = runST $ runMockT (map (\(a :-> b) -> a :-> return(runIdentity b)) actions) x
 
 instance (PrimMonad m, PrimState m ~ s) => MonadMock f (MockT_ s f m m) where
-  mockAction fnName action = MockT $ do
-    ref <- ask
+  mockAction fnName action = do
+    ref <- MockT ask
     results <- lift $ readMutVar ref
     case results of
       [] -> error'
         $ "runMockT: expected end of program, called " ++ fnName ++ "\n"
         ++ "  given action: " ++ showAction action ++ "\n"
+      SkipWhile f : actions
+        | Just res <- f action
+        -> lift res
+        | otherwise -> do
+            lift $ writeMutVar ref actions
+            mockAction fnName action
       (action' :-> r) : actions
         | Just Refl <- action `eqAction` action' -> do
             lift $ writeMutVar ref actions