module Test.Async.Combinators where

import Universum

import Control.Concurrent (killThread, myThreadId, threadDelay)
import Control.Exception (AsyncException (ThreadKilled))
import Control.Exception.Safe (catchAny, handleAsync, isAsyncException, throw, throwTo, tryAsync)

import Test.HUnit (Assertion, assertEqual, assertFailure, (@?))

import Async.Combinators


data TestException = TestException
    deriving Show

instance Exception TestException

-----------------------
-- withWorker
-----------------------

data WorkerTestResult = WorkerTestResult
    { wtrReturn     :: Either WorkerExited ()
    , wtrWorkerDone :: Bool
    , wtrGoDone     :: Bool
    }

withWorker_test :: IO () -> IO () -> IO WorkerTestResult
withWorker_test worker go = do
    workerDone <- newIORef False
    goDone     <- newIORef False

    r <- tryAsync $ withWorker "worker" (worker *> writeIORef workerDone True)
                                        (go *> writeIORef goDone True)

    WorkerTestResult r <$> readIORef workerDone <*> readIORef goDone

assertWtr :: Bool -> Bool -> Bool -> WorkerTestResult -> Assertion
assertWtr expectRight expectWorkerDone expectGoDone WorkerTestResult{..} = do
    assertEqual "return right" expectRight (isRight wtrReturn)
    assertEqual "worker done"  expectWorkerDone wtrWorkerDone
    assertEqual "go done"      expectGoDone wtrGoDone


--------
-- First, some sanity checks.
--------

unit_withWorker_exception_is_async :: Assertion
unit_withWorker_exception_is_async = do
    isAsyncException (WorkerExited "hi") @? "isAsync"

    me <- myThreadId
    r <- tryAsync $ catchAny (throwTo me $ WorkerExited "hi") (const pass)
    case r of
        Right _                  -> assertFailure $ "caught by safe `catchAny`"
        Left (_ :: WorkerExited) -> pass


unit_withWorker_exception_can_be_caught :: Assertion
unit_withWorker_exception_can_be_caught = do
    me <- myThreadId
    handleAsync (\(_ :: WorkerExited) -> pass) $ throwTo me (WorkerExited "bye")

--------
-- Now actual tests.
--------

-- | Main thread exits -> worker is killed.
unit_withWorker_main_exit :: Assertion
unit_withWorker_main_exit =
    withWorker_test (forever $ threadDelay 100) (threadDelay 100) >>=
    assertWtr True False True

-- | Main crashes -> worker is killed.
unit_withWorker_main_crash :: Assertion
unit_withWorker_main_crash =
    withWorker_test (forever $ threadDelay 100) (throw TestException `catch` \TestException -> putText "Hi") >>=
    assertWtr True False True

-- | Worker exits -> main thread is killed.
unit_withWorker_worker_exit :: Assertion
unit_withWorker_worker_exit =
    withWorker_test pass (threadDelay 100) >>=
    assertWtr False True False

-- | Worker crashes -> main thread is killed.
unit_withWorker_worker_crash :: Assertion
unit_withWorker_worker_crash =
    withWorker_test (throw TestException) (threadDelay 100000) >>=
    assertWtr False False False

-- | Worker receives an async exception -> main thread is killed.
unit_withWorker_worker_crash_async :: Assertion
unit_withWorker_worker_crash_async =
    withWorker_test worker (threadDelay 100000) >>=
    assertWtr False False False
  where
    worker = myThreadId >>= \me -> throwTo me TestException

-- | Worker receives ThreadKilled -> main thread is killed.
unit_withWorker_worker_crash_ThreadKilled :: Assertion
unit_withWorker_worker_crash_ThreadKilled =
    withWorker_test worker (threadDelay 100000) >>=
    assertWtr False False False
  where
    worker = myThreadId >>= \me -> throwTo me ThreadKilled

-- | Main thread kills worker -> main thread is killed.
unit_withWorker_main_kills_worker :: Assertion
unit_withWorker_main_kills_worker = do
    v <- newEmptyMVar
    wtr <- withWorker_test (myThreadId >>= putMVar v  >> threadDelay 1000000)
                           (takeMVar v >>= killThread >> threadDelay 1000000)
    assertWtr False False False wtr