content: {-# LANGUAGE CPP, OverloadedStrings #-}
  module PubSubTest (testPubSubThreaded) where
  
  import Control.Concurrent
  import Control.Monad
  import Control.Concurrent.Async
  import Control.Exception
  import qualified Data.List
  import Data.Text hiding (show)
  import Data.Typeable
  import Data.ByteString
  import Control.Concurrent.STM
  import qualified Test.Framework as Test
  import qualified Test.Framework.Providers.HUnit as Test (testCase)
  import qualified Test.HUnit as HUnit
  
  import Database.Redis
  
  testPubSubThreaded :: [Connection -> Test.Test]
  testPubSubThreaded = [removeAllTest, callbackErrorTest, removeFromUnregister]
  
  -- | A handler label to be able to distinguish the handlers from one another
  -- to help make sure we unregister the correct handler.
  type HandlerLabel = Text
  
  data TestMsg = MsgFromChannel HandlerLabel ByteString
               | MsgFromPChannel HandlerLabel RedisChannel ByteString
    deriving (Show, Eq)
  
  type MsgVar = TVar [TestMsg]
  
  -- | A handler that just writes the message into the TVar
  handler :: HandlerLabel -> MsgVar -> MessageCallback
  handler label ref msg = atomically $
    modifyTVar ref $ \x -> x ++ [MsgFromChannel label msg]
  
  -- | A pattern handler that just writes the message into the TVar
  phandler :: HandlerLabel -> MsgVar -> PMessageCallback
  phandler label ref chan msg = atomically $
    modifyTVar ref $ \x -> x ++ [MsgFromPChannel label chan msg]
  
  -- | Wait for a given message to be received
  waitForMessage :: MsgVar -> HandlerLabel -> ByteString -> IO ()
  waitForMessage ref label msg = atomically $ do
    let expected = MsgFromChannel label msg
    lst <- readTVar ref
    unless (expected `Prelude.elem` lst) retry
    writeTVar ref $ Prelude.filter (/= expected) lst
  
  -- | Wait for a given pattern message to be received
  waitForPMessage :: MsgVar -> HandlerLabel -> RedisChannel -> ByteString -> IO ()
  waitForPMessage ref label chan msg = atomically $ do
    let expected = MsgFromPChannel label chan msg
    lst <- readTVar ref
    unless (expected `Prelude.elem` lst) retry
    writeTVar ref $ Prelude.filter (/= expected) lst
  
  expectRedisChannels :: Connection -> [RedisChannel] -> IO ()
  expectRedisChannels conn expected = do
    actual <- runRedis conn $ sendRequest ["PUBSUB", "CHANNELS"]
    case actual of
      Left err -> HUnit.assertFailure $ "Error geting channels: " ++ show err
      Right s -> HUnit.assertEqual "redis channels" (Data.List.sort s) (Data.List.sort expected)
  
  -- | Test basic messages, plus using removeChannels
  removeAllTest :: Connection -> Test.Test
  removeAllTest conn = Test.testCase "Multithreaded Pub/Sub - basic" $ do
    msgVar <- newTVarIO []
    initialComplete <- newTVarIO False
    ctrl <- newPubSubController [("foo1", handler "InitialFoo1" msgVar), ("foo2", handler "InitialFoo2" msgVar)]
                                [("bar1:*", phandler "InitialBar1" msgVar), ("bar2:*", phandler "InitialBar2" msgVar)]
    withAsync (pubSubForever conn ctrl (atomically $ writeTVar initialComplete True)) $ \_ -> do
      -- wait for initial
      atomically $ readTVar initialComplete >>= \b -> if b then return () else retry
      expectRedisChannels conn ["foo1", "foo2"]
  
      runRedis conn $ publish "foo1" "Hello"
      waitForMessage msgVar "InitialFoo1" "Hello"
  
      runRedis conn $ publish "bar2:zzz" "World"
      waitForPMessage msgVar "InitialBar2" "bar2:zzz" "World"
  
      -- subscribe to foo1 and bar1 again
      addChannelsAndWait ctrl [("foo1", handler "NewFoo1" msgVar)] [("bar1:*", phandler "NewBar1" msgVar)]
      expectRedisChannels conn ["foo1", "foo2"]
  
      runRedis conn $ publish "foo1" "abcdef"
      waitForMessage msgVar "InitialFoo1" "abcdef"
      waitForMessage msgVar "NewFoo1" "abcdef"
  
      -- unsubscribe from foo1 and bar1
      removeChannelsAndWait ctrl ["foo1", "unusued"] ["bar1:*", "unused:*"]
      expectRedisChannels conn ["foo2"]
  
      -- foo2 and bar2 are still subscribed
      runRedis conn $ publish "foo2" "12345"
      waitForMessage msgVar "InitialFoo2" "12345"
  
      runRedis conn $ publish "bar2:aaa" "0987"
      waitForPMessage msgVar "InitialBar2" "bar2:aaa" "0987"
  
  data TestError = TestError ByteString
    deriving (Eq, Show)
  instance Exception TestError
  
  -- | Test an error thrown from a message handler
  callbackErrorTest :: Connection -> Test.Test
  callbackErrorTest conn = Test.testCase "Multithreaded Pub/Sub - error in handler" $ do
    initialComplete <- newTVarIO False
    ctrl <- newPubSubController [("foo", throwIO . TestError)] []
  
    thread <- async (pubSubForever conn ctrl (atomically $ writeTVar initialComplete True))
    atomically $ readTVar initialComplete >>= \b -> if b then return () else retry
  
    runRedis conn $ publish "foo" "Hello"
  
    ret <- waitCatch thread
    case ret of
      Left (SomeException e) | cast e == Just (TestError "Hello") -> return ()
      _ -> HUnit.assertFailure $ "Did not properly throw error from message thread " ++ show ret
  
  -- | Test removing channels by using the return value of 'addHandlersAndWait'.
  removeFromUnregister :: Connection -> Test.Test
  removeFromUnregister conn = Test.testCase "Multithreaded Pub/Sub - unregister handlers" $ do
    msgVar <- newTVarIO []
    initialComplete <- newTVarIO False
    ctrl <- newPubSubController [] []
    withAsync (pubSubForever conn ctrl (atomically $ writeTVar initialComplete True)) $ \_ -> do
      atomically $ readTVar initialComplete >>= \b -> if b then return () else retry
  
      -- register to some channels
      void $ addChannelsAndWait ctrl
          [("abc", handler "InitialAbc" msgVar), ("xyz", handler "InitialXyz" msgVar)]
          [("def:*", phandler "InitialDef" msgVar), ("uvw", phandler "InitialUvw" msgVar)]
      expectRedisChannels conn ["abc", "xyz"]
  
      runRedis conn $ publish "abc" "Hello"
      waitForMessage msgVar "InitialAbc" "Hello"
  
      -- register to some more channels
      unreg <- addChannelsAndWait ctrl
          [("abc", handler "SecondAbc"  msgVar), ("123", handler "Second123" msgVar)]
          [("def:*", phandler "SecondDef" msgVar), ("890:*", phandler "Second890" msgVar)]
      expectRedisChannels conn ["abc", "xyz", "123"]
  
      -- check messages on all channels
      runRedis conn $ publish "abc" "World"
      waitForMessage msgVar "InitialAbc" "World"
      waitForMessage msgVar "SecondAbc" "World"
  
      runRedis conn $ publish "123" "World2"
      waitForMessage msgVar "Second123" "World2"
  
      runRedis conn $ publish "def:bbbb" "World3"
      waitForPMessage msgVar "InitialDef" "def:bbbb" "World3"
      waitForPMessage msgVar "SecondDef" "def:bbbb" "World3"
  
      runRedis conn $ publish "890:tttt" "World4"
      waitForPMessage msgVar "Second890" "890:tttt" "World4"
  
      -- unregister
      unreg
  
      -- we have no way of waiting until unregister actually happened, so just delay and hope
      threadDelay $ 1000*1000 -- 1 second
      expectRedisChannels conn ["abc", "xyz"]
  
      -- now only initial should be around. In particular, abc should still be subscribed
      runRedis conn $ publish "abc" "World5"
      waitForMessage msgVar "InitialAbc" "World5"
  
      runRedis conn $ publish "def:cccc" "World6"
      waitForPMessage msgVar "InitialDef" "def:cccc" "World6"
  