@@ -1,4 +1,5 @@
11{-# LANGUAGE DeriveDataTypeable #-}
2+{-# LANGUAGE DoAndIfThenElse #-}
23{-# LANGUAGE ScopedTypeVariables #-}
34import Common
45import Database.PostgreSQL.Simple.FromField (FromField)
@@ -8,6 +9,7 @@ import Control.Applicative
89import Control.Exception as E
910import Control.Monad
1011import Data.ByteString (ByteString)
12+import Data.IORef
1113import Data.Typeable
1214import qualified Data.ByteString as B
1315import qualified Data.Map as Map
@@ -78,6 +80,48 @@ testFold TestEnv{..} = TestCase $ do
7880 [] $ \xs (Only x) -> return (x:xs)
7981 reverse xs @?= ([1..10000] :: [Int])
8082
83+ ref <- newIORef []
84+ forEach conn "SELECT * FROM generate_series(1,?) a, generate_series(1,?) b"
85+ (100 :: Int, 50 :: Int) $ \(a :: Int, b :: Int) -> do
86+ xs <- readIORef ref
87+ writeIORef ref $! (a,b):xs
88+ xs <- readIORef ref
89+ reverse xs @?= [(a,b) | a <- [1..100], b <- [1..50]]
90+
91+ -- Make sure fold propagates our exception.
92+ ref <- newIORef []
93+ True <- expectError (== TestException) $
94+ forEach_ conn "SELECT generate_series(1,10)" $ \(Only a) ->
95+ if a == 5 then do
96+ -- Cause a SQL error to trip up CLOSE.
97+ True <- expectError isSyntaxError $
98+ execute_ conn "asdf"
99+ True <- expectError ST.isFailedTransactionError $
100+ (query_ conn "SELECT 1" :: IO [(Only Int)])
101+ throwIO TestException
102+ else do
103+ xs <- readIORef ref
104+ writeIORef ref $! (a :: Int) : xs
105+ xs <- readIORef ref
106+ reverse xs @?= [1..4]
107+
108+ withTransaction conn $ replicateM_ 2 $ do
109+ xs <- fold_ conn "VALUES (1), (2), (3), (4), (5)"
110+ [] $ \xs (Only x) -> return (x:xs)
111+ reverse xs @?= ([1..5] :: [Int])
112+
113+ ref <- newIORef []
114+ forEach_ conn "SELECT generate_series(1,101)" $ \(Only a) ->
115+ forEach_ conn "SELECT generate_series(1,55)" $ \(Only b) -> do
116+ xs <- readIORef ref
117+ writeIORef ref $! (a :: Int, b :: Int) : xs
118+ xs <- readIORef ref
119+ reverse xs @?= [(a,b) | a <- [1..101], b <- [1..55]]
120+
121+ xs <- fold_ conn "SELECT 1 WHERE FALSE"
122+ [] $ \xs (Only x) -> return (x:xs)
123+ xs @?= ([] :: [Int])
124+
81125 -- TODO: add more complete tests, e.g.:
82126 --
83127 -- * Fold in a transaction
@@ -183,18 +227,23 @@ testSavepoint TestEnv{..} = TestCase $ do
183227
184228 return ()
185229
186− where
187− expectError p io =
188− (io >> return False) `E.catch` \ex ->
189− if p ex then return True else throwIO ex
190− isUniqueViolation SqlError{..} = sqlState == "23505"
191−
192230data TestException
193231 = TestException
194232 deriving (Eq, Show, Typeable)
195233
196234instance Exception TestException
197235
236+expectError :: Exception e => (e -> Bool) -> IO a -> IO Bool
237+expectError p io =
238+ (io >> return False) `E.catch` \ex ->
239+ if p ex then return True else throwIO ex
240+
241+isUniqueViolation :: SqlError -> Bool
242+isUniqueViolation SqlError{..} = sqlState == "23505"
243+
244+isSyntaxError :: SqlError -> Bool
245+isSyntaxError SqlError{..} = sqlState == "42601"
246+
198247------------------------------------------------------------------------
199248
200249-- | Action for connecting to the database that will be used for testing.