marionette-1.1.0: test/ClientSpec.hs

module Main where

import Data.Aeson qualified as Aeson
import Data.Binary qualified as Binary
import Network.Simple.TCP
    ( HostPreference (Host)
    , Socket
    , accept
    , closeSock
    , listen
    , sendLazy
    )
import Network.Socket (PortNumber, socketPort)
import Test.Hspec (describe, hspec, it)
import Test.Hspec.Expectations.Lifted (shouldReturn, shouldSatisfy)
import Test.Marionette.Class (sendCommand)
import Test.Marionette.Client
    ( MarionetteTimeout (..)
    , SocketClosed (..)
    , incoming
    , runMarionetteTWith
    )
import Test.Marionette.Protocol
import UnliftIO (TQueue, async, cancel, finally, timeout, try)
import UnliftIO.STM (atomically, readTQueue)
import Prelude

withServer :: (Socket -> TQueue MarionetteMessage -> IO ()) -> (PortNumber -> IO a) -> IO a
withServer serve body =
    listen (Host "127.0.0.1") "0" \(listening, _) -> do
        port <- socketPort listening
        handling <-
            async $ accept listening \(socket, _) ->
                serve socket =<< incoming socket \_ -> pure ()
        body port `finally` cancel handling

send :: (Aeson.ToJSON a) => Socket -> a -> IO ()
send socket = sendLazy socket . Binary.encode . MarionetteMessage . Aeson.encode

recv :: TQueue MarionetteMessage -> IO (Message Command)
recv frames = do
    MarionetteMessage lbs <- atomically $ readTQueue frames
    either fail pure $ Aeson.eitherDecode lbs

greeting :: Greeting
greeting = Greeting{applicationType = "fake", marionetteProtocol = 3}

reply :: Socket -> Message Command -> Result -> IO ()
reply socket Message{messageId} result =
    send socket Message{messageId, messageContent = result}

testCommandTimeout :: Int
testCommandTimeout = 200_000

main :: IO ()
main = hspec do
    describe "incoming" do
        it "reports SocketClosed when the server hangs up" $
            withServer
                ( \socket _frames -> do
                    send socket greeting
                    closeSock socket
                )
                \port -> do
                    outcome <-
                        timeout 2_000_000 . try $
                            runMarionetteTWith
                                "127.0.0.1"
                                port
                                testCommandTimeout
                                (sendCommand @_ @Aeson.Value "noop")

                    outcome `shouldSatisfy` \case
                        Just (Left SocketClosed) -> True
                        _ -> False

    describe "sendCommand" do
        it "throws MarionetteTimeout without killing the session" $
            withServer
                ( \socket frames -> do
                    send socket greeting
                    _silent <- recv frames
                    echoed <- recv frames
                    reply socket echoed $ Right (Aeson.String "ok")
                )
                \port ->
                    runMarionetteTWith "127.0.0.1" port testCommandTimeout do
                        silent <- try $ sendCommand @_ @Aeson.Value "silent"
                        silent `shouldSatisfy` \case
                            Left MarionetteTimeout{} -> True
                            _ -> False
                        sendCommand "echo" `shouldReturn` Aeson.String "ok"

        it "reports connection loss promptly, waking a command already in flight" $
            withServer
                ( \socket frames -> do
                    send socket greeting
                    _abandoned <- recv frames
                    closeSock socket
                )
                \port -> do
                    outcome <-
                        timeout 2_000_000 . try $
                            runMarionetteTWith
                                "127.0.0.1"
                                port
                                testCommandTimeout
                                (sendCommand @_ @Aeson.Value "abandoned")

                    outcome `shouldSatisfy` \case
                        Just (Left SocketClosed) -> True
                        _ -> False

        it "surfaces a server-reported Error rather than a transport failure" $
            withServer
                ( \socket frames -> do
                    send socket greeting
                    command <- recv frames
                    reply socket command $
                        Left Error{error = "no such element", message = "nope", stacktrace = ""}
                )
                \port -> do
                    outcome <-
                        try $
                            runMarionetteTWith
                                "127.0.0.1"
                                port
                                testCommandTimeout
                                (sendCommand @_ @Aeson.Value "find")

                    outcome `shouldSatisfy` \case
                        Left Error{message = "nope"} -> True
                        _ -> False

        it "discards a reply that arrives after its command already timed out" $
            withServer
                ( \socket frames -> do
                    send socket greeting
                    slow <- recv frames
                    echoed <- recv frames -- only received once the client gives up on "slow"
                    reply socket slow $ Right (Aeson.String "too late")
                    reply socket echoed $ Right (Aeson.String "ok")
                )
                \port ->
                    runMarionetteTWith "127.0.0.1" port testCommandTimeout do
                        slow <- try $ sendCommand @_ @Aeson.Value "slow"
                        slow `shouldSatisfy` \case
                            Left MarionetteTimeout{} -> True
                            _ -> False
                        sendCommand "echo" `shouldReturn` Aeson.String "ok"