content: {-# LANGUAGE CPP #-}
  {-# LANGUAGE DeriveGeneric #-}
  {-# LANGUAGE DeriveDataTypeable #-}
  {-# LANGUAGE DoAndIfThenElse #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  import Common
  import Database.PostgreSQL.Simple.FromField (FromField)
  import Database.PostgreSQL.Simple.Types(Query(..),Values(..))
  import Database.PostgreSQL.Simple.HStore
  import Database.PostgreSQL.Simple.Copy
  import qualified Database.PostgreSQL.Simple.Transaction as ST
  
  import Control.Applicative
  import Control.Exception as E
  import Control.Monad
  import Data.Char
  import Data.List (sort)
  import Data.IORef
  import Data.Typeable
  import GHC.Generics (Generic)
  
  import Data.Aeson
  import Data.ByteString (ByteString)
  import qualified Data.ByteString as B
  import qualified Data.ByteString.Lazy.Char8 as BL
  import Data.CaseInsensitive (CI)
  import qualified Data.CaseInsensitive as CI
  import Data.Map (Map)
  import qualified Data.Map as Map
  import Data.Text(Text)
  import qualified Data.Text.Encoding as T
  import qualified Data.Vector as V
  import System.FilePath
  import System.Timeout(timeout)
  import Data.Time(getCurrentTime, diffUTCTime)
  
  import Test.Tasty
  import Test.Tasty.Golden
  import Notify
  import Serializable
  import Time
  
  tests :: TestEnv -> TestTree
  tests env = testGroup "tests"
      $ map ($ env)
      [ testBytea
      , testCase "ExecuteMany"        . testExecuteMany
      , testCase "Fold"               . testFold
      , testCase "Notify"             . testNotify
      , testCase "Serializable"       . testSerializable
      , testCase "Time"               . testTime
      , testCase "Array"              . testArray
      , testCase "Array of nullables" . testNullableArray
      , testCase "HStore"             . testHStore
      , testCase "citext"             . testCIText
      , testCase "JSON"               . testJSON
      , testCase "Savepoint"          . testSavepoint
      , testCase "Unicode"            . testUnicode
      , testCase "Values"             . testValues
      , testCase "Copy"               . testCopy
      , testCopyFailures
      , testCase "Double"             . testDouble
      , testCase "1-ary generic"      . testGeneric1
      , testCase "2-ary generic"      . testGeneric2
      , testCase "3-ary generic"      . testGeneric3
      , testCase "Timeout"            . testTimeout
      ]
  
  testBytea :: TestEnv -> TestTree
  testBytea TestEnv{..} = testGroup "Bytea"
      [ testStr "empty"                  []
      , testStr "\"hello\""              $ map (fromIntegral . fromEnum) ("hello" :: String)
      , testStr "ascending"              [0..255]
      , testStr "descending"             [255,254..0]
      , testStr "ascending, doubled up"  $ doubleUp [0..255]
      , testStr "descending, doubled up" $ doubleUp [255,254..0]
      ]
    where
      testStr label bytes = testCase label $ do
          let bs = B.pack bytes
  
          [Only h] <- query conn "SELECT md5(?::bytea)" [Binary bs]
          assertBool "Haskell -> SQL conversion altered the string" $ md5 bs == h
  
          [Only (Binary r)] <- query conn "SELECT ?::bytea" [Binary bs]
          assertBool "SQL -> Haskell conversion altered the string" $ bs == r
  
      doubleUp = concatMap (\x -> [x, x])
  
  testExecuteMany :: TestEnv -> Assertion
  testExecuteMany TestEnv{..} = do
      execute_ conn "CREATE TEMPORARY TABLE tmp_executeMany (i INT, t TEXT, b BYTEA)"
  
      let rows :: [(Int, String, Binary ByteString)]
          rows = [ (1, "hello", Binary "bye")
                 , (2, "world", Binary "\0\r\t\n")
                 , (3, "?",     Binary "")
                 ]
  
      count <- executeMany conn "INSERT INTO tmp_executeMany VALUES (?, ?, ?)" rows
      count @?= fromIntegral (length rows)
  
      rows' <- query_ conn "SELECT * FROM tmp_executeMany"
      rows' @?= rows
  
      return ()
  
  testFold :: TestEnv -> Assertion
  testFold TestEnv{..} = do
      xs <- fold_ conn "SELECT generate_series(1,10000)"
              [] $ \xs (Only x) -> return (x:xs)
      reverse xs @?= ([1..10000] :: [Int])
  
      ref <- newIORef []
      forEach conn "SELECT * FROM generate_series(1,?) a, generate_series(1,?) b"
        (100 :: Int, 50 :: Int) $ \(a :: Int, b :: Int) -> do
          xs <- readIORef ref
          writeIORef ref $! (a,b):xs
      xs <- readIORef ref
      reverse xs @?= [(a,b) | a <- [1..100], b <- [1..50]]
  
      -- Make sure fold propagates our exception.
      ref <- newIORef []
      True <- expectError (== TestException) $
                forEach_ conn "SELECT generate_series(1,10)" $ \(Only a) ->
                  if a == 5 then do
                    -- Cause a SQL error to trip up CLOSE.
                    True <- expectError isSyntaxError $
                            execute_ conn "asdf"
                    True <- expectError ST.isFailedTransactionError $
                            (query_ conn "SELECT 1" :: IO [(Only Int)])
                    throwIO TestException
                  else do
                    xs <- readIORef ref
                    writeIORef ref $! (a :: Int) : xs
      xs <- readIORef ref
      reverse xs @?= [1..4]
  
      withTransaction conn $ replicateM_ 2 $ do
          xs <- fold_ conn "VALUES (1), (2), (3), (4), (5)"
                  [] $ \xs (Only x) -> return (x:xs)
          reverse xs @?= ([1..5] :: [Int])
  
      ref <- newIORef []
      forEach_ conn "SELECT generate_series(1,101)" $ \(Only a) ->
        forEach_ conn "SELECT generate_series(1,55)" $ \(Only b) -> do
          xs <- readIORef ref
          writeIORef ref $! (a :: Int, b :: Int) : xs
      xs <- readIORef ref
      reverse xs @?= [(a,b) | a <- [1..101], b <- [1..55]]
  
      xs <- fold_ conn "SELECT 1 WHERE FALSE"
              [] $ \xs (Only x) -> return (x:xs)
      xs @?= ([] :: [Int])
  
      -- TODO: add more complete tests, e.g.:
      --
      --  * Fold in a transaction
      --
      --  * Fold in a transaction after a previous fold has been performed
      --
      --  * Nested fold
  
      return ()
  
  queryFailure :: forall a. (FromField a, Typeable a, Show a)
               => Connection -> Query -> a -> Assertion
  queryFailure conn q resultType = do
    x :: Either SomeException [Only a] <- E.try $ query_ conn q
    case x of
      Left  _   -> return ()
      Right val -> assertFailure ("Did not fail as expected:  "
                                ++ show q
                                ++ " :: "
                                ++ show (typeOf resultType)
                                ++ " -> " ++ show val)
  
  testArray :: TestEnv -> Assertion
  testArray TestEnv{..} = do
      xs <- query_ conn "SELECT '{1,2,3,4}'::_int4"
      xs @?= [Only (V.fromList [1,2,3,4 :: Int])]
      xs <- query_ conn "SELECT '{{1,2},{3,4}}'::_int4"
      xs @?= [Only (V.fromList [V.fromList [1,2],
                                V.fromList [3,4 :: Int]])]
      queryFailure conn "SELECT '{1,2,3,4}'::_int4" (undefined :: V.Vector Bool)
      queryFailure conn "SELECT '{{1,2},{3,4}}'::_int4" (undefined :: V.Vector Int)
  
  testNullableArray :: TestEnv -> Assertion
  testNullableArray TestEnv{..} = do
      xs <- query_ conn "SELECT '{sometext, \"NULL\"}'::_text"
      xs @?= [Only (V.fromList ["sometext", "NULL" :: Text])]
      xs <- query_ conn "SELECT '{sometext, NULL}'::_text"
      xs @?= [Only (V.fromList [Just "sometext", Nothing :: Maybe Text])]
      queryFailure conn "SELECT '{sometext, NULL}'::_text" (undefined :: V.Vector Text)
  
  testHStore :: TestEnv -> Assertion
  testHStore TestEnv{..} = do
      execute_ conn "CREATE EXTENSION IF NOT EXISTS hstore"
      roundTrip []
      roundTrip [("foo","bar"),("bar","baz"),("baz","hello")]
      roundTrip [("fo\"o","bar"),("b\\ar","baz"),("baz","\"value\\with\"escapes")]
    where
      roundTrip :: [(Text,Text)] -> Assertion
      roundTrip xs = do
        let m = Only (HStoreMap (Map.fromList xs))
        m' <- query conn "SELECT ?::hstore" m
        [m] @?= m'
  
  testCIText :: TestEnv -> Assertion
  testCIText TestEnv{..} = do
      execute_ conn "CREATE EXTENSION IF NOT EXISTS citext"
      roundTrip (CI.mk "")
      roundTrip (CI.mk "UPPERCASE")
      roundTrip (CI.mk "lowercase")
    where
      roundTrip :: (CI Text) -> Assertion
      roundTrip cit = do
          let toPostgres = Only cit
          fromPostgres <- query conn "SELECT ?::citext" toPostgres
          [toPostgres] @?= fromPostgres
  
  testJSON :: TestEnv -> Assertion
  testJSON TestEnv{..} = do
      roundTrip (Map.fromList [] :: Map Text Text)
      roundTrip (Map.fromList [("foo","bar"),("bar","baz"),("baz","hello")] :: Map Text Text)
      roundTrip (Map.fromList [("fo\"o","bar"),("b\\ar","baz"),("baz","\"value\\with\"escapes")] :: Map Text Text)
      roundTrip (V.fromList [1,2,3,4,5::Int])
      roundTrip ("foo" :: Text)
      roundTrip (42 :: Int)
    where
      roundTrip :: ToJSON a => a -> Assertion
      roundTrip a = do
        let js = Only (toJSON a)
        js' <- query conn "SELECT ?::json" js
        [js] @?= js'
  
  testSavepoint :: TestEnv -> Assertion
  testSavepoint TestEnv{..} = do
      True <- expectError ST.isNoActiveTransactionError $
              withSavepoint conn $ return ()
  
      let getRows :: IO [Int]
          getRows = map fromOnly <$> query_ conn "SELECT a FROM tmp_savepoint ORDER BY a"
      withTransaction conn $ do
          execute_ conn "CREATE TEMPORARY TABLE tmp_savepoint (a INT UNIQUE)"
          execute_ conn "INSERT INTO tmp_savepoint VALUES (1)"
          [1] <- getRows
  
          withSavepoint conn $ do
              execute_ conn "INSERT INTO tmp_savepoint VALUES (2)"
              [1,2] <- getRows
              return ()
          [1,2] <- getRows
  
          withSavepoint conn $ do
              execute_ conn "INSERT INTO tmp_savepoint VALUES (3)"
              [1,2,3] <- getRows
              True <- expectError isUniqueViolation $
                  execute_ conn "INSERT INTO tmp_savepoint VALUES (2)"
              True <- expectError ST.isFailedTransactionError getRows
  
              -- Body returning successfully after handling error,
              -- but 'withSavepoint' will roll back without complaining.
              return ()
          -- Rolling back clears the error condition.
          [1,2] <- getRows
  
          -- 'withSavepoint' will roll back after an exception, even if the
          -- exception wasn't SQL-related.
          True <- expectError (== TestException) $
            withSavepoint conn $ do
              execute_ conn "INSERT INTO tmp_savepoint VALUES (3)"
              [1,2,3] <- getRows
              throwIO TestException
          [1,2] <- getRows
  
          -- Nested savepoint can be rolled back while the
          -- outer effects are retained.
          withSavepoint conn $ do
              execute_ conn "INSERT INTO tmp_savepoint VALUES (3)"
              True <- expectError isUniqueViolation $
                withSavepoint conn $ do
                  execute_ conn "INSERT INTO tmp_savepoint VALUES (4)"
                  [1,2,3,4] <- getRows
                  execute_ conn "INSERT INTO tmp_savepoint VALUES (4)"
              [1,2,3] <- getRows
              return ()
          [1,2,3] <- getRows
  
          return ()
  
      -- Transaction committed successfully, even though there were errors
      -- (but we rolled them back).
      [1,2,3] <- getRows
  
      return ()
  
  testUnicode :: TestEnv -> Assertion
  testUnicode TestEnv{..} = do
      let q = Query . T.encodeUtf8  -- Handle encoding ourselves to ensure
                                    -- the table gets created correctly.
      let messages = map Only ["привет","мир"] :: [Only Text]
      execute_ conn (q "CREATE TEMPORARY TABLE ру́сский (сообщение TEXT)")
      executeMany conn "INSERT INTO ру́сский (сообщение) VALUES (?)" messages
      messages' <- query_ conn "SELECT сообщение FROM ру́сский"
      sort messages @?= sort messages'
  
  testValues :: TestEnv -> Assertion
  testValues TestEnv{..} = do
      execute_ conn "CREATE TEMPORARY TABLE values_test (x int, y text)"
      test (Values ["int4","text"] [])
      test (Values ["int4","text"] [(1,"hello")])
      test (Values ["int4","text"] [(1,"hello"),(2,"world")])
      test (Values ["int4","text"] [(1,"hello"),(2,"world"),(3,"goodbye")])
      test (Values [] [(1,"hello")])
      test (Values [] [(1,"hello"),(2,"world")])
      test (Values [] [(1,"hello"),(2,"world"),(3,"goodbye")])
    where
      test :: Values (Int, Text) -> Assertion
      test table@(Values _ vals) = do
        execute conn "INSERT INTO values_test ?" (Only table)
        vals' <- query_  conn "DELETE FROM values_test RETURNING *"
        sort vals @?= sort vals'
  
  
  testCopy :: TestEnv -> Assertion
  testCopy TestEnv{..} = do
      execute_ conn "CREATE TEMPORARY TABLE copy_test (x int, y text)"
      copy_ conn "COPY copy_test FROM STDIN (FORMAT CSV)"
      mapM_ (putCopyData conn) copyRows
      putCopyEnd conn
      copy_ conn "COPY copy_test FROM STDIN (FORMAT CSV)"
      mapM_ (putCopyData conn) abortRows
      putCopyError conn "aborted"
      -- Hmm, does postgres always produce \n as an end-of-line here, or
      -- are there cases where it will use a \r\n as well?
      copy_ conn "COPY copy_test TO STDOUT (FORMAT CSV)"
      rows <- loop []
      sort rows @?= sort copyRows
      -- Now, let's just verify that the connection state is back to ready,
      -- so that we can issue more queries:
      [Only (x::Int)] <- query_ conn "SELECT 2 + 2"
      x @?= 4
    where
      copyRows  = ["1,foo\n"
                  ,"2,bar\n"]
      abortRows = ["3,baz\n"]
      loop rows = do
        mrow <- getCopyData conn
        case mrow of
          CopyOutDone _   -> return rows
          CopyOutRow  row -> loop (row:rows)
  
  testCopyFailures :: TestEnv -> TestTree
  testCopyFailures env = testGroup "Copy failures"
      $ map ($ env)
      [ testCopyUniqueConstraintError
      , testCopyMalformedError
      ]
  
  goldenTest :: TestName -> IO BL.ByteString -> TestTree
  goldenTest testName =
      goldenVsString testName (resultsDir </> fileName<.>"expected")
    where
      resultsDir = "test" </> "results"
      fileName = map normalize testName
      normalize c | not (isAlpha c) = '-'
                  | otherwise       = c
  
  -- | Test that we provide a sensible error message on failure
  testCopyUniqueConstraintError :: TestEnv -> TestTree
  testCopyUniqueConstraintError TestEnv{..} =
      goldenTest "unique constraint violation"
      $ handle (\(SomeException exc) -> return $ BL.pack $ show exc) $ do
          execute_ conn "CREATE TEMPORARY TABLE copy_unique_constraint_error_test (x int PRIMARY KEY, y text)"
          copy_ conn "COPY copy_unique_constraint_error_test FROM STDIN (FORMAT CSV)"
          mapM_ (putCopyData conn) copyRows
          _n <- putCopyEnd conn
          return BL.empty
    where
      copyRows  = ["1,foo\n"
                  ,"2,bar\n"
                  ,"1,baz\n"]
  
  testCopyMalformedError :: TestEnv -> TestTree
  testCopyMalformedError TestEnv{..} =
      goldenTest "malformed input"
      $ handle (\(SomeException exc) -> return $ BL.pack $ show exc) $ do
          execute_ conn "CREATE TEMPORARY TABLE copy_malformed_input_error_test (x int PRIMARY KEY, y text)"
          copy_ conn "COPY copy_unique_constraint_error_test FROM STDIN (FORMAT CSV)"
          mapM_ (putCopyData conn) copyRows
          _n <- putCopyEnd conn
          return BL.empty
    where
      copyRows  = ["1,foo\n"
                  ,"2,bar\n"
                  ,"z,baz\n"]
  
  testTimeout :: TestEnv -> Assertion
  testTimeout TestEnv{..} =
      withConn $ \c -> do
        start_t <- getCurrentTime
        res <- timeout 200000 $ do
                 withTransaction c $ do
                   query_ c "SELECT pg_sleep(1)" :: IO [Only ()]
        end_t <- getCurrentTime
        assertBool "Timeout did not occur" (res == Nothing)
  #if !defined(mingw32_HOST_OS)
  -- At the moment, you cannot timely abandon queries with async exceptions on
  -- Windows.
        let d = end_t `diffUTCTime` start_t
        assertBool "Timeout didn't work in a timely fashion" (0.1 < d && d < 0.6)
  #endif
  
  testDouble :: TestEnv -> Assertion
  testDouble TestEnv{..} = do
      [Only (x :: Double)] <- query_ conn "SELECT 'NaN'::float8"
      assertBool "expected NaN" (isNaN x)
      [Only (x :: Double)] <- query_ conn "SELECT 'Infinity'::float8"
      x @?= (1 / 0)
      [Only (x :: Double)] <- query_ conn "SELECT '-Infinity'::float8"
      x @?= (-1 / 0)
  
  
  testGeneric1 :: TestEnv -> Assertion
  testGeneric1 TestEnv{..} = do
      roundTrip conn (Gen1 123)
    where
      roundTrip conn x0 = do
          r <- query conn "SELECT ?::int" (x0 :: Gen1)
          r @?= [x0]
  
  testGeneric2 :: TestEnv -> Assertion
  testGeneric2 TestEnv{..} = do
      roundTrip conn (Gen2 123 "asdf")
    where
      roundTrip conn x0 = do
          r <- query conn "SELECT ?::int, ?::text" x0
          r @?= [x0]
  
  testGeneric3 :: TestEnv -> Assertion
  testGeneric3 TestEnv{..} = do
      roundTrip conn (Gen3 123 "asdf" True)
    where
      roundTrip conn x0 = do
          r <- query conn "SELECT ?::int, ?::text, ?::bool" x0
          r @?= [x0]
  
  data Gen1 = Gen1 Int
              deriving (Show,Eq,Generic)
  instance FromRow Gen1
  instance ToRow   Gen1
  
  data Gen2 = Gen2 Int Text
              deriving (Show,Eq,Generic)
  instance FromRow Gen2
  instance ToRow   Gen2
  
  data Gen3 = Gen3 Int Text Bool
              deriving (Show,Eq,Generic)
  instance FromRow Gen3
  instance ToRow   Gen3
  
  data TestException
    = TestException
    deriving (Eq, Show, Typeable)
  
  instance Exception TestException
  
  expectError :: Exception e => (e -> Bool) -> IO a -> IO Bool
  expectError p io =
      (io >> return False) `E.catch` \ex ->
      if p ex then return True else throwIO ex
  
  isUniqueViolation :: SqlError -> Bool
  isUniqueViolation SqlError{..} = sqlState == "23505"
  
  isSyntaxError :: SqlError -> Bool
  isSyntaxError SqlError{..} = sqlState == "42601"
  
  ------------------------------------------------------------------------
  
  -- | Action for connecting to the database that will be used for testing.
  --
  -- Note that some tests, such as Notify, use multiple connections, and assume
  -- that 'testConnect' connects to the same database every time it is called.
  testConnect :: IO Connection
  testConnect = connectPostgreSQL ""
  
  withTestEnv :: (TestEnv -> IO a) -> IO a
  withTestEnv cb =
      withConn $ \conn ->
          cb TestEnv
              { conn     = conn
              , withConn = withConn
              }
    where
      withConn = bracket testConnect close
  
  main :: IO ()
  main = withTestEnv $ defaultMain . tests
  