greskell-websocket-1.0.0.6: test/ServerTest/Connection.hs
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeApplications #-}
module ServerTest.Connection
( main
, spec
) where
import Control.Concurrent (threadDelay)
import Control.Concurrent.Async (Async, async, mapConcurrently, wait,
withAsync)
import Control.Concurrent.STM (TQueue, TVar, atomically, flushTQueue,
modifyTVar, newEmptyTMVarIO,
newTQueueIO, newTVarIO, putTMVar,
readTVar, takeTMVar, writeTQueue)
import Control.Exception.Safe (Exception, SomeException, bracket,
throwString, withException)
import Control.Monad (forM_, forever, mapM, when)
import Data.Aeson (FromJSON (..), Object,
ToJSON (toJSON), Value (Number))
import qualified Data.Aeson as Aeson
import qualified Data.Aeson.KeyMap as KM
import qualified Data.Aeson.Types as Aeson (parseEither)
import qualified Data.ByteString.Lazy as BSL
import Data.Greskell.GraphSON (GraphSON, gsonValue)
import Data.Maybe (fromJust, isJust, isNothing)
import Data.Monoid ((<>))
import Data.Text (Text, pack)
import Data.UUID (UUID)
import qualified Data.UUID as UUID
import Data.UUID.V4 (nextRandom)
import qualified Data.Vector as V
import GHC.Records (HasField (..))
import qualified Network.WebSockets as WS
import System.IO (hPutStrLn, stderr)
import Test.Hspec
import Network.Greskell.WebSocket.Connection (Connection, GeneralException (..),
Host, Port, RequestException (..),
ResponseHandle,
Settings (concurrency, onGeneralException, requestQueueSize, responseTimeout),
close, connect, defJSONSettings,
nextResponse, sendRequest,
sendRequest', slurpResponses)
import Network.Greskell.WebSocket.Request (RequestMessage (requestId),
makeRequestMessage, toRequestMessage)
import Network.Greskell.WebSocket.Request.Standard (OpEval (..))
import Network.Greskell.WebSocket.Response (ResponseCode (..),
ResponseMessage (requestId, result, status),
ResponseResult (resultData),
ResponseStatus (code))
import TestUtil.Env (withEnvForExtServer,
withEnvForIntServer)
import TestUtil.MockServer (receiveRequest, simpleRawResponse,
waitForServer, wsServer)
import qualified TestUtil.TCounter as TCounter
main :: IO ()
main = hspec spec
spec :: Spec
spec = describe "Connection" $ do
no_external_server_spec
withEnvForExtServer $ do
conn_basic_spec
conn_error_spec
conn_close_spec
withEnvForIntServer $ do
conn_bad_server_spec
ourSettings :: Settings Value
ourSettings = defJSONSettings
{ onGeneralException = \e -> error (show e)
-- basically we don't expect any GeneralException.
}
withConn :: (Connection Value -> IO a) -> (Host, Port) -> IO a
withConn act (host, port) = forConn host port act
forConn :: Host -> Port -> (Connection Value -> IO a) -> IO a
forConn = forConn' defJSONSettings
forConn' :: Settings Value -> Host -> Port -> (Connection Value -> IO a) -> IO a
forConn' settings host port act = bracket makeConn close act
where
makeConn = connect settings host port
parseValue :: FromJSON a => Value -> Either String a
parseValue v = Aeson.parseEither parseJSON v
opEval :: Text -> OpEval
opEval g = OpEval { batchSize = Nothing,
gremlin = g,
bindings = Nothing,
language = Nothing,
aliases = Nothing,
scriptEvaluationTimeout = Nothing
}
---- 'resultData' should be parsed with GraphSON wrappers to support v1, v2 and v3.
slurpParseEval :: FromJSON a => ResponseHandle Value -> IO [ResponseMessage (Either String (GraphSON [GraphSON a]))]
slurpParseEval rh = (fmap . fmap . fmap) parseValue $ fmap V.toList $ slurpResponses rh
responseValues :: ResponseMessage (Either String (GraphSON [GraphSON a])) -> Either String [a]
responseValues = fmap (map gsonValue . gsonValue) . resultData . result
slurpEvalValues :: FromJSON a => ResponseHandle Value -> IO [Either String [a]]
slurpEvalValues rh = (fmap . map) responseValues $ slurpParseEval rh
opSleep :: Int -> OpEval
opSleep time_ms = opSleep' time_ms time_ms
opSleep' :: Int -> Int -> OpEval
opSleep' time_ms val = opEval ("sleep " <> tshow time_ms <> "; " <> tshow val)
where
tshow = pack . show
inspectException :: IO a -> IO a
inspectException act = act `withException` printE
where
printE :: SomeException -> IO ()
printE e = hPutStrLn stderr ("[DEBUG] exception thrown: " ++ show e)
no_external_server_spec :: Spec
no_external_server_spec = describe "connect" $ do
it "should throw exception on failure" $ do
let act = connect ourSettings "this.should.not.be.real.server.com" 8000
expectSomeEx :: SomeException -> Bool
expectSomeEx _ = True
inspectException act `shouldThrow` expectSomeEx
conn_basic_spec :: SpecWith (Host, Port)
conn_basic_spec = do
specify "basic transaction" $ withConn $ \conn -> do
rid <- nextRandom
let op = opEval "123"
exp_val :: [Int]
exp_val = [123]
res <- sendRequest' conn $ toRequestMessage rid op
got <- slurpParseEval res
map (getField @"requestId") got `shouldBe` [rid]
map (code . status) got `shouldBe` [Success]
map responseValues got `shouldBe` [Right exp_val]
specify "continuous response with bindings" $ withConn $ \conn -> do
rid <- nextRandom
let op = (opEval "x") { batchSize = Just 2,
bindings = Just $ KM.fromList [("x", toJSON ([1 .. 10] :: [Int]))]
}
exp_vals :: [Either String [Int]]
exp_vals = map Right [[1,2], [3,4], [5,6], [7,8], [9,10]]
got <- slurpParseEval =<< (sendRequest' conn $ toRequestMessage rid op)
map (getField @"requestId") got `shouldBe` replicate 5 rid
map (code . status) got `shouldBe` ((replicate 4 PartialContent) ++ [Success])
map responseValues got `shouldBe` exp_vals
specify "concurrent requests" $ withConn $ \conn -> do
handles <- mapM (sendRequest conn) $ map opSleep $ map (* 100) $ reverse [1..5]
got <- mapM slurpEvalValues handles :: IO [[Either String [Int]]]
got `shouldBe` [ [Right [500]],
[Right [400]],
[Right [300]],
[Right [200]],
[Right [100]]
]
specify "make requests from multiple threads" $ withConn $ \conn -> do
let reqAndRes val = slurpEvalValues =<< (sendRequest conn $ opSleep' 200 val)
got <- mapConcurrently reqAndRes [1 .. 5]
:: IO [[Either String [Int]]]
got `shouldBe` [ [Right [1]],
[Right [2]],
[Right [3]],
[Right [4]],
[Right [5]]
]
specify "requestId should be cleared from the request pool once completed" $ withConn $ \conn -> do
req <- makeRequestMessage $ opSleep 30
let evalReq = slurpEvalValues =<< sendRequest' conn req :: IO [Either String [Int]]
evalReq `shouldReturn` [Right [30]]
evalReq `shouldReturn` [Right [30]]
specify "client concurrency should be bounded by 'concurrency' + 'requestQueueSize'" $ \(host, port) -> do
let input_concurrency = 2
input_qreq_size = 1
settings = defJSONSettings { concurrency = input_concurrency,
requestQueueSize = input_qreq_size
}
exp_concurrency_max = input_concurrency + input_qreq_size
forConn' settings host port $ \conn -> do
tcounter <- TCounter.new
let makeReq v = TCounter.count tcounter
(sendRequest conn $ opSleep' 500 v)
(\rh -> slurpEvalValues rh :: IO [Either String [Int]])
got <- mapConcurrently makeReq [1..10]
got `shouldBe` map (\v -> [Right [v]]) [1..10]
got_hist <- TCounter.history tcounter
length got_hist `shouldBe` (2 * 10)
forM_ got_hist $ \conc -> conc `shouldSatisfy` (<= exp_concurrency_max)
specify "multiple threads reading a single ResponseHandle" $ withConn $ \conn -> do
let op = (opEval "[1,2,3,4,5,6,7,8,9,10]")
{ batchSize = Just 1
}
rh <- sendRequest conn op
mgot <- mapConcurrently (const $ nextResponse rh) ([1..15] :: [Int])
(length $ filter isNothing mgot) `shouldBe` 5
let got = map fromJust $ filter isJust mgot
got_data = map (responseValues . fmap parseValue) got :: [Either String [Int]]
map (code . status) got `shouldMatchList` (Success : replicate 9 PartialContent)
got_data `shouldMatchList` map (\v -> Right [v]) [1..10]
conn_error_spec :: SpecWith (Host, Port)
conn_error_spec = do
specify "duplicate requests" $ withConn $ \conn -> do
req <- makeRequestMessage $ opSleep 300
let expEx (DuplicateRequestId got_rid) = got_rid == getField @"requestId" req
expEx _ = False
ok_rh <- sendRequest' conn req
ng_rh <- sendRequest' conn req
ok_res <- slurpEvalValues ok_rh :: IO [Either String [Int]]
ok_res `shouldBe` [Right [300]]
nextResponse ng_rh `shouldThrow` expEx
specify "request timeout" $ \(host, port) -> do
let settings = ourSettings { responseTimeout = 1 }
expEx ResponseTimeout = True
expEx _ = False
forConn' settings host port $ \conn -> do
rh <- sendRequest conn $ opSleep 2000
(nextResponse rh) `shouldThrow` expEx
conn_close_spec :: SpecWith (Host, Port)
conn_close_spec = describe "close" $ do
it "should wait for a pending request to finish" $ withConn $ \conn -> do
rh <- sendRequest conn $ opSleep 100
close conn
got <- slurpEvalValues rh :: IO [Either String [Int]]
got `shouldBe` [Right [100]]
it "should block for all pending requests" $ withConn $ \conn -> do
tc <- TCounter.new
let makeReq :: Int -> IO [Either String [Int]]
makeReq t = do
ret <- TCounter.count tc (sendRequest conn $ opSleep t) slurpEvalValues
close conn -- more than one close don't hurt anybody (after all requests are sent.)
return ret
req_threads <- mapM (async . makeReq) [200, 400, 600]
TCounter.waitFor tc (== 3)
close conn
TCounter.now tc `shouldReturn` 0
got_hist <- TCounter.history tc
length got_hist `shouldBe` 6
got <- mapM wait req_threads
got `shouldBe` map (\v -> [Right [v]]) [200, 400, 600]
it "should make nextResponse throw AlreadyClosed exception. nextResponse should throw it every time it's called" $ withConn $ \conn -> do
let expEx AlreadyClosed = True
expEx _ = False
close conn
rh <- sendRequest conn $ opEval "999"
nextResponse rh `shouldThrow` expEx
nextResponse rh `shouldThrow` expEx
nextResponse rh `shouldThrow` expEx
succUUID :: UUID -> UUID
succUUID orig = UUID.fromWords a b c d'
where
(a, b, c ,d) = UUID.toWords orig
d' = succ d
conn_bad_server_spec :: SpecWith Port
conn_bad_server_spec = do
describe "ResponseHandle" $ describe "nextResponse" $ do
it "should throw exception when the server closes the connection while there is a pending request" $ \port -> do
let server = wsServer port $ \wsconn -> do
_ <- WS.receiveDataMessage wsconn
throwString ( "Server connection abort. Seeing this message in a test console is OK."
++ " It's because of WS.runServer internals."
)
exp_ex ServerClosed = True
exp_ex _ = False
withAsync server $ \_ -> do
waitForServer
forConn "localhost" port $ \conn -> do
rh <- sendRequest conn $ opEval "100"
(inspectException $ nextResponse rh) `shouldThrow` exp_ex
it "should throw exception when the server sends Close request while there is a pending request" $ \port -> do
let server = wsServer port $ \wsconn -> do
_ <- WS.receiveDataMessage wsconn
WS.sendClose wsconn ("" :: Text)
(WS.ControlMessage (WS.Close wscode _)) <- WS.receive wsconn
when (wscode /= 1000) $ do
throwString ( "Fatal: expects WebSocket status code 1000, but got "
++ show wscode
)
exp_ex ServerClosed = True
exp_ex _ = False
withAsync server $ \_ -> do
waitForServer
forConn "localhost" port $ \conn -> do
rh <- sendRequest conn $ opEval "100"
(inspectException $ nextResponse rh) `shouldThrow` exp_ex
it "should be ok if the server actively closes the connection" $ \port -> do
let server = wsServer port $ \wsconn -> do
req <- receiveRequest wsconn
WS.sendBinaryData wsconn $ simpleRawResponse (getField @"requestId" req) 200 "[99]"
withAsync server $ \_ -> do
waitForServer
forConn "localhost" port $ \conn -> do
rh <- sendRequest conn $ opEval "99"
got <- slurpEvalValues rh :: IO [Either String [Int]]
got `shouldBe` [Right [99]]
it "should throw AlreadyClosed exception after the server actively closes the connection" $ \port -> do
let server = wsServer port $ \wsconn -> do
threadDelay 10000
WS.sendClose wsconn ("" :: Text)
expEx AlreadyClosed = True
expEx _ = False
withAsync server $ \_ -> do
waitForServer
forConn "localhost" port $ \conn -> do
threadDelay 40000
rh <- sendRequest conn $ opEval "256"
nextResponse rh `shouldThrow` expEx
describe "Settings" $ describe "onGeneralException" $ do
it "should be called on unexpected requestId" $ \port -> do
report_gex <- newEmptyTMVarIO
let server = wsServer port $ \wsconn -> do
req <- receiveRequest wsconn
let res_id = succUUID $ getField @"requestId" req -- deliberately send wrong requestId.
res = simpleRawResponse res_id 200 "[333]"
WS.sendBinaryData wsconn res
reportEx ex = atomically $ putTMVar report_gex ex
settings = defJSONSettings { onGeneralException = reportEx }
withAsync server $ \_ -> do
waitForServer
forConn' settings "localhost" port $ \conn -> do
req <- makeRequestMessage $ opEval "333"
let exp_res_id = succUUID $ getField @"requestId" req
_ <- sendRequest' conn req
got <- atomically $ takeTMVar report_gex
got `shouldBe` UnexpectedRequestId exp_res_id
it "should be called on failure to parse response" $ \port -> do
report_gex <- newEmptyTMVarIO
let server = wsServer port $ \wsconn -> do
_ <- receiveRequest wsconn
WS.sendBinaryData wsconn ("hoge hoge hoge" :: Text)
settings = defJSONSettings { onGeneralException = \e -> atomically $ putTMVar report_gex e }
withAsync server $ \_ -> do
waitForServer
forConn' settings "localhost" port $ \conn -> do
let expEx (ResponseParseFailure _) = True
expEx _ = False
_ <- sendRequest conn $ opEval "100"
got <- atomically $ takeTMVar report_gex
got `shouldSatisfy` expEx
it "should throw ResponseTimeout exception when the server endlessly sends responses" $ \port -> do
let server = wsServer port $ \wsconn -> do
req <- receiveRequest wsconn
let res_id = getField @"requestId" req
sendContinuousRes wsconn res_id 20
sendContinuousRes :: WS.Connection -> UUID -> Int -> IO ()
sendContinuousRes wsconn res_id n = do
let finish = n == 0
threadDelay 200000
WS.sendBinaryData wsconn $ simpleRawResponse res_id (if finish then 200 else 206) "[1]"
if finish
then return ()
else sendContinuousRes wsconn res_id (n - 1)
settings = ourSettings { responseTimeout = 1 }
expEx ResponseTimeout = True
expEx _ = False
withAsync server $ \_ -> do
waitForServer
forConn' settings "localhost" port $ \conn -> do
rh <- sendRequest conn $ opEval "99"
forM_ ([1..3] :: [Int]) $ \_ -> do
mgot <- (fmap . fmap) (responseValues . fmap parseValue) $ nextResponse rh :: IO (Maybe (Either String [Int]))
mgot `shouldBe` (Just $ Right [1])
slurpResponses rh `shouldThrow` expEx