content: {-# LANGUAGE GADTs #-}
  {-# LANGUAGE RankNTypes #-}
  {-# LANGUAGE OverloadedStrings #-}
  module Data.Conduit.BinarySpec (spec) where
  
  import Data.Conduit (runConduit, runConduitRes, (.|), runConduitPure, ConduitT, Void)
  import qualified Data.Conduit.Binary as CB
  import qualified Data.Conduit as C
  import qualified Data.Conduit.List as CL
  import Control.Monad.Trans.Resource
  import Control.Monad.IO.Class
  import Control.Exception (IOException)
  import qualified Data.ByteString.Lazy as L
  import Test.Hspec
  import Test.Hspec.QuickCheck
  import qualified Data.IORef as I
  import Control.Monad.Trans.Writer.Strict
  import qualified Data.ByteString as S
  import qualified Data.ByteString.Char8 as S8
  import Data.Functor.Identity
  import Test.QuickCheck.Arbitrary (Arbitrary, arbitrary)
  import Test.QuickCheck.Gen (Gen, oneof)
  import Data.Word (Word8)
  import Foreign.Storable (Storable, sizeOf, pokeByteOff, alignment)
  import Data.Typeable (Typeable)
  import Data.ByteString.Internal (createAndTrim')
  import Foreign.Ptr (alignPtr, minusPtr)
  import System.Directory (doesFileExist)
  import System.IO.Unsafe (unsafePerformIO)
  import Control.Applicative ((<$>))
  import qualified Data.ByteString.Builder as BB
  
  spec :: Spec
  spec = describe "Data.Conduit.Binary" $ do
  
      describe "file access" $ do
          it "read" $ do
              bs <- S.readFile "conduit-extra.cabal"
              bss <- runConduitRes $ CB.sourceFile "conduit-extra.cabal" .| CL.consume
              S.concat bss `shouldBe` bs
  
          it "read range" $ do
              S.writeFile "tmp" "0123456789"
              bss <- runConduitRes $ CB.sourceFileRange "tmp" (Just 2) (Just 3) .| CL.consume
              S.concat bss `shouldBe` "234"
  
          it "write" $ do
              runConduitRes $ CB.sourceFile "conduit-extra.cabal" .| CB.sinkFile "tmp"
              bs1 <- S.readFile "conduit-extra.cabal"
              bs2 <- S.readFile "tmp"
              bs2 `shouldBe` bs1
  
  
          it "write builder (withSinkFileBuilder)" $ do
              CB.withSinkFileBuilder "tmp" $ \sink ->
                CB.withSourceFile "conduit-extra.cabal" $ \src ->
                runConduit $ src .| CL.map BB.byteString .| sink
              bs1 <- S.readFile "conduit-extra.cabal"
              bs2 <- S.readFile "tmp"
              bs2 `shouldBe` bs1
  
          it "conduit" $ do
              runConduitRes
                  $ CB.sourceFile "conduit-extra.cabal"
                 .| CB.conduitFile "tmp"
                 .| CB.sinkFile "tmp2"
              bs1 <- S.readFile "conduit-extra.cabal"
              bs2 <- S.readFile "tmp"
              bs3 <- S.readFile "tmp2"
              bs2 `shouldBe` bs1
              bs3 `shouldBe` bs1
      describe "binary isolate" $ do
          it "works" $ do
              bss <- runConduitRes
                   $ CL.sourceList (replicate 1000 "X")
                  .| CB.isolate 6
                  .| CL.consume
              S.concat bss `shouldBe` "XXXXXX"
  
      describe "properly using binary file reading" $ do
          it "sourceFile" $ do
              x <- runConduitRes $ CB.sourceFile "test/random" .| CL.consume
              lbs <- L.readFile "test/random"
              L.fromChunks x `shouldBe` lbs
  
      describe "binary head" $ do
          let go lbs = do
                  x <- CB.head
                  case (x, L.uncons lbs) of
                      (Nothing, Nothing) -> return True
                      (Just y, Just (z, lbs'))
                          | y == z -> go lbs'
                      _ -> return False
  
          prop "works" $ \bss' ->
              let bss = map S.pack bss'
               in runConduitPure $ CL.sourceList bss .| go (L.fromChunks bss)
      describe "binary takeWhile" $ do
          prop "works" $ \bss' ->
              let bss = map S.pack bss'
               in runIdentity $ do
                  bss2 <- runConduit $ CL.sourceList bss .| CB.takeWhile (>= 5) .| CL.consume
                  return $ L.fromChunks bss2 == L.takeWhile (>= 5) (L.fromChunks bss)
          prop "leftovers present" $ \bss' ->
              let bss = map S.pack bss'
               in runIdentity $ do
                  result <- runConduit $ CL.sourceList bss .| do
                      x <- CB.takeWhile (>= 5) .| CL.consume
                      y <- CL.consume
                      return (S.concat x, S.concat y)
                  let expected = S.span (>= 5) $ S.concat bss
                  if result == expected
                      then return True
                      else error $ show (S.concat bss, result, expected)
  
      describe "binary dropWhile" $ do
          prop "works" $ \bss' ->
              let bss = map S.pack bss'
               in runIdentity $ do
                  bss2 <- runConduit $ CL.sourceList bss .| do
                      CB.dropWhile (< 5)
                      CL.consume
                  return $ L.fromChunks bss2 == L.dropWhile (< 5) (L.fromChunks bss)
  
      describe "binary take" $ do
        let go n l = runConduit $ CL.sourceList l .| do
                a <- CB.take n
                b <- CL.consume
                return (a, b)
  
        -- Taking nothing should result in an empty Bytestring
        it "nothing" $ do
          (a, b) <- runResourceT $ go 0 ["abc", "defg"]
          a              `shouldBe` L.empty
          L.fromChunks b `shouldBe` "abcdefg"
  
        it "normal" $ do
          (a, b) <- runResourceT $ go 4 ["abc", "defg"]
          a              `shouldBe` "abcd"
          L.fromChunks b `shouldBe` "efg"
  
        -- Taking exactly the data that is available should result in no
        -- leftover.
        it "all" $ do
          (a, b) <- runResourceT $ go 7 ["abc", "defg"]
          a `shouldBe` "abcdefg"
          b `shouldBe` []
  
        -- Take as much as possible.
        it "more" $ do
          (a, b) <- runResourceT $ go 10 ["abc", "defg"]
          a `shouldBe` "abcdefg"
          b `shouldBe` []
  
      describe "binary" $ do
          prop "lines" $ \bss' -> runIdentity $ do
              let bss = map S.pack bss'
                  bs = S.concat bss
                  src = CL.sourceList bss
              res <- runConduit $ src .| CB.lines .| CL.consume
              return $ S8.lines bs == res
  
      describe "sinkCacheLength" $ do
          it' "works" $ runResourceT $ do
              lbs <- liftIO $ L.readFile "test/Data/Conduit/BinarySpec.hs"
              (len, src) <- runConduit $ CB.sourceLbs lbs .| CB.sinkCacheLength
              lbs' <- runConduit $ src .| CB.sinkLbs
              liftIO $ do
                  fromIntegral len `shouldBe` L.length lbs
                  lbs' `shouldBe` lbs
                  fromIntegral len `shouldBe` L.length lbs'
  
      describe "sinkFileCautious" $ do
        it' "success" $ do
          runConduitRes $ CB.sourceFile "conduit-extra.cabal" .| CB.sinkFileCautious "tmp"
          bs1 <- S.readFile "conduit-extra.cabal"
          bs2 <- S.readFile "tmp"
          bs2 `shouldBe` bs1
        it' "failure" $ do
          let bs1 = "This is the original content"
          S.writeFile "tmp" bs1
          runConduitRes
                 ( (CB.sourceFile "conduit-extra.cabal" >> error "FIXME")
              .| CB.sinkFileCautious "tmp")
                 `shouldThrow` anyException
          bs2 <- S.readFile "tmp"
          bs2 `shouldBe` bs1
  
      it "sinkSystemTempFile" $ do
          let bs = "Hello World!"
          fp <- runResourceT $ do
            fp <- runConduit $ C.yield bs .| CB.sinkSystemTempFile "temp-file-test"
            liftIO $ do
              actual <- S.readFile fp
              actual `shouldBe` bs
            return fp
          exists <- doesFileExist fp
          exists `shouldBe` False
  
      describe "Data.Conduit.Binary.mapM_" $ do
          prop "telling works" $ \bytes ->
              let lbs = L.pack bytes
                  src = CB.sourceLbs lbs
                  sink = CB.mapM_ (tell . return . S.singleton)
                  bss = execWriter $ runConduit $ src .| sink
               in L.fromChunks bss == lbs
  
      describe "exception handling" $ do
          it "catchC" $ do
              ref <- I.newIORef 0
              let src = do
                      C.catchC (CB.sourceFile "some-file-that-does-not-exist") onErr
                      C.handleC onErr $ CB.sourceFile "conduit-extra.cabal"
                  onErr :: MonadIO m => IOException -> m ()
                  onErr _ = liftIO $ I.modifyIORef ref (+ 1)
              contents <- L.readFile "conduit-extra.cabal"
              res <- runConduitRes $ src .| CB.sinkLbs
              res `shouldBe` contents
              errCount <- I.readIORef ref
              errCount `shouldBe` (1 :: Int)
          it "tryC" $ do
              ref <- I.newIORef undefined
              let src = do
                      res1 <- C.tryC $ CB.sourceFile "some-file-that-does-not-exist"
                      res2 <- C.tryC $ CB.sourceFile "conduit-extra.cabal"
                      liftIO $ I.writeIORef ref (res1, res2)
              contents <- L.readFile "conduit-extra.cabal"
              res <- runConduitRes $ src .| CB.sinkLbs
              res `shouldBe` contents
              exc <- I.readIORef ref
              case exc :: (Either IOException (), Either IOException ()) of
                  (Left _, Right ()) ->
                      return ()
                  _ -> error $ show exc
  
      describe "normalFuseLeft" $ do
          it "does not double close conduit" $ do
              x <- runConduitRes $
                  let src = CL.sourceList ["foobarbazbin"]
                   in src .| CB.isolate 10 .| CL.head
              x `shouldBe` Just "foobarbazb"
  
      describe "Storable" $ do
          let test name func = describe name $ do
                  let test' size =
                        prop ("chunk size " ++ show size) $ \stores0 -> do
                          let src =
                                  loop (someStorables stores0)
                                where
                                  loop bs
                                      | S.null bs = return ()
                                      | otherwise = do
                                          let (x, y) = S.splitAt size bs
                                          C.yield x
                                          loop y
  
                              sink :: [SomeStorable]
                                   -> ConduitT S.ByteString Void IO ()
                              sink [] = do
                                  mw <- CB.head
                                  case mw of
                                      Nothing -> return ()
                                      Just _ -> error "trailing bytes"
                              sink (next:rest) = do
                                  withSomeStorable next checkOne
                                  sink rest
  
                              checkOne :: (Storable a, Eq a, Show a)
                                       => a
                                       -> ConduitT S.ByteString Void IO ()
                              checkOne expected = do
                                  mactual <-
                                      if func
                                          then CB.sinkStorable
                                          else fmap Just CB.sinkStorableEx
                                  actual <-
                                      case mactual of
                                          Nothing -> error "got Nothing"
                                          Just actual -> return actual
                                  liftIO $ actual `shouldBe` expected
  
                          runConduit $ src .| sink stores0 :: IO ()
                  mapM_ test' [1, 5, 10, 100]
  
          test "sink Maybe" True
          test "sink exception" False
  
          it' "insufficient bytes are leftovers, one chunk" $ do
              let src = C.yield $ S.singleton 1
              runConduit $ src .| do
                  mactual <- CB.sinkStorable
                  liftIO $ mactual `shouldBe` (Nothing :: Maybe Int)
                  lbs <- CB.sinkLbs
                  liftIO $ lbs `shouldBe` L.singleton 1
  
          it' "insufficient bytes are leftovers, multiple chunks" $ do
              let src = do
                      C.yield $ S.singleton 1
                      C.yield $ S.singleton 2
              runConduit $ src .| do
                  mactual <- CB.sinkStorable
                  liftIO $ mactual `shouldBe` (Nothing :: Maybe Int)
                  lbs <- CB.sinkLbs
                  liftIO $ lbs `shouldBe` L.pack [1, 2]
  
  data SomeStorable where
      SomeStorable :: (Storable a, Eq a, Show a, Typeable a) => a -> SomeStorable
  instance Show SomeStorable where
      show (SomeStorable x) = show x
  instance Arbitrary SomeStorable where
      arbitrary = oneof
          [ SomeStorable <$> (arbitrary :: Gen Int)
          , SomeStorable <$> (arbitrary :: Gen Word8)
          , SomeStorable <$> (arbitrary :: Gen Double)
          ]
  
  withSomeStorable :: SomeStorable
                   -> (forall a. (Storable a, Eq a, Show a) => a -> b)
                   -> b
  withSomeStorable (SomeStorable x) f = f x
  
  someStorable :: SomeStorable -> S.ByteString
  someStorable store =
      fst $ unsafePerformIO $ createAndTrim' (size + align) start
    where
      size = withSomeStorable store sizeOf
      align = withSomeStorable store alignment
      start ptr = do
          let off = minusPtr ptr (alignPtr ptr align)
          withSomeStorable store (pokeByteOff ptr off)
          return (off, size, ())
  
  someStorables :: [SomeStorable] -> S.ByteString
  someStorables = S.concat . map someStorable
  
  it' :: String -> IO () -> Spec
  it' = it
  