packages feed

monad-schedule-0.2: test/Yield.hs

{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Yield where

-- base
import Control.Monad (forever)
import Data.Foldable (forM_)
import Data.List.NonEmpty (NonEmpty, reverse)
import qualified Data.List.NonEmpty as NonEmpty
import Data.Maybe (fromJust, isJust, maybeToList)

-- transformers
import Control.Monad.Trans.Class
import Control.Monad.Trans.Reader
import Control.Monad.Trans.Writer (Writer, execWriter, runWriter, tell)

-- QuickCheck
import Test.QuickCheck (counterexample, (==>))

-- test-framework
import Test.Framework

-- test-framework-hunit
import Test.Framework.Providers.HUnit

-- test-framework-quickcheck2
import Test.Framework.Providers.QuickCheck2 (testProperty)

-- HUnit
import Test.HUnit hiding (Test)

-- monad-schedule
import Control.Monad.Schedule.Class (schedule, scheduleAndFinish)
import Control.Monad.Schedule.Trans (runScheduleT)
import Control.Monad.Schedule.Yield

sampleActions :: NonEmpty (MySchedule ())
sampleActions = [yield, yield]

tests =
  testGroup
    "Trans"
    [ testCase "Only leftover time is waited" $
        assertRunsLike sampleActions [Yielded]
    , testCase "Scheduling two waits" $
        assertRunsEqual sampleActions (Data.List.NonEmpty.reverse sampleActions)
    , testCase "Different number of waits" $
        assertRunsLike
          [ myLog "Thread 1 starts" >> yield >> myLog "Thread 1 action" >> yield >> myLog "Thread 1 done"
          , myLog "Thread 2 starts" >> yield >> myLog "Thread 2 done"
          ]
          [ Log "Thread 1 starts"
          , Log "Thread 2 starts"
          , Yielded
          , Log "Thread 1 action"
          , Log "Thread 2 done"
          , Yielded
          , Log "Thread 1 done"
          ]
    , testCase "Blocking thread doesn't starve other thread"
        $ assertRunContains
          [ forever $ myLog "Busy loop starts" >> yield >> myLog "Busy loop ends"
          , myLog "One off thread starts" >> yield >> myLog "One off thread does a thing" >> yield >> myLog "One off thread done"
          ]
        $ Log "One off thread done"
    , testCase "Programs with continuations can be scheduled"
        $ assertProgramsInitiallyRunsLike
          [[Log "Thread 1 active", Log "Thread 2 active", Yielded], [Log "Thread 1 active", Log "Thread 2 active", Yielded], [Log "Thread 1 active", Log "Thread 2 active", Yielded], [Log "Thread 1 active", Log "Thread 2 active", Yielded], [Log "Thread 1 active", Log "Thread 2 active", Yielded]]
          [foreverP (const ["Thread 1 active"]), foreverP (const ["Thread 2 active"])]
        $ repeat True
    , testCase "Two programs that tick alternately can be scheduled" $
        assertProgramsInitiallyRunsLike
          [[Log "1 Nope", Log "2 Yes", Yielded, Log "1 Nope"], [Log "2 Nope", Yielded, Log "2 Nope", Log "1 Yes", Yielded, Log "2 Nope"], [Log "1 Yes", Yielded, Log "2 Nope"], [Log "1 Nope", Yielded, Log "1 Nope", Log "2 Yes", Yielded, Log "1 Nope"], [Log "2 Yes", Yielded, Log "1 Nope"], [Log "2 Yes", Yielded, Log "1 Nope"], [Log "2 Yes", Yielded, Log "1 Nope"], [Log "2 Nope", Yielded, Log "2 Nope", Log "1 Yes", Yielded, Log "2 Nope"]]
          twoPrograms
          [True, False, False, True, True, True, True, False, False]
    , testProperty "Two programs that tick alternately can be scheduled with arbitrary input" $
        \(inputs :: [Bool]) skip ->
          let log = take (20 * length inputs) $ drop skip $ concat $ runProgramWith inputs $ schedulePrograms twoPrograms
              isContained expectedEntry = expectedEntry `elem` log
           in counterexample (show log) $
                all (`elem` drop skip inputs) ([True, False] :: [Bool]) ==>
                  all isContained ([Log "1 Yes", Log "2 Yes"] :: [Event])
    ]

assertRunsEqual :: NonEmpty (MySchedule a1) -> NonEmpty (MySchedule a2) -> Assertion
assertRunsEqual actions1 actions2 = assertEqual "Should run the same under scheduling" (runMySchedule actions1) (runMySchedule actions2)

assertRunsLike :: NonEmpty (MySchedule a) -> [Event] -> Assertion
assertRunsLike actions events = assertEqual "Should run like the following under scheduling" events (runMySchedule actions)

assertRunContains :: NonEmpty (MySchedule a) -> Event -> Assertion
assertRunContains actions event = assertBool ("The run should contain the event " ++ show event) $ event `elem` runMySchedule actions

assertInitiallyRunsLike :: NonEmpty (MySchedule a) -> [Event] -> Assertion
assertInitiallyRunsLike actions events = assertEqual "Should, at the beginning, run like the following under scheduling" events $ take (length events) $ runMySchedule actions

data Event
  = Log String
  | Yielded
  deriving (Eq, Show)

type MySchedule a = YieldT (Writer [Event]) a

myLog :: String -> MySchedule ()
myLog = lift . tell . pure . Log

runMySchedule :: NonEmpty (MySchedule a) -> [Event]
runMySchedule = execWriter . runScheduleT (tell . pure . const Yielded) . scheduleAndFinish

type MyReaderSchedule a = YieldT (ReaderT Bool (Writer [Event])) a

-- Ok this is basically ListT
newtype Program = Program {unProgram :: MyReaderSchedule (Maybe Program)}

foreverP :: (Bool -> [String]) -> Program
foreverP action = go
  where
    go = Program $ do
      input <- lift ask
      lift $ lift $ tell $ Log <$> action input
      yield
      return $ Just go

-- Returning Left means looping further, Right means going to the next continuation
wait :: (Bool -> Either String String) -> Program
wait f = go
  where
    go = Program loop
    loop = do
      input <- lift ask
      case f input of
        Left msg -> do
          myReaderLog msg
          yield
          loop
        Right msg -> do
          myReaderLog msg
          yield
          return $ Just go

twoPrograms =
  [ wait (\b -> if b then Left "1 Nope" else Right "1 Yes")
  , wait (\b -> if b then Right "2 Yes" else Left "2 Nope")
  ]

runProgram :: Program -> MyReaderSchedule ()
runProgram Program {..} = do
  tick <- unProgram
  forM_ tick runProgram

myReaderLog :: String -> MyReaderSchedule ()
myReaderLog = lift . lift . tell . pure . Log

schedulePrograms :: NonEmpty Program -> Program
schedulePrograms programs = Program $ do
  (done, running) <- schedule $ unProgram <$> programs
  return $ fmap schedulePrograms $ NonEmpty.nonEmpty $ (fromJust <$> NonEmpty.filter isJust done) ++ (Program <$> running)

-- Should be possible with some recursion scheme
runProgramWith :: [Bool] -> Program -> [[Event]]
runProgramWith [] _ = []
runProgramWith (input : inputs) Program {..} =
  let (cont, events) = runWriter $ flip runReaderT input $ runScheduleT (const $ lift $ tell [Yielded]) unProgram
   in events : (runProgramWith inputs =<< maybeToList cont)

assertProgramsInitiallyRunsLike :: [[Event]] -> NonEmpty Program -> [Bool] -> Assertion
assertProgramsInitiallyRunsLike events programs inputs =
  assertEqual "The programs, when scheduled, should run like" events $
    take (length events) $
      runProgramWith inputs $
        schedulePrograms programs