Adds mock return typing

6 files changed, 229 insertions, 119 deletions

diff --git a/TODO.norg b/TODO.norg
index 1ed3db6..18edf65 100644
--- a/TODO.norg
+++ b/TODO.norg
@@ -1,13 +1,13 @@
 * Current
-  - ( ) Optimize speed of ocr
-  --- Load incrementally?
   - ( ) Add hjkl for search mode
   - ( ) Middle click
+  - (-) Optimize speed of ocr
 
 * Later
+  - ( ) Sort ocr match result by position on screen (top to bottom, left to right)
+  - ( ) Look into making mouse controls (click/mouse down/mouse up) cross-platform
   - ( ) Look into making controls cross-platform
   - ( ) Lens-ey setup for Mode access
-  - ( ) Switch to [test-fixture]{https://hackage.haskell.org/package/test-fixture}?
 
 * Maybe
   - ( ) Scroll
diff --git a/chelleport.cabal b/chelleport.cabal
index de6008e..e06cbc6 100644
--- a/chelleport.cabal
+++ b/chelleport.cabal
@@ -26,20 +26,23 @@ common common-config
     NumericUnderscores
     OverloadedStrings
     QuasiQuotes
+    RankNTypes
+    StandaloneDeriving
     TemplateHaskell
     TupleSections
     UndecidableInstances
   default-language: Haskell2010
   build-depends:
+    array,
     base,
-    text,
-    time,
+    containers,
+    data-default,
+    directory,
     mtl == 2.3.1,
     sdl2 == 2.5.5.0,
-    array,
     temporary,
-    directory,
-    containers
+    text,
+    time
 
 common warnings
   ghc-options:
diff --git a/specs/Mock.hs b/specs/Mock.hs
index 00cb64a..4d7f3f2 100644
--- a/specs/Mock.hs
+++ b/specs/Mock.hs
@@ -1,7 +1,28 @@
 module Mock where
 
 import Control.Monad (join)
+import Control.Monad.IO.Class (MonadIO)
+import Control.Monad.State (MonadState (state), StateT)
+import Data.Default (Default (def))
+import Data.Typeable (Typeable)
 import Language.Haskell.TH
+import Test.Hspec
+
+callTypeName :: Name
+callTypeName = mkName "Call"
+
+generateMock :: [Name] -> Q [Dec]
+generateMock typeClassNames = do
+  functions <- join <$> mapM typeclassFunctions typeClassNames
+  let deriveClause = DerivClause Nothing []
+  let callDataDeclr =
+        DataD [] callTypeName [PlainTV (mkName "ret") ()] Nothing (toGadtCtor <$> functions) [deriveClause]
+  instances <- join <$> mapM createInstanceOnType typeClassNames
+  extras <- staticExtrasForMockCalls
+  pure $ [callDataDeclr] ++ extras ++ instances
+
+toMockCtorName :: Name -> Name
+toMockCtorName = mkName . ("Mock_" ++) . nameBase
 
 typeclassFunctions :: Name -> Q [(Name, Type)]
 typeclassFunctions cls = do
@@ -13,21 +34,100 @@ typeclassFunctions cls = do
     functionName (SigD name typ) = [(name, typ)]
     functionName _ = []
 
-generateMock :: [Name] -> Q [Dec]
-generateMock clss = do
-  functions <- join <$> mapM typeclassFunctions clss
-  pure [DataD [] (mkName "Call") [] Nothing (toCtor <$> functions) [deriveClause]]
+countArgsInType :: Type -> Int
+countArgsInType (AppT (AppT ArrowT _) rest) = 1 + countArgsInType rest
+countArgsInType _ = 0
+
+toGadtCtor :: (Name, Type) -> Con
+toGadtCtor (name, typ) = GadtC [toMockCtorName name] args retType
   where
-    deriveClause = DerivClause Nothing (ConT <$> [''Show, ''Eq])
+    (args, retType) = extractArgsAndReturnTypes typ
 
-    toCtor :: (Name, Type) -> Con
-    toCtor (name, typ) = NormalC (toCtorName name) ((noBang,) <$> getArgs typ)
+extractArgsAndReturnTypes :: Type -> ([BangType], Type)
+extractArgsAndReturnTypes (AppT (AppT ArrowT typ) rest) = ((noBang, typ) : args, ret)
+  where
+    (args, ret) = extractArgsAndReturnTypes rest
+    noBang = Bang NoSourceUnpackedness NoSourceStrictness
+extractArgsAndReturnTypes (AppT _ ret) = ([], AppT (ConT callTypeName) ret)
+extractArgsAndReturnTypes ret = ([], AppT (ConT callTypeName) ret)
 
-    getArgs :: Type -> [Type]
-    getArgs (AppT (AppT ArrowT typ) rest) = typ : getArgs rest
-    getArgs _ = []
+createInstanceOnType :: Name -> Q [Dec]
+createInstanceOnType name =
+  typeclassFunctions name >>= createInstance name
 
-    noBang = Bang NoSourceUnpackedness NoSourceStrictness
+createInstance :: Name -> [(Name, Type)] -> Q [Dec]
+createInstance name funcs = do
+  let ctx = AppT (ConT ''MonadIO) (VarT $ mkName "m")
+  let typ = AppT (ConT name) (ConT (mkName "TestM") `AppT` VarT (mkName "a") `AppT` VarT (mkName "m"))
+  funcDeclrs <- mapM (\(n, t) -> toInstanceMethodDef n $ countArgsInType t) funcs
+  let inst = InstanceD Nothing [ctx] typ funcDeclrs
+  pure [inst]
+
+toInstanceMethodDef :: Name -> Int -> Q Dec
+toInstanceMethodDef name argCount = do
+  let argNames = mkName . ("a" ++) . show <$> [1 .. argCount]
+  let callExp = foldl AppE (ConE $ toMockCtorName name) (VarE <$> argNames)
+  bodyExp <- instanceMethod callExp
+  pure $ FunD name [Clause (VarP <$> argNames) (NormalB bodyExp) []]
+
+instanceMethod :: Exp -> Q Exp
+instanceMethod mockExp = do
+  [e|
+    do
+      declarations <- gets (fmap unsafeUnpackDeclr . mockDeclarations)
+      let call = $(pure mockExp)
+      getMockValue declarations call <$ registerMockCall call
+    |]
+
+staticExtrasForMockCalls :: Q [Dec]
+staticExtrasForMockCalls =
+  [d|
+    deriving instance Show ($(conT callTypeName) a)
+
+    deriving instance Eq ($(conT callTypeName) a)
+
+    data CallWrapper where
+      CallWrapper :: (Typeable a, Show a, Eq a, Eq ($(conT callTypeName) a)) => $(conT callTypeName) a -> CallWrapper
+
+    data CallMockDeclaration where
+      CallMockDeclaration :: (Typeable a, Show a, Eq a, Eq ($(conT callTypeName) a)) => $(conT callTypeName) a -> a -> CallMockDeclaration
+
+    deriving instance Show CallWrapper
+
+    -- deriving instance Eq CallWrapper
+    instance Eq CallWrapper where
+      (CallWrapper a) == (CallWrapper b) =
+        case cast a of
+          Just a' -> a' == b
+          Nothing -> False
+
+    data MockCalls = MockCalls {calls :: [CallWrapper], mockDeclarations :: [CallMockDeclaration]}
+
+    newtype TestM ret m a = TestM {runTestM :: StateT MockCalls m a}
+      deriving (Functor, Applicative, Monad, MonadIO, MonadState MockCalls)
+
+    runTestMWithMocks :: (MonadIO m) => TestM x m a -> m (a, MockCalls)
+    runTestMWithMocks action = runStateT (runTestM action) (MockCalls [] [])
+
+    registerMockCall :: (MonadState MockCalls m, Typeable a, Show a, Eq a) => $(conT callTypeName) a -> m ()
+    registerMockCall call =
+      void $ state (\mock -> ((), mock {calls = calls mock ++ [CallWrapper call]}))
+
+    getMockValue :: (Typeable a, Default a, Show a, Eq a) => [($(conT callTypeName) a, a)] -> $(conT callTypeName) a -> a
+    getMockValue [] _ = def
+    getMockValue ((fn, ret) : _) call | call == fn = ret
+    getMockValue (_ : rest) call = getMockValue rest call
+
+    unsafeUnpackDeclr :: CallMockDeclaration -> ($(conT callTypeName) a, a)
+    unsafeUnpackDeclr (CallMockDeclaration f r) = unsafeCoerce (f, r)
+
+    mockReturns :: (MonadState MockCalls m, Typeable a, Show a, Eq a) => $(conT callTypeName) a -> a -> m ()
+    mockReturns call ret =
+      state (\mock -> ((), mock {mockDeclarations = mockDeclarations mock ++ [CallMockDeclaration call ret]}))
+
+    shouldHaveCalled :: (HasCallStack, Typeable a, Show a, Eq a) => MockCalls -> $(conT callTypeName) a -> Expectation
+    shouldHaveCalled mock call = calls mock `shouldContain` [CallWrapper call]
 
-    toCtorName :: Name -> Name
-    toCtorName = mkName . ("Mock_" ++) . nameBase
+    shouldContainCalls :: (HasCallStack, Typeable a, Show a, Eq a) => MockCalls -> [$(conT callTypeName) a] -> Expectation
+    shouldContainCalls mock ls = calls mock `shouldContain` (CallWrapper <$> ls)
+    |]
diff --git a/specs/Specs/AppStateSpec.hs b/specs/Specs/AppStateSpec.hs
index 7eeb082..245ee9c 100644
--- a/specs/Specs/AppStateSpec.hs
+++ b/specs/Specs/AppStateSpec.hs
@@ -35,7 +35,7 @@ test = do
 
       it "hides window, triggers mouse click and shows the window again" $ do
         (_, mock) <- runWithMocks $ update currentState $ ChainMouseClick LeftClick
-        calls mock `shouldBe` [Mock_hideWindow, Mock_clickMouseButton LeftClick, Mock_showWindow]
+        mock `shouldContainCalls` [Mock_hideWindow, Mock_clickMouseButton LeftClick, Mock_showWindow]
 
       it "continues with action ResetKeys without updating state" $ do
         ((nextState, action), _) <- runWithMocks $ update currentState $ ChainMouseClick LeftClick
@@ -51,23 +51,23 @@ test = do
 
         it "clicks multiple times" $ do
           (_, mock) <- runWithMocks $ update currentState $ ChainMouseClick LeftClick
-          calls mock
-            `shouldBe` [ Mock_hideWindow,
-                         Mock_clickMouseButton LeftClick,
-                         Mock_clickMouseButton LeftClick,
-                         Mock_clickMouseButton LeftClick,
-                         Mock_showWindow
-                       ]
+          mock
+            `shouldContainCalls` [ Mock_hideWindow,
+                                   Mock_clickMouseButton LeftClick,
+                                   Mock_clickMouseButton LeftClick,
+                                   Mock_clickMouseButton LeftClick,
+                                   Mock_showWindow
+                                 ]
 
       context "when repetition is 0" $ do
         let currentState = defaultState {stateRepetition = 0}
 
         it "clicks just once" $ do
           (_, mock) <- runWithMocks $ update currentState $ ChainMouseClick LeftClick
-          calls mock `shouldBe` [Mock_hideWindow, Mock_clickMouseButton LeftClick, Mock_showWindow]
+          mock `shouldContainCalls` [Mock_hideWindow, Mock_clickMouseButton LeftClick, Mock_showWindow]
 
     context "with action IncrementHighlightIndex" $ do
-      let currentState = defaultState
+      -- let currentState = defaultState
 
       it "todo: implement" $ do
         1 `shouldBe` 1
@@ -76,7 +76,9 @@ test = do
       let currentState = defaultState
 
       it "continues with MoveMousePosition" $ do
-        ((_, action), _) <- runWithMocks $ update currentState $ IncrementMouseCursor (10, -5)
+        ((_, action), _) <- runWithMocks $ do
+          Mock_getMousePointerPosition `mockReturns` (42, 42)
+          update currentState $ IncrementMouseCursor (10, -5)
         action `shouldBe` Just (MoveMousePosition (52, 37))
 
       it "does update state" $ do
@@ -87,14 +89,18 @@ test = do
         let currentState = defaultState {stateRepetition = 5}
 
         it "multiplies increment" $ do
-          ((_, action), _) <- runWithMocks $ update currentState $ IncrementMouseCursor (10, -5)
+          ((_, action), _) <- runWithMocks $ do
+            Mock_getMousePointerPosition `mockReturns` (42, 42)
+            update currentState $ IncrementMouseCursor (10, -5)
           action `shouldBe` Just (MoveMousePosition (92, 17))
 
       context "when repetition is 0" $ do
         let currentState = defaultState {stateRepetition = 0}
 
         it "increments just once" $ do
-          ((_, action), _) <- runWithMocks $ update currentState $ IncrementMouseCursor (10, -5)
+          ((_, action), _) <- runWithMocks $ do
+            Mock_getMousePointerPosition `mockReturns` (42, 42)
+            update currentState $ IncrementMouseCursor (10, -5)
           action `shouldBe` Just (MoveMousePosition (52, 37))
 
     context "with action MouseDragEnd" $ do
@@ -102,7 +108,7 @@ test = do
 
       it "hides window, stops dragging and shows the window again" $ do
         (_, mock) <- runWithMocks $ update currentState MouseDragEnd
-        calls mock `shouldContain` [Mock_hideWindow, Mock_releaseMouseButton, Mock_showWindow]
+        mock `shouldContainCalls` [Mock_hideWindow, Mock_releaseMouseButton, Mock_showWindow]
 
       it "does not continue or update state" $ do
         (result, _) <- runWithMocks $ update currentState MouseDragStart
@@ -113,7 +119,7 @@ test = do
 
       it "hides window, starts dragging and shows the window again" $ do
         (_, mock) <- runWithMocks $ update currentState MouseDragStart
-        calls mock `shouldContain` [Mock_hideWindow, Mock_pressMouseButton, Mock_showWindow]
+        mock `shouldContainCalls` [Mock_hideWindow, Mock_pressMouseButton, Mock_showWindow]
 
       it "does not continue or update state" $ do
         (result, _) <- runWithMocks $ update currentState MouseDragStart
@@ -172,7 +178,10 @@ test = do
               nextState `shouldBe` currentState {stateKeySequence = "DEF", stateIsMatched = True}
 
             it "continues with MoveMousePosition action at center of matched cell" $ do
-              ((_, action), _) <- runWithMocks $ update currentState $ HandleKeyInput SDL.KeycodeF
+              ((_, action), _) <- runWithMocks $ do
+                Mock_windowSize `mockReturns` mockWindowSize
+                Mock_windowPosition `mockReturns` mockWindowPosition
+                update currentState $ HandleKeyInput SDL.KeycodeF
               action `shouldBe` Just (MoveMousePosition (1640, 370))
 
     context "with action MoveMouseInDirection" $ do
@@ -180,19 +189,27 @@ test = do
 
       context "when direction is up" $ do
         it "continues to increment movement" $ do
-          ((_, action), _) <- runWithMocks $ update currentState $ MoveMouseInDirection DirUp
+          ((_, action), _) <- runWithMocks $ do
+            Mock_windowSize `mockReturns` mockWindowSize
+            update currentState $ MoveMouseInDirection DirUp
           action `shouldBe` Just (IncrementMouseCursor (0, -33))
       context "when direction is down" $ do
         it "continues to increment movement" $ do
-          ((_, action), _) <- runWithMocks $ update currentState $ MoveMouseInDirection DirDown
+          ((_, action), _) <- runWithMocks $ do
+            Mock_windowSize `mockReturns` mockWindowSize
+            update currentState $ MoveMouseInDirection DirDown
           action `shouldBe` Just (IncrementMouseCursor (0, 33))
       context "when direction is left" $ do
         it "continues to increment movement" $ do
-          ((_, action), _) <- runWithMocks $ update currentState $ MoveMouseInDirection DirLeft
+          ((_, action), _) <- runWithMocks $ do
+            Mock_windowSize `mockReturns` mockWindowSize
+            update currentState $ MoveMouseInDirection DirLeft
           action `shouldBe` Just (IncrementMouseCursor (-60, 0))
       context "when direction is right" $ do
         it "continues to increment movement" $ do
-          ((_, action), _) <- runWithMocks $ update currentState $ MoveMouseInDirection DirRight
+          ((_, action), _) <- runWithMocks $ do
+            Mock_windowSize `mockReturns` mockWindowSize
+            update currentState $ MoveMouseInDirection DirRight
           action `shouldBe` Just (IncrementMouseCursor (60, 0))
 
     context "with action MoveMousePosition" $ do
@@ -225,13 +242,28 @@ test = do
 
       context "when mode is ModeSearch" $ do
         it "captures screenshot for word search" $ do
-          ((_, _), mock) <- runWithMocks $ update currentState $ SetMode defaultSearchMode
+          ((_, _), mock) <- runWithMocks $ do
+            Mock_windowSize `mockReturns` mockWindowSize
+            Mock_windowPosition `mockReturns` mockWindowPosition
+            update currentState $ SetMode defaultSearchMode
           mock `shouldHaveCalled` Mock_captureScreenshot (mockWindowOffsetX, mockWindowOffsetY) (mockWindowWidth, mockWindowHeight)
 
         it "updates mode in state with ocr words" $ do
-          ((nextState, _), _) <- runWithMocks $ update currentState $ SetMode defaultSearchMode
           let matchWord = OCRMatch {matchStartX = 40, matchStartY = 5, matchEndX = 100, matchEndY = 20, matchText = "Wow"}
-          nextState `shouldBe` currentState {stateMode = defaultSearchMode {searchWords = [matchWord], searchFilteredWords = [matchWord]}}
+          ((nextState, _), _) <- runWithMocks $ do
+            Mock_windowSize `mockReturns` mockWindowSize
+            Mock_windowPosition `mockReturns` mockWindowPosition
+            Mock_captureScreenshot mockWindowPosition mockWindowSize `mockReturns` "mock-filename"
+            Mock_getWordsInImage "mock-filename" `mockReturns` [matchWord]
+            update currentState $ SetMode defaultSearchMode
+          nextState
+            `shouldBe` currentState
+              { stateMode =
+                  defaultSearchMode
+                    { searchWords = [matchWord],
+                      searchFilteredWords = [matchWord]
+                    }
+              }
 
     context "with action ShutdownApp" $ do
       let currentState = defaultState
@@ -249,7 +281,7 @@ test = do
 
       it "hides window and triggers mouse click" $ do
         (_, mock) <- runWithMocks $ update currentState $ TriggerMouseClick LeftClick
-        calls mock `shouldContain` [Mock_hideWindow, Mock_clickMouseButton LeftClick]
+        mock `shouldContainCalls` [Mock_hideWindow, Mock_clickMouseButton LeftClick]
 
       it "continues with action ShutdownApp without updating state" $ do
         ((nextState, action), _) <- runWithMocks $ update currentState $ TriggerMouseClick LeftClick
@@ -265,19 +297,19 @@ test = do
 
         it "clicks multiple times" $ do
           (_, mock) <- runWithMocks $ update currentState $ TriggerMouseClick LeftClick
-          calls mock
-            `shouldBe` [ Mock_hideWindow,
-                         Mock_clickMouseButton LeftClick,
-                         Mock_clickMouseButton LeftClick,
-                         Mock_clickMouseButton LeftClick
-                       ]
+          mock
+            `shouldContainCalls` [ Mock_hideWindow,
+                                   Mock_clickMouseButton LeftClick,
+                                   Mock_clickMouseButton LeftClick,
+                                   Mock_clickMouseButton LeftClick
+                                 ]
 
       context "when repetition is 0" $ do
         let currentState = defaultState {stateRepetition = 0}
 
         it "clicks just once" $ do
           (_, mock) <- runWithMocks $ update currentState $ TriggerMouseClick LeftClick
-          calls mock `shouldBe` [Mock_hideWindow, Mock_clickMouseButton LeftClick]
+          mock `shouldContainCalls` [Mock_hideWindow, Mock_clickMouseButton LeftClick]
 
     context "with action UpdateRepetition" $ do
       let currentState = defaultState
diff --git a/specs/Specs/ViewSpec.hs b/specs/Specs/ViewSpec.hs
index 1e34811..79c8173 100644
--- a/specs/Specs/ViewSpec.hs
+++ b/specs/Specs/ViewSpec.hs
@@ -9,46 +9,52 @@ import TestUtils
 test :: SpecWith ()
 test = do
   let defaultState = defaultAppState {stateGrid = [["ABC", "DEF"], ["DJK", "JKL"]]}
-  let drawTextCalls = filter (\case Mock_drawText {} -> True; _ -> False) . calls
 
   describe "#render" $ do
     context "when key sequence is empty" $ do
       let currentState = defaultState {stateKeySequence = ""}
 
       it "draws matching text labels" $ do
-        (_, mock) <- runWithMocks $ render currentState
-        drawTextCalls mock
-          `shouldBe` [ Mock_drawText (460, 10) colorWhite FontLG "ABC",
-                       Mock_drawText (1420, 10) colorWhite FontLG "DEF",
-                       Mock_drawText (460, 550) colorWhite FontLG "DJK",
-                       Mock_drawText (1420, 550) colorWhite FontLG "JKL"
-                     ]
+        (_, mock) <- runWithMocks $ do
+          Mock_windowSize `mockReturns` mockWindowSize
+          render currentState
+        mock `shouldHaveCalled` Mock_drawText (460, 10) colorWhite FontLG "ABC"
+        mock `shouldHaveCalled` Mock_drawText (1420, 10) colorWhite FontLG "DEF"
+        mock `shouldHaveCalled` Mock_drawText (460, 550) colorWhite FontLG "DJK"
+        mock `shouldHaveCalled` Mock_drawText (1420, 550) colorWhite FontLG "JKL"
 
     context "when there is a partial match" $ do
       let currentState = defaultState {stateKeySequence = "D"}
 
       it "draws matching text labels" $ do
-        (_, mock) <- runWithMocks $ render currentState
-        drawTextCalls mock
-          `shouldBe` [ Mock_drawText (1420, 10) colorLightGray FontLG "D",
-                       Mock_drawText (1430, 10) colorAccent FontLG "EF",
-                       Mock_drawText (460, 550) colorLightGray FontLG "D",
-                       Mock_drawText (470, 550) colorAccent FontLG "JK"
-                     ]
+        (_, mock) <- runWithMocks $ do
+          Mock_windowSize `mockReturns` mockWindowSize
+          Mock_drawText (1420, 10) colorLightGray FontLG "D" `mockReturns` (10, 0)
+          Mock_drawText (460, 550) colorLightGray FontLG "D" `mockReturns` (10, 0)
+          render currentState
+        mock `shouldHaveCalled` Mock_drawText (1420, 10) colorLightGray FontLG "D"
+        mock `shouldHaveCalled` Mock_drawText (1430, 10) colorAccent FontLG "EF"
+        mock `shouldHaveCalled` Mock_drawText (460, 550) colorLightGray FontLG "D"
+        mock `shouldHaveCalled` Mock_drawText (470, 550) colorAccent FontLG "JK"
 
     context "when key sequence is complete match" $ do
       let currentState = defaultState {stateKeySequence = "DEF"}
 
       it "draws only the matching label" $ do
-        (_, mock) <- runWithMocks $ render currentState
-        drawTextCalls mock `shouldBe` [Mock_drawText (1420, 10) colorLightGray FontLG "DEF"]
+        (_, mock) <- runWithMocks $ do
+          Mock_windowSize `mockReturns` mockWindowSize
+          render currentState
+        mock `shouldHaveCalled` Mock_drawText (1420, 10) colorLightGray FontLG "DEF"
 
   describe "#renderKeySequence" $ do
     context "when there is a partial match" $ do
       it "draws the matched section and highlights the remaining characters" $ do
-        (_, mock) <- runWithMocks $ renderKeySequence "ABC" "ABCDE" (0, 0)
-        calls mock
-          `shouldBe` [Mock_drawText (0, 0) colorLightGray FontLG "ABC", Mock_drawText (3 * 10, 0) colorAccent FontLG "DE"]
+        (_, mock) <- runWithMocks $ do
+          Mock_windowSize `mockReturns` mockWindowSize
+          Mock_drawText (0, 0) colorLightGray FontLG "ABC" `mockReturns` (30, 0)
+          renderKeySequence "ABC" "ABCDE" (0, 0)
+        mock `shouldHaveCalled` Mock_drawText (0, 0) colorLightGray FontLG "ABC"
+        mock `shouldHaveCalled` Mock_drawText (30, 0) colorAccent FontLG "DE"
 
       it "return true as the text is visible" $ do
         (isVisible, _) <- runWithMocks $ renderKeySequence "ABC" "ABCDE" (0, 0)
@@ -57,7 +63,7 @@ test = do
     context "when there is no input key sequence" $ do
       it "draws text as a single chunk" $ do
         (_, mock) <- runWithMocks $ renderKeySequence "" "ABCD" (0, 0)
-        calls mock `shouldBe` [Mock_drawText (0, 0) colorWhite FontLG "ABCD"]
+        mock `shouldHaveCalled` Mock_drawText (0, 0) colorWhite FontLG "ABCD"
 
       it "return true as the text is visible" $ do
         (isVisible, _) <- runWithMocks $ renderKeySequence "" "ABCD" (0, 0)
diff --git a/specs/TestUtils.hs b/specs/TestUtils.hs
index 59a59c7..d6f3cd2 100644
--- a/specs/TestUtils.hs
+++ b/specs/TestUtils.hs
@@ -4,71 +4,40 @@ import Chelleport.AppShell (MonadAppShell (..))
 import Chelleport.Control (MonadControl (..))
 import Chelleport.Draw (MonadDraw (..))
 import Chelleport.OCR (MonadOCR (..))
-import Chelleport.Types
 import Control.Monad (void)
 import Control.Monad.IO.Class (MonadIO)
-import Control.Monad.State (MonadState (state), StateT (runStateT))
-import qualified Data.Text as Text
+import Control.Monad.State (StateT (runStateT), gets)
+import Data.Typeable (cast)
 import Foreign.C (CInt)
 import Mock
-import Test.Hspec
+import Unsafe.Coerce (unsafeCoerce)
 
 $(generateMock [''MonadDraw, ''MonadControl, ''MonadAppShell, ''MonadOCR])
 
-newtype MockCalls = MockCalls {calls :: [Call]}
-  deriving (Show)
-
-shouldHaveCalled :: (HasCallStack) => MockCalls -> Call -> Expectation
-shouldHaveCalled mock call = calls mock `shouldContain` [call]
-
-runWithMocks :: (MonadIO m) => TestM m a -> m (a, MockCalls)
-runWithMocks action = runStateT (runTestM action) (MockCalls [])
-
-newtype TestM m a = TestM {runTestM :: StateT MockCalls m a}
-  deriving (Functor, Applicative, Monad, MonadIO, MonadState MockCalls)
-
 mockWindowWidth :: CInt
 mockWindowWidth = 1920
 
 mockWindowHeight :: CInt
 mockWindowHeight = 1080
 
+mockWindowSize :: (CInt, CInt)
+mockWindowSize = (mockWindowWidth, mockWindowHeight)
+
 mockWindowOffsetX :: CInt
 mockWindowOffsetX = 200
 
 mockWindowOffsetY :: CInt
 mockWindowOffsetY = 100
 
+mockWindowPosition :: (CInt, CInt)
+mockWindowPosition = (mockWindowOffsetX, mockWindowOffsetY)
+
 mockTextWidth :: Int
 mockTextWidth = 10
 
-registerMockCall :: (MonadState MockCalls m) => Call -> m ()
-registerMockCall call =
-  void $ state (\mock -> ((), MockCalls {calls = calls mock ++ [call]}))
-
-instance (MonadIO m) => MonadControl (TestM m) where
-  clickMouseButton btn = registerMockCall $ Mock_clickMouseButton btn
-  moveMousePointer x y = registerMockCall $ Mock_moveMousePointer x y
-  getMousePointerPosition = (42, 42) <$ registerMockCall Mock_getMousePointerPosition
-  pressMouseButton = registerMockCall Mock_pressMouseButton
-  releaseMouseButton = registerMockCall Mock_releaseMouseButton
-
-instance (MonadIO m) => MonadDraw (TestM m) where
-  drawLine p1 p2 = registerMockCall $ Mock_drawLine p1 p2
-  fillRect p size = registerMockCall $ Mock_fillRect p size
-  drawText p color size text = (fromIntegral $ mockTextWidth * Text.length text, 0) <$ registerMockCall (Mock_drawText p color size text)
-  drawCircle radius p = registerMockCall $ Mock_drawCircle radius p
-  setDrawColor color = registerMockCall $ Mock_setDrawColor color
-  windowSize = (mockWindowWidth, mockWindowHeight) <$ registerMockCall Mock_windowSize
-  windowPosition = (mockWindowOffsetX, mockWindowOffsetY) <$ registerMockCall Mock_windowPosition
-
-instance (MonadIO m) => MonadAppShell (TestM m) where
-  hideWindow = registerMockCall Mock_hideWindow
-  showWindow = registerMockCall Mock_showWindow
-  shutdownApp = registerMockCall Mock_shutdownApp
-
-instance (MonadIO m) => MonadOCR (TestM m) where
-  captureScreenshot p size = "" <$ registerMockCall (Mock_captureScreenshot p size)
-  getWordsInImage filePath = [match] <$ registerMockCall (Mock_getWordsInImage filePath)
-    where
-      match = OCRMatch {matchStartX = 40, matchStartY = 5, matchEndX = 100, matchEndY = 20, matchText = "Wow"}
+runWithMocks :: (MonadIO m) => TestM x m a -> m (a, MockCalls)
+runWithMocks act = runTestMWithMocks $ do
+  -- Default mocks
+  Mock_windowSize `mockReturns` mockWindowSize
+  Mock_windowPosition `mockReturns` mockWindowPosition
+  act