-- | Polling helpers for tests.
module Arbiter.Test.Poll
  ( waitUntil
  , withLinkedAsync
  ) where

import Control.Concurrent (threadDelay)
import Control.Monad (unless)
import GHC.Stack (HasCallStack)
import Test.Hspec (expectationFailure)
import UnliftIO (MonadUnliftIO)
import UnliftIO.Async (Async, link, withAsync)

-- | Poll every 100 ms until the predicate returns 'True'.
-- Fails with 'expectationFailure' after @timeoutMs@ milliseconds.
waitUntil :: (HasCallStack) => Int -> IO Bool -> IO ()
waitUntil :: HasCallStack => Int -> IO Bool -> IO ()
waitUntil Int
timeoutMs IO Bool
check = HasCallStack => Int -> IO ()
Int -> IO ()
go (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 (Int
timeoutMs Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
100))
  where
    go :: (HasCallStack) => Int -> IO ()
    go :: HasCallStack => Int -> IO ()
go Int
0 = HasCallStack => String -> IO ()
String -> IO ()
expectationFailure String
"waitUntil: timed out waiting for condition"
    go Int
remaining = do
      satisfied <- IO Bool
check
      unless satisfied $ do
        threadDelay 100_000
        go (remaining - 1)

-- | Run 'withAsync' with a linked handle. Raise a child-thread failure in the
-- caller and stop the associated 'waitUntil'.
withLinkedAsync :: (MonadUnliftIO m) => m a -> (Async a -> m b) -> m b
withLinkedAsync :: forall (m :: * -> *) a b.
MonadUnliftIO m =>
m a -> (Async a -> m b) -> m b
withLinkedAsync m a
run Async a -> m b
body = m a -> (Async a -> m b) -> m b
forall (m :: * -> *) a b.
MonadUnliftIO m =>
m a -> (Async a -> m b) -> m b
withAsync m a
run ((Async a -> m b) -> m b) -> (Async a -> m b) -> m b
forall a b. (a -> b) -> a -> b
$ \Async a
handle -> Async a -> m ()
forall (m :: * -> *) a. MonadIO m => Async a -> m ()
link Async a
handle m () -> m b -> m b
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Async a -> m b
body Async a
handle