diff --git a/ChangeLog.md b/ChangeLog.md
--- a/ChangeLog.md
+++ b/ChangeLog.md
@@ -1,2 +1,4 @@
 # Changelog for wait-on-port
 - 0.0.0.1 First version
+
+- 0.2.0.0 Rework internals and change `EventHandlers` to `EventHandlers m`. Fix bug if host could not be found. It would loop forever and now throws.
diff --git a/port-utils.cabal b/port-utils.cabal
--- a/port-utils.cabal
+++ b/port-utils.cabal
@@ -2,10 +2,10 @@
 --
 -- see: https://github.com/sol/hpack
 --
--- hash: 1d9819c36e0182c73f52e502f0bdd4d406aa01dd948ec440d6852ba1e7bd482b
+-- hash: 86eb2c8b9fe7b1d2f8dcce2f6eae318261062f193924cf7fe0119f3191c4088b
 
 name:           port-utils
-version:        0.1.0.4
+version:        0.2.0.0
 synopsis:       Utilities for creating and waiting on ports
 description:    Utilities for creating and waiting on ports.
  .
@@ -13,7 +13,6 @@
  .
  @wait@ will attempt to connect to given host and port repeatedly until successful.
  .
-
 homepage:       https://github.com/jfischoff/port-utils#readme
 bug-reports:    https://github.com/jfischoff/port-utils/issues
 author:         Jonathan Fischoff
@@ -35,11 +34,12 @@
   exposed-modules:
       Network.Socket.Free
       Network.Socket.Wait
+      Network.Socket.Wait.Internal
   other-modules:
       Paths_port_utils
   hs-source-dirs:
       src
-  default-extensions: ScopedTypeVariables RecordWildCards LambdaCase
+  default-extensions: ScopedTypeVariables RecordWildCards LambdaCase UndecidableInstances TypeSynonymInstances FlexibleInstances
   build-depends:
       base >=4.7 && <5
     , network
@@ -53,7 +53,7 @@
       Paths_port_utils
   hs-source-dirs:
       test
-  default-extensions: ScopedTypeVariables RecordWildCards LambdaCase
+  default-extensions: ScopedTypeVariables RecordWildCards LambdaCase UndecidableInstances TypeSynonymInstances FlexibleInstances
   ghc-options: -threaded -rtsopts -with-rtsopts=-N
   build-depends:
       async
@@ -62,4 +62,5 @@
     , network
     , port-utils
     , stm
+    , transformers
   default-language: Haskell2010
diff --git a/src/Network/Socket/Wait.hs b/src/Network/Socket/Wait.hs
--- a/src/Network/Socket/Wait.hs
+++ b/src/Network/Socket/Wait.hs
@@ -6,117 +6,4 @@
   , EventHandlers (..)
   , defaultDelay
   ) where
-import qualified Network.Socket as N
-import qualified Control.Concurrent as C
-import qualified Control.Exception as E
-
--------------------------------------------------------------------------------
--- Simple Api
--------------------------------------------------------------------------------
-
--- | 'wait' will attempt to connect to the given host and port repeated every
--- 10 milliseconds until it is successful. It will throw an 'IOError' if the
--- host cannot be resolved.
---
--- A typical use case is to call 'wait' in test code to wait for a server to
--- start before trying to connect. For example:
---
--- @
---    void $ forkIO $ Warp.run 7000 app
---    -- Wait for the server to start listening on the socket
---    wait "127.0.0.1" 7000
---    -- Communicate with the server
--- @
---
--- If you would like to control the delay or understand how many connection
--- attempts were made use 'waitWith'.
---
--- Since 0.0.0.1
-wait :: String
-     -- ^ Host
-     -> Int
-     -- ^ Port
-     -> IO ()
-wait = waitWith mempty defaultDelay
-
--------------------------------------------------------------------------------
--- Advanced Api
--------------------------------------------------------------------------------
--- | The default delay between retries is 10000 microseconds (10 ms)
-defaultDelay :: Int
-defaultDelay = 10000
-
--- | The 'EventHandlers' is a record of 'IO' actions that are called when
---   interesting events occur in the lifecycle of the 'waitWith' loop.
---   One can pass in custom 'EventHandlers' values to implement logging
---   and other forms of instrumentation.
---
---   For example for debug one could print out each step using:
---
--- @
---    printHandlers = EventHandlers
---      { createdSocket = putStrLn "createdSocket"
---      , delaying      = putStrLn "delaying"
---      , restarting    = putStrLn "restarting"
---      }
--- @
---
-data EventHandlers = EventHandlers
-  { createdSocket :: IO ()
-  -- ^ Called after the socket is created
-  , delaying      :: IO ()
-  -- ^ Called after a failed attempt to connect before the thread
-  -- is put to sleep.
-  , restarting    :: IO ()
-  -- ^ Called before a recursive call to restart the connection attempt
-  }
-
-instance Semigroup EventHandlers where
-  x <> y = EventHandlers
-    { createdSocket = createdSocket x <> createdSocket y
-    , delaying      = delaying      x <> delaying      y
-    , restarting    = restarting    x <> restarting    y
-    }
-
-instance Monoid EventHandlers where
-  mempty  = EventHandlers mempty mempty mempty
-  mappend = (<>)
-
--- | Advanced usage. In most situations calling 'wait' will suffice. This allows
--- one to customize the delay between retries and debug the behavior of the
--- function. 'wait' is defined as
---
--- @
---    wait = waitWith mempty defaultDelay
--- @
---
--- Since 0.0.0.1
-waitWith :: EventHandlers
-         -- ^ A record of IO actions that are called during sp
-         -> Int
-         -- ^ Microseconds to delay
-         -> String
-         -- ^ Host
-         -> Int
-         -- ^ Port
-         -> IO ()
-waitWith eh@EventHandlers {..} delay host port = do
-  res <- E.try $ do
-    let hints = N.defaultHints { N.addrSocketType = N.Stream }
-    -- getAddrInfo returns a non-empty array or throws per the doc
-    addr:_ <- N.getAddrInfo (Just hints) (Just host) (Just $ show port)
-    E.bracket
-      (N.socket (N.addrFamily addr) (N.addrSocketType addr) (N.addrProtocol addr))
-      N.close
-      $ \sock -> do
-        createdSocket
-        N.connect sock $ N.addrAddress addr
-
-  case res of
-    Left (_ :: E.IOException) -> do
-      delaying
-      C.threadDelay delay
-
-      restarting
-      waitWith eh delay host port
-    Right _ -> pure ()
+import Network.Socket.Wait.Internal
diff --git a/src/Network/Socket/Wait/Internal.hs b/src/Network/Socket/Wait/Internal.hs
new file mode 100644
--- /dev/null
+++ b/src/Network/Socket/Wait/Internal.hs
@@ -0,0 +1,141 @@
+{-|
+  This internal modules exposes additional functions for testing.
+  Use 'Network.Socket.Wait' instead
+|-}
+module Network.Socket.Wait.Internal
+  ( -- * Simple Wait Api
+    wait
+    -- * Advanced Wait Api
+  , waitWith
+  , EventHandlers (..)
+  , defaultDelay
+  , connectAction
+  , waitM
+  ) where
+import qualified Network.Socket as N
+import qualified Control.Concurrent as C
+import qualified Control.Exception as E
+
+-------------------------------------------------------------------------------
+-- Simple Api
+-------------------------------------------------------------------------------
+
+-- | 'wait' will attempt to connect to the given host and port repeated every
+-- 10 milliseconds until it is successful. It will throw an 'IOError' if the
+-- host cannot be resolved.
+--
+-- A typical use case is to call 'wait' in test code to wait for a server to
+-- start before trying to connect. For example:
+--
+-- @
+--    void $ forkIO $ Warp.run 7000 app
+--    -- Wait for the server to start listening on the socket
+--    wait "127.0.0.1" 7000
+--    -- Communicate with the server
+-- @
+--
+-- If you would like to control the delay or understand how many connection
+-- attempts were made use 'waitWith'.
+--
+-- Since 0.0.0.1
+wait :: String
+     -- ^ Host
+     -> Int
+     -- ^ Port
+     -> IO ()
+wait = waitWith mempty defaultDelay
+
+-------------------------------------------------------------------------------
+-- Advanced Api
+-------------------------------------------------------------------------------
+-- | The default delay between retries is 10000 microseconds (10 ms)
+defaultDelay :: Int
+defaultDelay = 10000
+
+-- | The 'EventHandlers' is a record of 'IO' actions that are called when
+--   interesting events occur in the lifecycle of the 'waitWith' loop.
+--   One can pass in custom 'EventHandlers' values to implement logging
+--   and other forms of instrumentation.
+--
+--   For example for debug one could print out each step using:
+--
+-- @
+--    printHandlers = EventHandlers
+--      { acting     = putStrLn "acting"
+--      , delaying   = putStrLn "delaying"
+--      , restarting = putStrLn "restarting"
+--      }
+-- @
+--
+data EventHandlers m = EventHandlers
+  { connecting :: m ()
+  -- ^ Called before the socket is created for connection
+  , delaying :: m ()
+  -- ^ Called after a failed attempt to connect before the thread
+  -- is put to sleep.
+  , restarting :: m ()
+  -- ^ Called before a recursive call to restart the connection attempt
+  }
+
+instance Semigroup (m ()) => Semigroup (EventHandlers m) where
+  x <> y = EventHandlers
+    { connecting = connecting x <> connecting y
+    , delaying   = delaying   x <> delaying   y
+    , restarting = restarting x <> restarting y
+    }
+
+instance Monoid (m ()) => Monoid (EventHandlers m) where
+  mempty  = EventHandlers mempty mempty mempty
+  mappend = (<>)
+
+-- | Advanced usage. In most situations calling 'wait' will suffice. This allows
+-- one to customize the delay between retries and debug the behavior of the
+-- function. 'wait' is defined as
+--
+-- @
+--    wait = waitWith mempty defaultDelay
+-- @
+--
+-- Since 0.0.0.1
+waitWith :: EventHandlers IO
+         -- ^ A record of IO actions that are called during sp
+         -> Int
+         -- ^ Microseconds to delay
+         -> String
+         -- ^ Host
+         -> Int
+         -- ^ Port
+         -> IO ()
+waitWith eh delay host port
+  = waitM eh (C.threadDelay delay) (connectAction host port)
+
+-- | This function loops if the second argument returns 'False'. Between the
+--   recursive calls it will call it's first argument.
+waitM :: Monad m => EventHandlers m -> m () -> m Bool -> m ()
+waitM eh@EventHandlers {..} delayer action = do
+  connecting
+  action >>= \case
+    True -> pure ()
+    False -> do
+      delaying
+      delayer
+
+      restarting
+      waitM eh delayer action
+
+-- | This function attempts to initiate a TCP connection to a given host and
+--   port. If the host does not exist it throws an IOError. If the connection
+--   fails it returns a 'False'. If it connects it returns 'True'.
+connectAction :: String -> Int -> IO Bool
+connectAction host port = do
+  let hints = N.defaultHints { N.addrSocketType = N.Stream }
+  -- getAddrInfo returns a non-empty array or throws per the doc
+  addr:_ <- N.getAddrInfo (Just hints) (Just host) (Just $ show port)
+  e <- E.try $ E.bracket
+        (N.socket (N.addrFamily addr) (N.addrSocketType addr) (N.addrProtocol addr))
+        N.close
+        $ \sock -> N.connect sock $ N.addrAddress addr
+
+  pure $ case e of
+    Left (_ :: IOError) -> False
+    Right _ -> True
diff --git a/test/Tests/Network/Socket/WaitSpec.hs b/test/Tests/Network/Socket/WaitSpec.hs
--- a/test/Tests/Network/Socket/WaitSpec.hs
+++ b/test/Tests/Network/Socket/WaitSpec.hs
@@ -1,5 +1,6 @@
+{-# OPTIONS_GHC -fno-warn-orphans #-}
 module Tests.Network.Socket.WaitSpec where
-import qualified Network.Socket.Wait as W
+import qualified Network.Socket.Wait.Internal as W
 import qualified Network.Socket.Free as F
 import qualified Network.Socket as N
 import qualified Control.Concurrent as C
@@ -8,80 +9,58 @@
 import qualified Test.Hspec as H
 import qualified System.Timeout as T
 import qualified Control.Concurrent.Async as A
+import qualified Control.Monad.Trans.State as St
 
+instance Monad m => Semigroup (St.StateT s m ()) where
+  (<>) = (>>)
+
+instance Monad m => Monoid (St.StateT s m ()) where
+  mempty = pure ()
+  mappend = (<>)
+
 main :: IO ()
 main = H.hspec spec
 
-data WaitState
-  = Start
-  | Retried
-  | Open
-  deriving (Show, Eq)
-
 throwTimeout :: Int -> String -> IO a -> IO a
 throwTimeout delay msg action = T.timeout delay action >>= \case
   Nothing -> fail $ "timed out " ++ msg
   Just x -> pure x
 
 spec :: H.Spec
-spec = H.before F.openFreePort $ H.after (N.close . snd) $ H.describe "wait" $ do
-  H.it "retries if a socket is not ready" $ \(port, sock) -> do
-    -- There is a small race some what might connect to this
-    N.close sock
-
-    retryRef <- S.newTVarIO 0
-
-    let retryCount = 2
+spec = do
+  H.describe "connectAction" $ H.before F.openFreePort $ H.after (N.close . snd) $ do
+    H.it "returns False if it fails to connect" $ \(port, sock) -> do
+      N.close sock
+      -- There is a small race that another process might bind to this port
+      W.connectAction "127.0.0.1" port `H.shouldReturn` False
 
-        handlers = mempty { W.restarting = S.atomically $ S.modifyTVar' retryRef (+1) }
+    H.it "returns True if it connects" $ \(port, sock) -> do
+      N.listen sock 128
 
-    let start = W.waitWith handlers W.defaultDelay "127.0.0.1" port
-    E.bracket (C.forkIO start) C.killThread $ \_ ->
-      T.timeout
-        (retryCount * W.defaultDelay * 10) -- wait more than long enough
-        (S.atomically $ S.check . (==2) =<< S.readTVar retryRef) >>= \case
-          Nothing
-            -> fail
-            $ "timed out waiting for the wait to retry "
-            ++ show retryCount
-            ++ " times"
-          Just _  -> pure ()
+      W.connectAction "127.0.0.1" port `H.shouldReturn` True
 
-  H.it "connects immediantly to socket when it is ready" $ \(port, sock) -> do
-    N.listen sock 128
+    H.it "returns True if connects after failing" $ \(port, sock) -> do
+      N.close sock
 
-    retryRef <- S.newTVarIO 0
+      W.connectAction "127.0.0.1" port `H.shouldReturn` False
 
-    let handlers = mempty { W.restarting = S.atomically $ S.modifyTVar' retryRef (+1) }
+      E.bracket (N.socket N.AF_INET N.Stream N.defaultProtocol) N.close $ \sock' -> do
+         N.bind sock' $ N.SockAddrInet (fromIntegral port) $ N.tupleToHostAddress (127,0,0,1)
+         N.listen sock' 128
 
-    T.timeout
-      (W.defaultDelay * 10) -- wait more than long enough
-      (W.waitWith handlers W.defaultDelay "127.0.0.1" port) >>= \case
-        Nothing -> fail "timed out waiting to connect"
-        Just _  -> S.atomically (S.readTVar retryRef) `H.shouldReturn` 0
+         W.connectAction "127.0.0.1" port `H.shouldReturn` True
 
-  H.it "connects even if the socket is not open at first" $ \(port, sock) -> do
-    N.close sock
-    retryRef <- S.newTVarIO Start
+    H.it "throws if the host does not exist" $ \(_, _) -> do
+      W.connectAction "invalid." 3000 `H.shouldThrow` (\(_ :: IOError) -> True)
 
-    let handlers = mempty
-          { W.restarting    = S.atomically $ S.writeTVar retryRef Retried
-          , W.createdSocket = S.atomically $ S.readTVar retryRef >>= \case
-              -- Pass through and first and fail
-              Start   -> pure ()
-              -- Block until listening
-              Retried -> S.retry
-              -- Pass through
-              Open    -> pure ()
-          }
+  H.describe "waitM" $ do
+    H.it "returns immediantly if the action returns True" $
+      flip St.execState False (W.waitM mempty (St.put True) (pure True)) `H.shouldBe` False
 
-    let start = W.waitWith handlers W.defaultDelay "127.0.0.1" port
-    throwTimeout 100000 "waiting for wait to return"
-      $ E.bracket (A.async start) A.wait $ \_ -> do
-        S.atomically $ S.check . (==Retried) =<< S.readTVar retryRef
+    H.it "loops until the action returns True" $ do
+      let theAction = do
+            St.get >>= \case
+              0 -> St.put 1 >> pure False
+              _ -> St.put 2 >> pure True
 
-        E.bracketOnError (N.socket N.AF_INET N.Stream N.defaultProtocol) N.close
-          $ \sock' -> do
-            S.atomically $ S.writeTVar retryRef Open
-            N.bind sock' $ N.SockAddrInet (fromIntegral port) $ N.tupleToHostAddress (127,0,0,1)
-            N.listen sock' 128
+      flip St.execState 0 (W.waitM mempty (pure ()) theAction) `H.shouldBe` 2