{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- This file is part of the Wire Server implementation.
--
-- Copyright (C) 2022 Wire Swiss GmbH <opensource@wire.com>
--
-- This program is free software: you can redistribute it and/or modify it under
-- the terms of the GNU Affero General Public License as published by the Free
-- Software Foundation, either version 3 of the License, or (at your option) any
-- later version.
--
-- This program is distributed in the hope that it will be useful, but WITHOUT
-- ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
-- FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more
-- details.
--
-- You should have received a copy of the GNU Affero General Public License along
-- with this program. If not, see <https://www.gnu.org/licenses/>.

module Testlib.Cannon
  ( WebSocket (..),
    WSConnect (..),
    ToWSConnect (..),
    AwaitResult (..),
    withWebSocket,
    withWebSockets,
    awaitNMatchesResult,
    awaitNMatches,
    awaitMatch,
    awaitAnyEvent,
    awaitAtLeastNMatchesResult,
    awaitAtLeastNMatches,
    awaitNToMMatchesResult,
    awaitNToMMatches,
    assertAwaitResult,
    nPayload,
    printAwaitResult,
    printAwaitAtLeastResult,
    waitForResponse,
    assertNoEvent,
  )
where

import Control.Concurrent
import Control.Concurrent.Async
import Control.Concurrent.STM.TChan
import Control.Exception (throwIO)
import qualified Control.Exception as E
import Control.Monad
import Control.Monad.Catch hiding (bracket)
import qualified Control.Monad.Catch as Catch
import Control.Monad.IO.Class
import Control.Monad.Reader (asks)
import Control.Monad.STM
import Data.Aeson hiding ((.=))
import Data.ByteString (ByteString)
import Data.ByteString.Conversion (fromByteString)
import Data.ByteString.Conversion.To
import Data.Function
import Data.Maybe
import Data.Traversable
import Data.Word
import GHC.Records
import GHC.Stack
import qualified Network.HTTP.Client as HTTP
import qualified Network.HTTP.Client as Http
import qualified Network.WebSockets as WS
import System.Random (randomIO)
import System.Timeout (timeout)
import Testlib.App
import Testlib.Assertions
import Testlib.Env
import Testlib.HTTP
import Testlib.JSON
import Testlib.Printing
import Testlib.Types
import UnliftIO (withRunInIO)
import Prelude

data WebSocket = WebSocket
  { WebSocket -> WSConnect
wsConnect :: WSConnect,
    WebSocket -> TChan Value
wsChan :: TChan Value,
    WebSocket -> MVar ()
wsCloseLatch :: MVar (),
    WebSocket -> Async ()
wsAppThread :: Async ()
  }

instance HasField "client" WebSocket (Maybe ClientIdentity) where
  getField :: WebSocket -> Maybe ClientIdentity
getField WebSocket
ws = do
    c <- WebSocket
ws.wsConnect.client
    pure
      ClientIdentity
        { domain = ws.wsConnect.domain,
          user = ws.wsConnect.user,
          client = c
        }

instance HasField "user" WebSocket Value where
  getField :: WebSocket -> Value
getField WebSocket
ws =
    [Pair] -> Value
object
      [ String
"domain" String -> String -> Pair
forall a. ToJSON a => String -> a -> Pair
.= WebSocket
ws.wsConnect.domain,
        String
"id" String -> String -> Pair
forall a. ToJSON a => String -> a -> Pair
.= WebSocket
ws.wsConnect.user
      ]

-- Specifies how a Websocket at cannon should be opened
data WSConnect = WSConnect
  { WSConnect -> String
user :: String,
    WSConnect -> String
domain :: String,
    WSConnect -> Maybe String
client :: Maybe String,
    -- | If this is Nothing then a random Z-Connection will be used
    WSConnect -> Maybe String
conn :: Maybe String
  }

class ToWSConnect a where
  toWSConnect :: (HasCallStack) => a -> App WSConnect

instance {-# OVERLAPPING #-} ToWSConnect WSConnect where
  toWSConnect :: HasCallStack => WSConnect -> App WSConnect
toWSConnect = WSConnect -> App WSConnect
forall a. a -> App a
forall (f :: * -> *) a. Applicative f => a -> f a
pure

instance {-# OVERLAPPABLE #-} (MakesValue user) => ToWSConnect user where
  toWSConnect :: HasCallStack => user -> App WSConnect
toWSConnect user
u = do
    (domain, uid) <- user -> App (String, String)
forall a. (HasCallStack, MakesValue a) => a -> App (String, String)
objQid user
u
    mc <- lookupField u "client_id"
    c <- traverse asString mc
    pure (WSConnect uid domain c Nothing)

instance (MakesValue user, MakesValue conn) => ToWSConnect (user, conn) where
  toWSConnect :: HasCallStack => (user, conn) -> App WSConnect
toWSConnect (user
u, conn
c) = do
    (domain, uid) <- user -> App (String, String)
forall a. (HasCallStack, MakesValue a) => a -> App (String, String)
objQid user
u
    conn <- make c & asString
    pure (WSConnect uid domain Nothing (Just conn))

instance (MakesValue user, MakesValue conn, MakesValue client) => ToWSConnect (user, conn, client) where
  toWSConnect :: HasCallStack => (user, conn, client) -> App WSConnect
toWSConnect (user
u, conn
c, client
cl) = do
    (domain, uid) <- user -> App (String, String)
forall a. (HasCallStack, MakesValue a) => a -> App (String, String)
objQid user
u
    client <- make cl & asString
    conn <- make c & asString
    pure (WSConnect uid domain (Just client) (Just conn))

connect :: (HasCallStack) => WSConnect -> App WebSocket
connect :: HasCallStack => WSConnect -> App WebSocket
connect WSConnect
wsConnect = do
  nchan <- IO (TChan Value) -> App (TChan Value)
forall a. IO a -> App a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO (TChan Value)
forall a. IO (TChan a)
newTChanIO
  latch <- liftIO newEmptyMVar
  wsapp <- run wsConnect (clientApp nchan latch)
  pure $ WebSocket wsConnect nchan latch wsapp

clientApp :: (HasCallStack) => TChan Value -> MVar () -> WS.ClientApp ()
clientApp :: HasCallStack => TChan Value -> MVar () -> ClientApp ()
clientApp TChan Value
wsChan MVar ()
latch Connection
conn = do
  r <- IO (ZonkAny 2) -> IO (Async (ZonkAny 2))
forall a. IO a -> IO (Async a)
async IO (ZonkAny 2)
wsRead
  w <- async wsWrite
  void $ waitEitherCancel r w
  where
    wsRead :: IO (ZonkAny 2)
wsRead = IO () -> IO (ZonkAny 2)
forall (f :: * -> *) a b. Applicative f => f a -> f b
forever (IO () -> IO (ZonkAny 2)) -> IO () -> IO (ZonkAny 2)
forall a b. (a -> b) -> a -> b
$ do
      bs <- Connection -> IO ByteString
forall a. WebSocketsData a => Connection -> IO a
WS.receiveData Connection
conn
      case decodeStrict' bs of
        Just Value
n -> STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TChan Value -> Value -> STM ()
forall a. TChan a -> a -> STM ()
writeTChan TChan Value
wsChan Value
n
        Maybe Value
Nothing -> String -> IO ()
putStrLn (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String
"Failed to decode notification: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ ByteString -> String
forall a. Show a => a -> String
show ByteString
bs
    wsWrite :: IO (ZonkAny 1)
wsWrite = IO () -> IO (ZonkAny 1)
forall (f :: * -> *) a b. Applicative f => f a -> f b
forever (IO () -> IO (ZonkAny 1)) -> IO () -> IO (ZonkAny 1)
forall a b. (a -> b) -> a -> b
$ do
      MVar () -> IO ()
forall a. MVar a -> IO a
takeMVar MVar ()
latch
      Connection -> ByteString -> IO ()
forall a. WebSocketsData a => Connection -> a -> IO ()
WS.sendClose Connection
conn (ByteString
"close" :: ByteString)

-- | Start a client thread in 'Async' that opens a web socket to a Cannon, wait
--   for the connection to register with Gundeck, and return the 'Async' thread.
run ::
  (HasCallStack) =>
  WSConnect ->
  WS.ClientApp () ->
  App (Async ())
run :: HasCallStack => WSConnect -> ClientApp () -> App (Async ())
run WSConnect
wsConnect ClientApp ()
app = do
  domain <- String -> App String
forall a. (HasCallStack, MakesValue a) => a -> App String
asString WSConnect
wsConnect.domain
  serviceMap <- getServiceMap domain

  let HostPort caHost caPort = serviceHostPort serviceMap Cannon
  latch <- liftIO newEmptyMVar

  connId <- case wsConnect.conn of
    Just String
c -> String -> App String
forall a. a -> App a
forall (f :: * -> *) a. Applicative f => a -> f a
pure String
c
    Maybe String
Nothing -> Word32 -> String
forall a. Show a => a -> String
show (Word32 -> String) -> App Word32 -> App String
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO Word32 -> App Word32
forall a. IO a -> App a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Word32
forall a (m :: * -> *). (Random a, MonadIO m) => m a
randomIO :: IO Word32)

  let path =
        String
"/await"
          String -> String -> String
forall a. Semigroup a => a -> a -> a
<> ( case WSConnect
wsConnect.client of
                 Maybe String
Nothing -> String
""
                 Just String
client -> Maybe String -> String
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe String -> String)
-> (ByteString -> Maybe String) -> ByteString -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Maybe String
forall a. FromByteString a => ByteString -> Maybe a
fromByteString (ByteString -> String) -> ByteString -> String
forall a b. (a -> b) -> a -> b
$ Request -> ByteString
Http.queryString ([(ByteString, Maybe ByteString)] -> Request -> Request
Http.setQueryString [(ByteString
"client", ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just (String -> ByteString
forall a. ToByteString a => a -> ByteString
toByteString' String
client))] Request
Http.defaultRequest)
             )
      caHdrs =
        [ (HeaderName
"Z-User", String -> ByteString
forall a. ToByteString a => a -> ByteString
toByteString' (WSConnect
wsConnect.user)),
          (HeaderName
"Z-Connection", String -> ByteString
forall a. ToByteString a => a -> ByteString
toByteString' String
connId)
        ]
  request <- do
    r <- rawBaseRequest domain Cannon Versioned path
    pure r {HTTP.requestHeaders = caHdrs}

  wsapp <-
    liftIO
      $ async
      $ catch
        ( WS.runClientWith
            caHost
            (fromIntegral caPort)
            path
            WS.defaultConnectionOptions
            caHdrs
            app
        )
      $ \(SomeException
e :: SomeException) -> MVar SomeException -> SomeException -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar SomeException
latch SomeException
e

  presenceRequest <-
    baseRequest domain Cannon Unversioned $
      "/i/presences/" <> wsConnect.user <> "/" <> connId

  waitForPresence <- appToIO $ retryT $ do
    response <- submit "HEAD" presenceRequest
    status response `shouldMatchInt` 200
  let waitForException = do
        ex <- MVar SomeException -> IO SomeException
forall a. MVar a -> IO a
takeMVar MVar SomeException
latch
        -- Construct a "fake" response. We do not really have access to the
        -- websocket connection requests and response, unfortunately, but it is
        -- useful to display some information about the request in case an
        -- exception occurs.
        let r =
              Response
                { jsonBody :: Maybe Value
jsonBody = Maybe Value
forall a. Maybe a
Nothing,
                  body :: ByteString
body = ByteString
"This is a fake response. The actual response from cannon is not available.",
                  status :: Int
status = Int
101,
                  headers :: [Header]
headers = [Header]
forall a. Monoid a => a
mempty,
                  request :: Request
request = Request
request
                }
        throwIO (AssertionFailure callStack (Just r) Nothing (displayException ex))

  liftIO $ race_ waitForPresence waitForException
  pure wsapp

close :: (MonadIO m) => WebSocket -> m ()
close :: forall (m :: * -> *). MonadIO m => WebSocket -> m ()
close WebSocket
ws = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ do
  MVar () -> () -> IO ()
forall a. MVar a -> a -> IO ()
putMVar (WebSocket -> MVar ()
wsCloseLatch WebSocket
ws) ()
  IO (Either SomeException ()) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Either SomeException ()) -> IO ())
-> IO (Either SomeException ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ Async () -> IO (Either SomeException ())
forall a. Async a -> IO (Either SomeException a)
waitCatch (WebSocket -> Async ()
wsAppThread WebSocket
ws)

withWebSocket :: (HasCallStack, ToWSConnect w) => w -> (WebSocket -> App a) -> App a
withWebSocket :: forall w a.
(HasCallStack, ToWSConnect w) =>
w -> (WebSocket -> App a) -> App a
withWebSocket w
w WebSocket -> App a
k = do
  wsConnect <- w -> App WSConnect
forall a. (ToWSConnect a, HasCallStack) => a -> App WSConnect
toWSConnect w
w
  Catch.bracket (connect wsConnect) close k

withWebSockets :: forall a w. (HasCallStack, ToWSConnect w) => [w] -> ([WebSocket] -> App a) -> App a
withWebSockets :: forall a w.
(HasCallStack, ToWSConnect w) =>
[w] -> ([WebSocket] -> App a) -> App a
withWebSockets [w]
twcs [WebSocket] -> App a
k = do
  wcs <- [w] -> (w -> App WSConnect) -> App [WSConnect]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for [w]
twcs w -> App WSConnect
forall a. (ToWSConnect a, HasCallStack) => a -> App WSConnect
toWSConnect
  go wcs []
  where
    go :: (HasCallStack) => [WSConnect] -> [WebSocket] -> App a
    go :: HasCallStack => [WSConnect] -> [WebSocket] -> App a
go [] [WebSocket]
wss = [WebSocket] -> App a
k ([WebSocket] -> [WebSocket]
forall a. [a] -> [a]
reverse [WebSocket]
wss)
    go (WSConnect
wc : [WSConnect]
wcs) [WebSocket]
wss = WSConnect -> (WebSocket -> App a) -> App a
forall w a.
(HasCallStack, ToWSConnect w) =>
w -> (WebSocket -> App a) -> App a
withWebSocket WSConnect
wc (\WebSocket
ws -> HasCallStack => [WSConnect] -> [WebSocket] -> App a
[WSConnect] -> [WebSocket] -> App a
go [WSConnect]
wcs (WebSocket
ws WebSocket -> [WebSocket] -> [WebSocket]
forall a. a -> [a] -> [a]
: [WebSocket]
wss))

data AwaitResult = AwaitResult
  { AwaitResult -> Bool
success :: Bool,
    AwaitResult -> Int
nMatchesExpected :: Int,
    AwaitResult -> [Value]
matches :: [Value],
    AwaitResult -> [Value]
nonMatches :: [Value]
  }

data AwaitAtLeastResult = AwaitAtLeastResult
  { AwaitAtLeastResult -> Bool
success :: Bool,
    AwaitAtLeastResult -> Int
nMatchesExpectedMin :: Int,
    AwaitAtLeastResult -> Maybe Int
nMatchesExpectedMax :: Maybe Int,
    AwaitAtLeastResult -> [Value]
matches :: [Value],
    AwaitAtLeastResult -> [Value]
nonMatches :: [Value]
  }

prettyAwaitResult :: AwaitResult -> App String
prettyAwaitResult :: AwaitResult -> App String
prettyAwaitResult AwaitResult
r = do
  matchesS <- [Value] -> (Value -> App String) -> App [String]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for AwaitResult
r.matches Value -> App String
forall a. MakesValue a => a -> App String
prettyJSON
  nonMatchesS <- for r.nonMatches prettyJSON
  pure $
    "AwaitResult\n"
      <> indent
        4
        ( unlines
            [ "success: " <> show (r.success),
              "matches:\n" <> unlines matchesS,
              "non-matches:\n" <> unlines nonMatchesS
            ]
        )

prettyAwaitAtLeastResult :: AwaitAtLeastResult -> App String
prettyAwaitAtLeastResult :: AwaitAtLeastResult -> App String
prettyAwaitAtLeastResult AwaitAtLeastResult
r = do
  matchesS <- [Value] -> (Value -> App String) -> App [String]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for AwaitAtLeastResult
r.matches Value -> App String
forall a. MakesValue a => a -> App String
prettyJSON
  nonMatchesS <- for r.nonMatches prettyJSON
  pure $
    "AwaitAtLeastResult\n"
      <> indent
        4
        ( unlines
            [ "min expected:" <> show r.nMatchesExpectedMin,
              "max expected:" <> show r.nMatchesExpectedMax,
              "success: " <> show (r.success),
              "matches:\n" <> unlines matchesS,
              "non-matches:\n" <> unlines nonMatchesS
            ]
        )

printAwaitResult :: AwaitResult -> App ()
printAwaitResult :: AwaitResult -> App ()
printAwaitResult = AwaitResult -> App String
prettyAwaitResult (AwaitResult -> App String)
-> (String -> App ()) -> AwaitResult -> App ()
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> IO () -> App ()
forall a. IO a -> App a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> App ()) -> (String -> IO ()) -> String -> App ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> IO ()
putStrLn

printAwaitAtLeastResult :: AwaitAtLeastResult -> App ()
printAwaitAtLeastResult :: AwaitAtLeastResult -> App ()
printAwaitAtLeastResult = AwaitAtLeastResult -> App String
prettyAwaitAtLeastResult (AwaitAtLeastResult -> App String)
-> (String -> App ()) -> AwaitAtLeastResult -> App ()
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> IO () -> App ()
forall a. IO a -> App a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> App ()) -> (String -> IO ()) -> String -> App ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> IO ()
putStrLn

awaitAnyEvent :: Int -> WebSocket -> App (Maybe Value)
awaitAnyEvent :: Int -> WebSocket -> App (Maybe Value)
awaitAnyEvent Int
tSecs = IO (Maybe Value) -> App (Maybe Value)
forall a. IO a -> App a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe Value) -> App (Maybe Value))
-> (WebSocket -> IO (Maybe Value))
-> WebSocket
-> App (Maybe Value)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> IO Value -> IO (Maybe Value)
forall a. Int -> IO a -> IO (Maybe a)
timeout (Int
tSecs Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1000 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1000) (IO Value -> IO (Maybe Value))
-> (WebSocket -> IO Value) -> WebSocket -> IO (Maybe Value)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. STM Value -> IO Value
forall a. STM a -> IO a
atomically (STM Value -> IO Value)
-> (WebSocket -> STM Value) -> WebSocket -> IO Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TChan Value -> STM Value
forall a. TChan a -> STM a
readTChan (TChan Value -> STM Value)
-> (WebSocket -> TChan Value) -> WebSocket -> STM Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. WebSocket -> TChan Value
wsChan

-- | 'await' an expected number of notification events on the websocket that
-- satisfy the provided predicate. If there isn't any new event (matching or
-- non-matching) for a 'tSecs' seconds then AwaitResult is a failure. This
-- function will never terminate if there is a constant stream of events
-- received. When this functions returns it will push any non-matching
-- events back to the websocket.
awaitNMatchesResult ::
  (HasCallStack) =>
  -- | Number of matches
  Int ->
  -- | Selection function. Exceptions are *not* caught.
  (Value -> App Bool) ->
  WebSocket ->
  App AwaitResult
awaitNMatchesResult :: HasCallStack =>
Int -> (Value -> App Bool) -> WebSocket -> App AwaitResult
awaitNMatchesResult Int
nExpected Value -> App Bool
checkMatch WebSocket
ws = Int -> [Value] -> [Value] -> App AwaitResult
go Int
nExpected [] []
  where
    go :: Int -> [Value] -> [Value] -> App AwaitResult
go Int
0 [Value]
nonMatches [Value]
matches = do
      [Value] -> App ()
refill [Value]
nonMatches
      AwaitResult -> App AwaitResult
forall a. a -> App a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AwaitResult -> App AwaitResult) -> AwaitResult -> App AwaitResult
forall a b. (a -> b) -> a -> b
$
        AwaitResult
          { success :: Bool
success = Bool
True,
            nMatchesExpected :: Int
nMatchesExpected = Int
nExpected,
            matches :: [Value]
matches = [Value] -> [Value]
forall a. [a] -> [a]
reverse [Value]
matches,
            nonMatches :: [Value]
nonMatches = [Value] -> [Value]
forall a. [a] -> [a]
reverse [Value]
nonMatches
          }
    go Int
nLeft [Value]
nonMatches [Value]
matches = do
      tSecs <- (Env -> Int) -> App Int
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Env -> Int
timeOutSeconds
      mEvent <- awaitAnyEvent tSecs ws
      case mEvent of
        Just Value
event ->
          do
            isMatch <- Value -> App Bool
checkMatch Value
event
            if isMatch
              then go (nLeft - 1) nonMatches (event : matches)
              else go nLeft (event : nonMatches) matches
        Maybe Value
Nothing -> do
          [Value] -> App ()
refill [Value]
nonMatches
          AwaitResult -> App AwaitResult
forall a. a -> App a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AwaitResult -> App AwaitResult) -> AwaitResult -> App AwaitResult
forall a b. (a -> b) -> a -> b
$
            AwaitResult
              { success :: Bool
success = Bool
False,
                nMatchesExpected :: Int
nMatchesExpected = Int
nExpected,
                matches :: [Value]
matches = [Value] -> [Value]
forall a. [a] -> [a]
reverse [Value]
matches,
                nonMatches :: [Value]
nonMatches = [Value] -> [Value]
forall a. [a] -> [a]
reverse [Value]
nonMatches
              }
    refill :: [Value] -> App ()
refill = (Value -> App ()) -> [Value] -> App ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (IO () -> App ()
forall a. IO a -> App a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> App ()) -> (Value -> IO ()) -> Value -> App ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> (Value -> STM ()) -> Value -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TChan Value -> Value -> STM ()
forall a. TChan a -> a -> STM ()
writeTChan (WebSocket -> TChan Value
wsChan WebSocket
ws))

awaitAtLeastNMatchesResult ::
  (HasCallStack) =>
  -- | Minimum number of matches
  Int ->
  -- | Selection function. Exceptions are *not* caught.
  (Value -> App Bool) ->
  WebSocket ->
  App AwaitAtLeastResult
awaitAtLeastNMatchesResult :: HasCallStack =>
Int -> (Value -> App Bool) -> WebSocket -> App AwaitAtLeastResult
awaitAtLeastNMatchesResult Int
nExpected Value -> App Bool
checkMatch WebSocket
ws = Int -> [Value] -> [Value] -> App AwaitAtLeastResult
go Int
0 [] []
  where
    go :: Int -> [Value] -> [Value] -> App AwaitAtLeastResult
go Int
nSeen [Value]
nonMatches [Value]
matches = do
      tSecs <- (Env -> Int) -> App Int
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Env -> Int
timeOutSeconds
      mEvent <- awaitAnyEvent tSecs ws
      case mEvent of
        Just Value
event ->
          do
            isMatch <- Value -> App Bool
checkMatch Value
event
            if isMatch
              then go (nSeen + 1) nonMatches (event : matches)
              else go nSeen (event : nonMatches) matches
        Maybe Value
Nothing -> do
          [Value] -> App ()
refill [Value]
nonMatches
          AwaitAtLeastResult -> App AwaitAtLeastResult
forall a. a -> App a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AwaitAtLeastResult -> App AwaitAtLeastResult)
-> AwaitAtLeastResult -> App AwaitAtLeastResult
forall a b. (a -> b) -> a -> b
$
            AwaitAtLeastResult
              { success :: Bool
success = Int
nSeen Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
nExpected,
                nMatchesExpectedMin :: Int
nMatchesExpectedMin = Int
nExpected,
                nMatchesExpectedMax :: Maybe Int
nMatchesExpectedMax = Maybe Int
forall a. Maybe a
Nothing,
                matches :: [Value]
matches = [Value] -> [Value]
forall a. [a] -> [a]
reverse [Value]
matches,
                nonMatches :: [Value]
nonMatches = [Value] -> [Value]
forall a. [a] -> [a]
reverse [Value]
nonMatches
              }
    refill :: [Value] -> App ()
refill = (Value -> App ()) -> [Value] -> App ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (IO () -> App ()
forall a. IO a -> App a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> App ()) -> (Value -> IO ()) -> Value -> App ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> (Value -> STM ()) -> Value -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TChan Value -> Value -> STM ()
forall a. TChan a -> a -> STM ()
writeTChan (WebSocket -> TChan Value
wsChan WebSocket
ws))

awaitNToMMatchesResult ::
  (HasCallStack) =>
  -- | Minimum number of matches
  Int ->
  -- | Maximum number of matches
  Int ->
  -- | Selection function. Exceptions are *not* caught.
  (Value -> App Bool) ->
  WebSocket ->
  App AwaitAtLeastResult
awaitNToMMatchesResult :: HasCallStack =>
Int
-> Int
-> (Value -> App Bool)
-> WebSocket
-> App AwaitAtLeastResult
awaitNToMMatchesResult Int
nMin Int
nMax Value -> App Bool
checkMatch WebSocket
ws = Int -> [Value] -> [Value] -> App AwaitAtLeastResult
go Int
0 [] []
  where
    go :: Int -> [Value] -> [Value] -> App AwaitAtLeastResult
go Int
nSeen [Value]
nonMatches [Value]
matches = do
      tSecs <- (Env -> Int) -> App Int
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Env -> Int
timeOutSeconds
      mEvent <- awaitAnyEvent tSecs ws
      case mEvent of
        Just Value
event ->
          do
            isMatch <- Value -> App Bool
checkMatch Value
event
            if isMatch
              then go (nSeen + 1) nonMatches (event : matches)
              else go nSeen (event : nonMatches) matches
        Maybe Value
Nothing -> do
          [Value] -> App ()
refill [Value]
nonMatches
          AwaitAtLeastResult -> App AwaitAtLeastResult
forall a. a -> App a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AwaitAtLeastResult -> App AwaitAtLeastResult)
-> AwaitAtLeastResult -> App AwaitAtLeastResult
forall a b. (a -> b) -> a -> b
$
            AwaitAtLeastResult
              { success :: Bool
success = Int
nMin Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
nSeen Bool -> Bool -> Bool
&& Int
nSeen Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
nMax,
                nMatchesExpectedMin :: Int
nMatchesExpectedMin = Int
nMin,
                nMatchesExpectedMax :: Maybe Int
nMatchesExpectedMax = Int -> Maybe Int
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
nMax,
                matches :: [Value]
matches = [Value] -> [Value]
forall a. [a] -> [a]
reverse [Value]
matches,
                nonMatches :: [Value]
nonMatches = [Value] -> [Value]
forall a. [a] -> [a]
reverse [Value]
nonMatches
              }
    refill :: [Value] -> App ()
refill = (Value -> App ()) -> [Value] -> App ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (IO () -> App ()
forall a. IO a -> App a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> App ()) -> (Value -> IO ()) -> Value -> App ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> (Value -> STM ()) -> Value -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TChan Value -> Value -> STM ()
forall a. TChan a -> a -> STM ()
writeTChan (WebSocket -> TChan Value
wsChan WebSocket
ws))

awaitNMatches ::
  (HasCallStack) =>
  -- | Number of matches
  Int ->
  -- | Selection function. Should not throw any exceptions
  (Value -> App Bool) ->
  WebSocket ->
  App [Value]
awaitNMatches :: HasCallStack =>
Int -> (Value -> App Bool) -> WebSocket -> App [Value]
awaitNMatches Int
nExpected Value -> App Bool
checkMatch WebSocket
ws = do
  res <- HasCallStack =>
Int -> (Value -> App Bool) -> WebSocket -> App AwaitResult
Int -> (Value -> App Bool) -> WebSocket -> App AwaitResult
awaitNMatchesResult Int
nExpected Value -> App Bool
checkMatch WebSocket
ws
  withWebSocketFailureContext ws $
    assertAwaitResult res

withWebSocketFailureContext :: WebSocket -> App a -> App a
withWebSocketFailureContext :: forall a. WebSocket -> App a -> App a
withWebSocketFailureContext WebSocket
ws = String -> App a -> App a
forall a. String -> App a -> App a
addFailureContext (String
"on websocket for user: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> WebSocket
ws.wsConnect.user String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"@" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> WebSocket
ws.wsConnect.domain)

assertAwaitResult :: (HasCallStack) => AwaitResult -> App [Value]
assertAwaitResult :: HasCallStack => AwaitResult -> App [Value]
assertAwaitResult AwaitResult
res = do
  if AwaitResult
res.success
    then [Value] -> App [Value]
forall a. a -> App a
forall (f :: * -> *) a. Applicative f => a -> f a
pure AwaitResult
res.matches
    else do
      let msgHeader :: String
msgHeader = String
"Expected " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show AwaitResult
res.nMatchesExpected String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" matching events, but got " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show ([Value] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length AwaitResult
res.matches) String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"."
      details <- (String
"Details:\n" String -> String -> String
forall a. Semigroup a => a -> a -> a
<>) (String -> String) -> App String -> App String
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> AwaitResult -> App String
prettyAwaitResult AwaitResult
res
      assertFailure $ unlines [msgHeader, details]

awaitAtLeastNMatches ::
  (HasCallStack) =>
  -- | Minumum number of matches
  Int ->
  -- | Selection function. Should not throw any exceptions
  (Value -> App Bool) ->
  WebSocket ->
  App [Value]
awaitAtLeastNMatches :: HasCallStack =>
Int -> (Value -> App Bool) -> WebSocket -> App [Value]
awaitAtLeastNMatches Int
nExpected Value -> App Bool
checkMatch WebSocket
ws = do
  res <- HasCallStack =>
Int -> (Value -> App Bool) -> WebSocket -> App AwaitAtLeastResult
Int -> (Value -> App Bool) -> WebSocket -> App AwaitAtLeastResult
awaitAtLeastNMatchesResult Int
nExpected Value -> App Bool
checkMatch WebSocket
ws
  if res.success
    then pure res.matches
    else do
      let msgHeader = String
"Expected " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
nExpected String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" matching events, but got " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show ([Value] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length AwaitAtLeastResult
res.matches) String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"."
      details <- ("Details:\n" <>) <$> prettyAwaitAtLeastResult res
      assertFailure $ unlines [msgHeader, details]

awaitNToMMatches ::
  (HasCallStack) =>
  -- | Minimum Number of matches
  Int ->
  -- | Maximum Number of matches
  Int ->
  -- | Selection function. Should not throw any exceptions
  (Value -> App Bool) ->
  WebSocket ->
  App [Value]
awaitNToMMatches :: HasCallStack =>
Int -> Int -> (Value -> App Bool) -> WebSocket -> App [Value]
awaitNToMMatches Int
nMin Int
nMax Value -> App Bool
checkMatch WebSocket
ws = do
  res <- HasCallStack =>
Int
-> Int
-> (Value -> App Bool)
-> WebSocket
-> App AwaitAtLeastResult
Int
-> Int
-> (Value -> App Bool)
-> WebSocket
-> App AwaitAtLeastResult
awaitNToMMatchesResult Int
nMin Int
nMax Value -> App Bool
checkMatch WebSocket
ws
  if res.success
    then pure res.matches
    else do
      let msgHeader = String
"Expected between" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
nMin String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" to " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
nMax String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" matching events, but got " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show ([Value] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length AwaitAtLeastResult
res.matches) String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"."
      details <- ("Details:\n" <>) <$> prettyAwaitAtLeastResult res
      assertFailure $ unlines [msgHeader, details]

awaitMatch ::
  (HasCallStack) =>
  -- | Selection function. Should not throw any exceptions
  (Value -> App Bool) ->
  WebSocket ->
  App Value
awaitMatch :: HasCallStack => (Value -> App Bool) -> WebSocket -> App Value
awaitMatch Value -> App Bool
checkMatch WebSocket
ws = [Value] -> Value
forall a. HasCallStack => [a] -> a
head ([Value] -> Value) -> App [Value] -> App Value
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> HasCallStack =>
Int -> (Value -> App Bool) -> WebSocket -> App [Value]
Int -> (Value -> App Bool) -> WebSocket -> App [Value]
awaitNMatches Int
1 Value -> App Bool
checkMatch WebSocket
ws

assertNoEvent ::
  (HasCallStack) =>
  Int ->
  WebSocket ->
  App ()
assertNoEvent :: HasCallStack => Int -> WebSocket -> App ()
assertNoEvent Int
to WebSocket
ws = WebSocket -> App () -> App ()
forall a. WebSocket -> App a -> App a
withWebSocketFailureContext WebSocket
ws (App () -> App ()) -> App () -> App ()
forall a b. (a -> b) -> a -> b
$ do
  mEvent <- Int -> WebSocket -> App (Maybe Value)
awaitAnyEvent Int
to WebSocket
ws
  case mEvent of
    Just Value
event -> String -> App ()
forall a. HasCallStack => String -> App a
assertFailure (String -> App ()) -> String -> App ()
forall a b. (a -> b) -> a -> b
$ String
"Expected no event, but got: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Value -> String
forall a. Show a => a -> String
show Value
event
    Maybe Value
Nothing -> () -> App ()
forall a. a -> App a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

nPayload :: (MakesValue a) => a -> App Value
nPayload :: forall a. MakesValue a => a -> App Value
nPayload a
event = do
  payloads <- a
event a -> String -> App Value
forall a. (HasCallStack, MakesValue a) => a -> String -> App Value
%. String
"payload" App Value -> (App Value -> App [Value]) -> App [Value]
forall a b. a -> (a -> b) -> b
& App Value -> App [Value]
forall a. (HasCallStack, MakesValue a) => a -> App [Value]
asList
  assertOne payloads

-- | waits for an http response to satisfy a predicate
waitForResponse :: (HasCallStack) => App Response -> (Response -> App r) -> App r
waitForResponse :: forall r.
HasCallStack =>
App Response -> (Response -> App r) -> App r
waitForResponse App Response
act Response -> App r
p = do
  tSecs <- (Env -> Int) -> App Int
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Env -> Int
timeOutSeconds
  r <- withRunInIO \forall a. App a -> IO a
inIO ->
    Int -> IO r -> IO (Maybe r)
forall a. Int -> IO a -> IO (Maybe a)
timeout (Int
1000 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1000 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
tSecs) do
      let go :: IO r
go = do
            App r -> IO r
forall a. App a -> IO a
inIO (App Response -> (Response -> App r) -> App r
forall r.
HasCallStack =>
App Response -> (Response -> App r) -> App r
bindResponse App Response
act Response -> App r
p) IO r -> (AssertionFailure -> IO r) -> IO r
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`E.catch` \(AssertionFailure
_ :: AssertionFailure) -> do
              Int -> IO ()
threadDelay Int
1000
              IO r
go
      IO r
go
  let err = [String] -> String
unwords [String
"Expected event didn't come true after", Int -> String
forall a. Show a => a -> String
show Int
tSecs, String
"seconds"]
  maybe (assertFailure err) pure r