content: --------------------------------------------------------------------------------
  {-# LANGUAGE OverloadedStrings #-}
  {-# OPTIONS_GHC -fno-warn-orphans #-}
  module Network.WebSockets.Tests
      ( tests
      ) where
  
  
  --------------------------------------------------------------------------------
  import qualified Data.ByteString.Builder               as Builder
  import           Control.Applicative                   ((<$>))
  import           Control.Concurrent                    (forkIO)
  import           Control.Exception                     (try)
  import           Control.Monad                         (replicateM)
  import           Data.Binary.Get                       (runGetOrFail)
  import qualified Data.ByteString.Lazy                  as BL
  import           Data.List                             (intersperse)
  import           Data.Maybe                            (catMaybes)
  import           Data.Monoid                           (mempty, mconcat)
  import           Network.WebSockets
  import qualified Network.WebSockets.Hybi13             as Hybi13
  import           Network.WebSockets.Hybi13.Demultiplex
  import           Network.WebSockets.Protocol
  import qualified Network.WebSockets.Stream             as Stream
  import           Network.WebSockets.Tests.Util
  import           Network.WebSockets.Types
  import           Test.Framework                        (Test, testGroup)
  import           Test.Framework.Providers.HUnit        (testCase)
  import           Test.Framework.Providers.QuickCheck2  (testProperty)
  import           Test.HUnit                            ((@=?))
  import           Test.QuickCheck                       (Arbitrary (..), Gen,
                                                          Property)
  import qualified Test.QuickCheck                       as QC
  import qualified Test.QuickCheck.Monadic               as QC
  import           Prelude
  
  
  --------------------------------------------------------------------------------
  tests :: Test
  tests = testGroup "Network.WebSockets.Test"
      [ testProperty "simple encode/decode Hybi13" (testSimpleEncodeDecode Hybi13)
      , testProperty "fragmented Hybi13"           testFragmentedHybi13
      , testRfc_6455_5_5_1
      , testRfc_6455_5_5_2
      , testFramePayloadSizeLimit
      ]
  
  --------------------------------------------------------------------------------
  testSimpleEncodeDecode :: Protocol -> Property
  testSimpleEncodeDecode protocol = QC.monadicIO $
      QC.forAllM QC.arbitrary $ \msgs -> QC.run $ do
          echo  <- Stream.makeEchoStream
          parse <- decodeMessages protocol mempty mempty echo
          write <- encodeMessages protocol ClientConnection echo
          _     <- forkIO $ write msgs
          msgs' <- catMaybes <$> replicateM (length msgs) parse
          Stream.close echo
          msgs @=? msgs'
  
  
  --------------------------------------------------------------------------------
  testFragmentedHybi13 :: Property
  testFragmentedHybi13 = QC.monadicIO $
      QC.forAllM QC.arbitrary $ \fragmented -> QC.run $ do
          echo     <- Stream.makeEchoStream
          parse    <- Hybi13.decodeMessages mempty mempty echo
          -- is'      <- Streams.filter isDataMessage =<< Hybi13.decodeMessages is
  
          -- Simple hacky encoding of all frames
          _ <- forkIO $ do
              mapM_ (Stream.write echo)
                  [ Builder.toLazyByteString (Hybi13.encodeFrame Nothing f)
                  | FragmentedMessage _ frames <- fragmented
                  , f                          <- frames
                  ]
              Stream.close echo
  
          -- Check if we got all data
          msgs <- filter isDataMessage <$> parseAll parse
          [msg | FragmentedMessage msg _ <- fragmented] @=? msgs
    where
      isDataMessage (ControlMessage _)    = False
      isDataMessage (DataMessage _ _ _ _) = True
  
      parseAll parse = do
          mbMsg <- try parse
          case mbMsg of
              Left  ConnectionClosed -> return []
              Left  _                -> return []
              Right (Just msg)       -> (msg :) <$> parseAll parse
              Right Nothing          -> return []
  
  --------------------------------------------------------------------------------
  testRfc_6455_5_5_1 :: Test
  testRfc_6455_5_5_1 =
      testCase "RFC 6455, 5.5: Frame encoder shall truncate control frame payload to 125 bytes" $ do
          260 @=? BL.length (encodedFrame ContinuationFrame)
          260 @=? BL.length (encodedFrame TextFrame)
          260 @=? BL.length (encodedFrame BinaryFrame)
          127 @=? BL.length (encodedFrame CloseFrame)
          127 @=? BL.length (encodedFrame PingFrame)
          127 @=? BL.length (encodedFrame PongFrame)
      where
          payload256 = BL.replicate 256 0
          encodedFrame ft
              = Builder.toLazyByteString
              $ Hybi13.encodeFrame Nothing (Frame True False False False ft payload256)
  
  --------------------------------------------------------------------------------
  testRfc_6455_5_5_2 :: Test
  testRfc_6455_5_5_2 =
      testCase "RFC 6455, 5.5: Frame decoder shall fail if control frame payload length > 125 bytes" $
          Left (BL.drop 4 ping126, 4, errMsg) @=? runGetOrFail (Hybi13.parseFrame mempty) ping126
      where
          errMsg = "Control Frames must not carry payload > 125 bytes!"
          ping126 = mconcat
             [ "\137\254\NUL~\219\SI\190\252\219\SI\190\252\219\SI\190\252\219"
             , "\SI\190\252\219\SI\190\252\219\SI\190\252\219\SI\190\252\219"
             , "\SI\190\252\219\SI\190\252\219\SI\190\252\219\SI\190\252\219"
             , "\SI\190\252\219\SI\190\252\219\SI\190\252\219\SI\190\252\219"
             , "\SI\190\252\219\SI\190\252\219\SI\190\252\219\SI\190\252\219"
             , "\SI\190\252\219\SI\190\252\219\SI\190\252\219\SI\190\252\219"
             , "\SI\190\252\219\SI\190\252\219\SI\190\252\219\SI\190\252\219"
             , "\SI\190\252\219\SI\190\252\219\SI\190\252\219\SI\190\252\219"
             , "\SI\190\252\219\SI"
             ]
  
  testFramePayloadSizeLimit :: Test
  testFramePayloadSizeLimit = testGroup "FramePayloadSizeLimit Hybi13"
      [ testCase "OK 1" $ case parse (frame 99) of
          Right _ -> return ()
          Left _  -> fail "Expecting successful parse."
      , testCase "OK 2" $ case parse (frame 100) of
          Right _ -> return ()
          Left _  -> fail "Expecting successful parse."
      , testCase "Exceed" $ case parse (frame 101) of
          Right _ -> fail "Expecting parse to fail."
          Left _  -> return ()
      ]
    where
      parse   = runGetOrFail (Hybi13.parseFrame (SizeLimit 100))
      frame n = Builder.toLazyByteString $ Hybi13.encodeFrame Nothing $
          Frame True False False False BinaryFrame (BL.replicate n 20)
  
  
  --------------------------------------------------------------------------------
  instance Arbitrary Message where
      arbitrary = QC.oneof
          [ do
              payload <- BL.take 125 . BL.pack <$> arbitrary
              return $ ControlMessage (Ping payload)
          , do
              payload <- BL.take 125 . BL.pack <$> arbitrary
              return $ ControlMessage (Pong payload)
          , do
              payload <- BL.pack <$> arbitrary
              return $ DataMessage False False False (Text payload Nothing)
          , do
              payload <- BL.pack <$> arbitrary
              return $ DataMessage False False False (Binary payload)
          ]
  
  
  --------------------------------------------------------------------------------
  data FragmentedMessage = FragmentedMessage Message [Frame]
      deriving (Show)
  
  
  --------------------------------------------------------------------------------
  instance Arbitrary FragmentedMessage where
      arbitrary = do
          -- Pick a frametype and a corresponding random payload
          ft        <- QC.elements [TextFrame, BinaryFrame]
          payload   <- case ft of
              TextFrame -> arbitraryUtf8
              _         -> arbitraryByteString
  
          fragments <- arbitraryFragmentation payload
          let fs  = makeFrames $ zip (ft : repeat ContinuationFrame) fragments
              msg = case ft of
                  TextFrame   -> DataMessage False False False (Text payload Nothing)
                  BinaryFrame -> DataMessage False False False (Binary payload)
                  _           -> error "Arbitrary FragmentedMessage crashed"
  
          interleaved <- arbitraryInterleave genControlFrame fs
          return $ FragmentedMessage msg interleaved
          -- return $ FragmentedMessage msg fs
        where
          makeFrames []              = []
          makeFrames [(ft, pl)]      = [Frame True False False False ft pl]
          makeFrames ((ft, pl) : fr) =
              Frame False False False False ft pl : makeFrames fr
  
          genControlFrame = QC.elements
              [ Frame True False False False PingFrame "Herp"
              , Frame True False False False PongFrame "Derp"
              ]
  
  
  --------------------------------------------------------------------------------
  arbitraryFragmentation :: BL.ByteString -> Gen [BL.ByteString]
  arbitraryFragmentation bs = arbitraryFragmentation' bs
    where
      len :: Int
      len = fromIntegral $ BL.length bs
      arbitraryFragmentation' bs' = do
          -- TODO: we currently can't send packets of length 0. We should
          -- investigate why (regardless of the spec).
          n <- QC.choose (1, len - 1)
          let (l, r) = BL.splitAt (fromIntegral n) bs'
          case r of
              "" -> return [l]
              _  -> (l :) <$> arbitraryFragmentation' r
  
  
  --------------------------------------------------------------------------------
  arbitraryInterleave :: Gen a -> [a] -> Gen [a]
  arbitraryInterleave sep xs = fmap concat $ sequence $
      [sep'] ++ intersperse sep' [return [x] | x <- xs] ++ [sep']
    where
      sep' = QC.sized $ \size -> do
          num <- QC.choose (1, size)
          replicateM num sep
  