agent: rename functions, remove unused code, stabilize tests (#830)

* agent: rename functions, stabilize tests

* test pattern match

src/Simplex/FileTransfer/Agent.hs

@@ -11,15 +11,14 @@
 {-# LANGUAGE TypeApplications #-}
 
 module Simplex.FileTransfer.Agent
-  ( startWorkers,
+  ( startXFTPWorkers,
     closeXFTPAgent,
     toFSFilePath,
     -- Receiving files
-    receiveFile,
-    deleteRcvFile,
+    xftpReceiveFile',
+    xftpDeleteRcvFile',
     -- Sending files
-    sendFileExperimental,
-    sendFile,
+    xftpSendFile',
     deleteSndFileInternal,
     deleteSndFileRemote,
   )
@@ -30,16 +29,10 @@ import Control.Logger.Simple (logError)
 import Control.Monad
 import Control.Monad.Except
 import Control.Monad.Reader
-import Crypto.Random (ChaChaDRG, randomBytesGenerate)
-import Data.Bifunctor (first)
-import Data.ByteString (ByteString)
-import qualified Data.ByteString.Base64.URL as U
 import qualified Data.ByteString.Char8 as B
 import qualified Data.ByteString.Lazy.Char8 as LB
-import Data.Composition ((.:))
 import Data.Int (Int64)
-import Data.List (foldl', isSuffixOf, partition, sortOn)
-import Data.List.NonEmpty (nonEmpty)
+import Data.List (foldl', sortOn)
 import qualified Data.List.NonEmpty as L
 import Data.Map (Map)
 import qualified Data.Map.Strict as M
@@ -50,7 +43,7 @@ import Simplex.FileTransfer.Client (XFTPChunkSpec (..))
 import Simplex.FileTransfer.Client.Main
 import Simplex.FileTransfer.Crypto
 import Simplex.FileTransfer.Description
-import Simplex.FileTransfer.Protocol (FileParty (..), FilePartyI, SFileParty (..))
+import Simplex.FileTransfer.Protocol (FileParty (..), SFileParty (..))
 import Simplex.FileTransfer.Transport (XFTPRcvChunkSpec (..))
 import Simplex.FileTransfer.Types
 import Simplex.FileTransfer.Util (removePath, uniqueCombine)
@@ -62,18 +55,16 @@ import Simplex.Messaging.Agent.Store.SQLite
 import qualified Simplex.Messaging.Crypto as C
 import qualified Simplex.Messaging.Crypto.Lazy as LC
 import Simplex.Messaging.Encoding
-import Simplex.Messaging.Encoding.String
-import Simplex.Messaging.Protocol (EntityId, XFTPServer, XFTPServerWithAuth)
+import Simplex.Messaging.Protocol (EntityId, XFTPServer)
 import Simplex.Messaging.TMap (TMap)
 import qualified Simplex.Messaging.TMap as TM
-import Simplex.Messaging.Util (liftError, liftIOEither, tshow, unlessM, whenM)
+import Simplex.Messaging.Util (liftError, tshow, unlessM, whenM)
 import System.FilePath (takeFileName, (</>))
 import UnliftIO
-import UnliftIO.Concurrent
 import UnliftIO.Directory
 
-startWorkers :: AgentMonad m => AgentClient -> Maybe FilePath -> m ()
-startWorkers c workDir = do
+startXFTPWorkers :: AgentMonad m => AgentClient -> Maybe FilePath -> m ()
+startXFTPWorkers c workDir = do
   wd <- asks $ xftpWorkDir . xftpAgent
   atomically $ writeTVar wd workDir
   startRcvFiles
@@ -108,8 +99,8 @@ closeXFTPAgent XFTPAgent {xftpRcvWorkers, xftpSndWorkers} = do
       ws <- atomically $ stateTVar wsSel (,M.empty)
       mapM_ (uninterruptibleCancel . snd) ws
 
-receiveFile :: AgentMonad m => AgentClient -> UserId -> ValidFileDescription 'FRecipient -> m RcvFileId
-receiveFile c userId (ValidFileDescription fd@FileDescription {chunks}) = do
+xftpReceiveFile' :: AgentMonad m => AgentClient -> UserId -> ValidFileDescription 'FRecipient -> m RcvFileId
+xftpReceiveFile' c userId (ValidFileDescription fd@FileDescription {chunks}) = do
   g <- asks idsDrg
   prefixPath <- getPrefixPath "rcv.xftp"
   createDirectory prefixPath
@@ -274,8 +265,8 @@ runXFTPRcvLocalWorker c doWork = do
         getChunkPaths (RcvFileChunk {chunkTmpPath = Nothing} : _cs) =
           throwError $ INTERNAL "no chunk path"
 
-deleteRcvFile :: AgentMonad m => AgentClient -> RcvFileId -> m ()
-deleteRcvFile c rcvFileEntityId = do
+xftpDeleteRcvFile' :: AgentMonad m => AgentClient -> RcvFileId -> m ()
+xftpDeleteRcvFile' c rcvFileEntityId = do
   RcvFile {rcvFileId, prefixPath, status} <- withStore c $ \db -> getRcvFileByEntityId db rcvFileEntityId
   if status == RFSComplete || status == RFSError
     then do
@@ -283,66 +274,11 @@ deleteRcvFile c rcvFileEntityId = do
       withStore' c (`deleteRcvFile'` rcvFileId)
     else withStore' c (`updateRcvFileDeleted` rcvFileId)
 
-sendFileExperimental :: forall m. AgentMonad m => AgentClient -> UserId -> FilePath -> Int -> m SndFileId
-sendFileExperimental c@AgentClient {xftpServers} userId filePath numRecipients = do
-  g <- asks idsDrg
-  sndFileId <- liftIO $ randomId g 12
-  xftpSrvs <- atomically $ TM.lookup userId xftpServers
-  void $ forkIO $ sendCLI sndFileId $ maybe [] L.toList xftpSrvs
-  pure sndFileId
-  where
-    randomId :: TVar ChaChaDRG -> Int -> IO ByteString
-    randomId gVar n = U.encode <$> (atomically . stateTVar gVar $ randomBytesGenerate n)
-    sendCLI :: SndFileId -> [XFTPServerWithAuth] -> m ()
-    sendCLI sndFileId xftpSrvs = do
-      let fileName = takeFileName filePath
-      workPath <- getXFTPWorkPath
-      outputDir <- uniqueCombine workPath $ fileName <> ".descr"
-      createDirectory outputDir
-      let tempPath = workPath </> "snd"
-      createDirectoryIfMissing False tempPath
-      runSend fileName outputDir tempPath `catchAgentError` \e -> do
-        cleanup outputDir tempPath
-        notify c sndFileId $ SFERR e
-      where
-        runSend :: String -> FilePath -> FilePath -> m ()
-        runSend fileName outputDir tempPath = do
-          let sendOptions =
-                SendOptions
-                  { filePath,
-                    outputDir = Just outputDir,
-                    numRecipients,
-                    xftpServers = xftpSrvs,
-                    retryCount = 3,
-                    tempPath = Just tempPath,
-                    verbose = False
-                  }
-          liftCLI $ cliSendFileOpts sendOptions False $ notify c sndFileId .: SFPROG
-          (sndDescr, rcvDescrs) <- readDescrs outputDir fileName
-          cleanup outputDir tempPath
-          notify c sndFileId $ SFDONE sndDescr rcvDescrs
-        cleanup :: FilePath -> FilePath -> m ()
-        cleanup outputDir tempPath = do
-          removePath tempPath
-          removePath outputDir
-    liftCLI :: ExceptT CLIError IO () -> m ()
-    liftCLI = either (throwError . INTERNAL . show) pure <=< liftIO . runExceptT
-    readDescrs :: FilePath -> FilePath -> m (ValidFileDescription 'FSender, [ValidFileDescription 'FRecipient])
-    readDescrs outDir fileName = do
-      let descrDir = outDir </> (fileName <> ".xftp")
-      files <- listDirectory descrDir
-      let (sdFiles, rdFiles) = partition ("snd.xftp.private" `isSuffixOf`) files
-          sdFile = maybe "" (\l -> descrDir </> L.head l) (nonEmpty sdFiles)
-          rdFiles' = map (descrDir </>) rdFiles
-      (,) <$> readDescr sdFile <*> mapM readDescr rdFiles'
-    readDescr :: FilePartyI p => FilePath -> m (ValidFileDescription p)
-    readDescr f = liftIOEither $ first INTERNAL . strDecode <$> B.readFile f
-
 notify :: forall m e. (MonadUnliftIO m, AEntityI e) => AgentClient -> EntityId -> ACommand 'Agent e -> m ()
 notify c entId cmd = atomically $ writeTBQueue (subQ c) ("", entId, APC (sAEntity @e) cmd)
 
-sendFile :: AgentMonad m => AgentClient -> UserId -> FilePath -> Int -> m SndFileId
-sendFile c userId filePath numRecipients = do
+xftpSendFile' :: AgentMonad m => AgentClient -> UserId -> FilePath -> Int -> m SndFileId
+xftpSendFile' c userId filePath numRecipients = do
   g <- asks idsDrg
   prefixPath <- getPrefixPath "snd.xftp"
   createDirectory prefixPath

src/Simplex/FileTransfer/Client/Main.hs

@@ -417,10 +417,10 @@ getChunkDigest XFTPChunkSpec {filePath = chunkPath, chunkOffset, chunkSize} =
 
 cliReceiveFile :: ReceiveOptions -> ExceptT CLIError IO ()
 cliReceiveFile ReceiveOptions {fileDescription, filePath, retryCount, tempPath, verbose, yes} =
-  getFileDescription' fileDescription >>= receiveFile
+  getFileDescription' fileDescription >>= receive
   where
-    receiveFile :: ValidFileDescription 'FRecipient -> ExceptT CLIError IO ()
-    receiveFile (ValidFileDescription FileDescription {size, digest, key, nonce, chunks}) = do
+    receive :: ValidFileDescription 'FRecipient -> ExceptT CLIError IO ()
+    receive (ValidFileDescription FileDescription {size, digest, key, nonce, chunks}) = do
       encPath <- getEncPath tempPath "xftp"
       createDirectory encPath
       a <- atomically $ newXFTPAgent defaultXFTPClientAgentConfig

src/Simplex/Messaging/Agent.hs

@@ -122,7 +122,7 @@ import Data.Text (Text)
 import qualified Data.Text as T
 import Data.Time.Clock
 import Data.Time.Clock.System (systemToUTCTime)
-import Simplex.FileTransfer.Agent (closeXFTPAgent, deleteRcvFile, deleteSndFileInternal, deleteSndFileRemote, receiveFile, sendFile, startWorkers, toFSFilePath)
+import Simplex.FileTransfer.Agent (closeXFTPAgent, xftpDeleteRcvFile', deleteSndFileInternal, deleteSndFileRemote, xftpReceiveFile', xftpSendFile', startXFTPWorkers, toFSFilePath)
 import Simplex.FileTransfer.Description (ValidFileDescription)
 import Simplex.FileTransfer.Protocol (FileParty (..))
 import Simplex.FileTransfer.Util (removePath)
@@ -351,19 +351,19 @@ toggleConnectionNtfs :: AgentErrorMonad m => AgentClient -> ConnId -> Bool -> m
 toggleConnectionNtfs c = withAgentEnv c .: toggleConnectionNtfs' c
 
 xftpStartWorkers :: AgentErrorMonad m => AgentClient -> Maybe FilePath -> m ()
-xftpStartWorkers c = withAgentEnv c . startWorkers c
+xftpStartWorkers c = withAgentEnv c . startXFTPWorkers c
 
 -- | Receive XFTP file
 xftpReceiveFile :: AgentErrorMonad m => AgentClient -> UserId -> ValidFileDescription 'FRecipient -> m RcvFileId
-xftpReceiveFile c = withAgentEnv c .: receiveFile c
+xftpReceiveFile c = withAgentEnv c .: xftpReceiveFile' c
 
 -- | Delete XFTP rcv file (deletes work files from file system and db records)
 xftpDeleteRcvFile :: AgentErrorMonad m => AgentClient -> RcvFileId -> m ()
-xftpDeleteRcvFile c = withAgentEnv c . deleteRcvFile c
+xftpDeleteRcvFile c = withAgentEnv c . xftpDeleteRcvFile' c
 
 -- | Send XFTP file
 xftpSendFile :: AgentErrorMonad m => AgentClient -> UserId -> FilePath -> Int -> m SndFileId
-xftpSendFile c = withAgentEnv c .:. sendFile c
+xftpSendFile c = withAgentEnv c .:. xftpSendFile' c
 
 -- | Delete XFTP snd file internally (deletes work files from file system and db records)
 xftpDeleteSndFileInternal :: AgentErrorMonad m => AgentClient -> SndFileId -> m ()

tests/AgentTests/FunctionalAPITests.hs

@@ -339,10 +339,19 @@ testServerMatrix2 t runTest = do
   it "2 servers" $ withSmpServer t . withSmpServerOn t testPort2 $ runTest initAgentServers2
 
 runTestCfg2 :: AgentConfig -> AgentConfig -> AgentMsgId -> (AgentClient -> AgentClient -> AgentMsgId -> IO ()) -> IO ()
-runTestCfg2 aliceCfg bobCfg baseMsgId runTest = do
-  alice <- getSMPAgentClient' aliceCfg initAgentServers testDB
-  bob <- getSMPAgentClient' bobCfg initAgentServers testDB2
-  runTest alice bob baseMsgId
+runTestCfg2 aCfg bCfg baseMsgId runTest =
+  withAgentClientsCfg2 aCfg bCfg $ \a b -> runTest a b baseMsgId
+
+withAgentClientsCfg2 :: AgentConfig -> AgentConfig -> (AgentClient -> AgentClient -> IO ()) -> IO ()
+withAgentClientsCfg2 aCfg bCfg runTest = do
+  a <- getSMPAgentClient' aCfg initAgentServers testDB
+  b <- getSMPAgentClient' bCfg initAgentServers testDB2
+  runTest a b
+  disconnectAgentClient a
+  disconnectAgentClient b
+
+withAgentClients2 :: (AgentClient -> AgentClient -> IO ()) -> IO ()
+withAgentClients2 = withAgentClientsCfg2 agentCfg agentCfg
 
 runAgentClientTest :: HasCallStack => AgentClient -> AgentClient -> AgentMsgId -> IO ()
 runAgentClientTest alice bob baseId = do
@@ -425,10 +434,8 @@ noMessages c err = tryGet `shouldReturn` ()
         _ -> return ()
 
 testAsyncInitiatingOffline :: HasCallStack => IO ()
-testAsyncInitiatingOffline = do
-  alice <- getSMPAgentClient' agentCfg initAgentServers testDB
-  bob <- getSMPAgentClient' agentCfg initAgentServers testDB2
-  runRight_ $ do
+testAsyncInitiatingOffline =
+  withAgentClients2 $ \alice bob -> runRight_ $ do
     (bobId, cReq) <- createConnection alice 1 True SCMInvitation Nothing
     disconnectAgentClient alice
     aliceId <- joinConnection bob 1 True cReq "bob's connInfo"
@@ -442,10 +449,8 @@ testAsyncInitiatingOffline = do
     exchangeGreetings alice' bobId bob aliceId
 
 testAsyncJoiningOfflineBeforeActivation :: HasCallStack => IO ()
-testAsyncJoiningOfflineBeforeActivation = do
-  alice <- getSMPAgentClient' agentCfg initAgentServers testDB
-  bob <- getSMPAgentClient' agentCfg initAgentServers testDB2
-  runRight_ $ do
+testAsyncJoiningOfflineBeforeActivation =
+  withAgentClients2 $ \alice bob -> runRight_ $ do
     (bobId, qInfo) <- createConnection alice 1 True SCMInvitation Nothing
     aliceId <- joinConnection bob 1 True qInfo "bob's connInfo"
     disconnectAgentClient bob
@@ -459,10 +464,8 @@ testAsyncJoiningOfflineBeforeActivation = do
     exchangeGreetings alice bobId bob' aliceId
 
 testAsyncBothOffline :: HasCallStack => IO ()
-testAsyncBothOffline = do
-  alice <- getSMPAgentClient' agentCfg initAgentServers testDB
-  bob <- getSMPAgentClient' agentCfg initAgentServers testDB2
-  runRight_ $ do
+testAsyncBothOffline =
+  withAgentClients2 $ \alice bob -> runRight_ $ do
     (bobId, cReq) <- createConnection alice 1 True SCMInvitation Nothing
     disconnectAgentClient alice
     aliceId <- joinConnection bob 1 True cReq "bob's connInfo"
@@ -479,9 +482,7 @@ testAsyncBothOffline = do
     exchangeGreetings alice' bobId bob' aliceId
 
 testAsyncServerOffline :: HasCallStack => ATransport -> IO ()
-testAsyncServerOffline t = do
-  alice <- getSMPAgentClient' agentCfg initAgentServers testDB
-  bob <- getSMPAgentClient' agentCfg initAgentServers testDB2
+testAsyncServerOffline t = withAgentClients2 $ \alice bob -> do
   -- create connection and shutdown the server
   (bobId, cReq) <- withSmpServerStoreLogOn t testPort $ \_ ->
     runRight $ createConnection alice 1 True SCMInvitation Nothing
@@ -510,9 +511,7 @@ testAsyncHelloTimeout = do
   let vr11 = mkVersionRange 1 1
       smpCfgV1 = (smpCfg agentCfg) {serverVRange = vr11}
       agentCfgV1 = agentCfg {smpAgentVRange = vr11, smpClientVRange = vr11, e2eEncryptVRange = vr11, smpCfg = smpCfgV1}
-  alice <- getSMPAgentClient' agentCfgV1 initAgentServers testDB
-  bob <- getSMPAgentClient' agentCfg {helloTimeout = 1} initAgentServers testDB2
-  runRight_ $ do
+  withAgentClientsCfg2 agentCfgV1 agentCfg {helloTimeout = 1} $ \alice bob -> runRight_ $ do
     (_, cReq) <- createConnection alice 1 True SCMInvitation Nothing
     disconnectAgentClient alice
     aliceId <- joinConnection bob 1 True cReq "bob's connInfo"
@@ -555,6 +554,8 @@ testAllowConnectionClientRestart t = do
         get bob ##> ("", aliceId, CON)
 
         exchangeGreetingsMsgId 4 alice2 bobId bob aliceId
+    disconnectAgentClient alice2
+    disconnectAgentClient bob
 
 testIncreaseConnAgentVersion :: HasCallStack => ATransport -> IO ()
 testIncreaseConnAgentVersion t = do
@@ -609,6 +610,8 @@ testIncreaseConnAgentVersion t = do
       exchangeGreetingsMsgId 12 alice3 bobId bob3 aliceId
       checkVersion alice3 bobId 3
       checkVersion bob3 aliceId 3
+    disconnectAgentClient alice3
+    disconnectAgentClient bob3
 
 checkVersion :: AgentClient -> ConnId -> Version -> ExceptT AgentErrorType IO ()
 checkVersion c connId v = do
@@ -640,6 +643,8 @@ testIncreaseConnAgentVersionMaxCompatible t = do
       exchangeGreetingsMsgId 6 alice2 bobId bob2 aliceId
       checkVersion alice2 bobId 3
       checkVersion bob2 aliceId 3
+    disconnectAgentClient alice2
+    disconnectAgentClient bob2
 
 testIncreaseConnAgentVersionStartDifferentVersion :: HasCallStack => ATransport -> IO ()
 testIncreaseConnAgentVersionStartDifferentVersion t = do
@@ -663,6 +668,8 @@ testIncreaseConnAgentVersionStartDifferentVersion t = do
       exchangeGreetingsMsgId 6 alice2 bobId bob aliceId
       checkVersion alice2 bobId 3
       checkVersion bob aliceId 3
+    disconnectAgentClient alice2
+    disconnectAgentClient bob
 
 testDeliverClientRestart :: HasCallStack => ATransport -> IO ()
 testDeliverClientRestart t = do
@@ -692,6 +699,8 @@ testDeliverClientRestart t = do
 
       get bob2 ##> ("", aliceId, SENT 6)
       get alice =##> \case ("", c, Msg "hello") -> c == bobId; _ -> False
+  disconnectAgentClient alice
+  disconnectAgentClient bob2
 
 testDuplicateMessage :: HasCallStack => ATransport -> IO ()
 testDuplicateMessage t = do
@@ -740,6 +749,8 @@ testDuplicateMessage t = do
       6 <- sendMessage alice2 bobId SMP.noMsgFlags "hello 3"
       get alice2 ##> ("", bobId, SENT 6)
       get bob2 =##> \case ("", c, Msg "hello 3") -> c == aliceId; _ -> False
+  disconnectAgentClient alice2
+  disconnectAgentClient bob2
 
 testSkippedMessages :: HasCallStack => ATransport -> IO ()
 testSkippedMessages t = do
@@ -787,25 +798,20 @@ testSkippedMessages t = do
       get alice2 ##> ("", bobId, SENT 9)
       get bob2 =##> \case ("", c, Msg "hello 6") -> c == aliceId; _ -> False
       ackMessage bob2 aliceId 6 Nothing
+  disconnectAgentClient alice2
+  disconnectAgentClient bob2
 
 testRatchetSync :: HasCallStack => ATransport -> IO ()
-testRatchetSync t = do
-  alice <- getSMPAgentClient' agentCfg initAgentServers testDB
-  bob <- getSMPAgentClient' agentCfg initAgentServers testDB2
+testRatchetSync t = withAgentClients2 $ \alice bob ->
   withSmpServerStoreMsgLogOn t testPort $ \_ -> do
     (aliceId, bobId, bob2) <- setupDesynchronizedRatchet alice bob
     runRight $ do
       ConnectionStats {ratchetSyncState} <- synchronizeRatchet bob2 aliceId False
       liftIO $ ratchetSyncState `shouldBe` RSStarted
-
       get alice =##> ratchetSyncP bobId RSAgreed
-
       get bob2 =##> ratchetSyncP aliceId RSAgreed
-
       get alice =##> ratchetSyncP bobId RSOk
-
       get bob2 =##> ratchetSyncP aliceId RSOk
-
       exchangeGreetingsMsgIds alice bobId 12 bob2 aliceId 9
 
 setupDesynchronizedRatchet :: HasCallStack => AgentClient -> AgentClient -> IO (ConnId, ConnId, AgentClient)
@@ -868,9 +874,7 @@ ratchetSyncP' cId rss = \case
   _ -> False
 
 testRatchetSyncServerOffline :: HasCallStack => ATransport -> IO ()
-testRatchetSyncServerOffline t = do
-  alice <- getSMPAgentClient' agentCfg initAgentServers testDB
-  bob <- getSMPAgentClient' agentCfg initAgentServers testDB2
+testRatchetSyncServerOffline t = withAgentClients2 $ \alice bob -> do
   (aliceId, bobId, bob2) <- withSmpServerStoreMsgLogOn t testPort $ \_ ->
     setupDesynchronizedRatchet alice bob
 
@@ -886,16 +890,12 @@ testRatchetSyncServerOffline t = do
         [ ratchetSyncP' bobId RSAgreed,
           serverUpP
         ]
-
       liftIO . getInAnyOrder bob2 $
         [ ratchetSyncP' aliceId RSAgreed,
           serverUpP
         ]
-
       get alice =##> ratchetSyncP bobId RSOk
-
       get bob2 =##> ratchetSyncP aliceId RSOk
-
       exchangeGreetingsMsgIds alice bobId 12 bob2 aliceId 9
 
 serverUpP :: ATransmission 'Agent -> Bool
@@ -909,32 +909,24 @@ testRatchetSyncClientRestart t = do
   bob <- getSMPAgentClient' agentCfg initAgentServers testDB2
   (aliceId, bobId, bob2) <- withSmpServerStoreMsgLogOn t testPort $ \_ ->
     setupDesynchronizedRatchet alice bob
-
   ("", "", DOWN _ _) <- nGet alice
   ("", "", DOWN _ _) <- nGet bob2
-
   ConnectionStats {ratchetSyncState} <- runRight $ synchronizeRatchet bob2 aliceId False
   liftIO $ ratchetSyncState `shouldBe` RSStarted
-
   disconnectAgentClient bob2
-
   bob3 <- getSMPAgentClient' agentCfg initAgentServers testDB2
-
   withSmpServerStoreMsgLogOn t testPort $ \_ -> do
     runRight_ $ do
       ("", "", UP _ _) <- nGet alice
-
       subscribeConnection bob3 aliceId
-
       get alice =##> ratchetSyncP bobId RSAgreed
-
       get bob3 =##> ratchetSyncP aliceId RSAgreed
-
       get alice =##> ratchetSyncP bobId RSOk
-
       get bob3 =##> ratchetSyncP aliceId RSOk
-
       exchangeGreetingsMsgIds alice bobId 12 bob3 aliceId 9
+  disconnectAgentClient alice
+  disconnectAgentClient bob
+  disconnectAgentClient bob3
 
 testRatchetSyncSuspendForeground :: HasCallStack => ATransport -> IO ()
 testRatchetSyncSuspendForeground t = do
@@ -959,17 +951,16 @@ testRatchetSyncSuspendForeground t = do
         [ ratchetSyncP' bobId RSAgreed,
           serverUpP
         ]
-
       liftIO . getInAnyOrder bob2 $
         [ ratchetSyncP' aliceId RSAgreed,
           serverUpP
         ]
-
       get alice =##> ratchetSyncP bobId RSOk
-
       get bob2 =##> ratchetSyncP aliceId RSOk
-
       exchangeGreetingsMsgIds alice bobId 12 bob2 aliceId 9
+  disconnectAgentClient alice
+  disconnectAgentClient bob
+  disconnectAgentClient bob2
 
 testRatchetSyncSimultaneous :: HasCallStack => ATransport -> IO ()
 testRatchetSyncSimultaneous t = do
@@ -993,17 +984,16 @@ testRatchetSyncSimultaneous t = do
         [ ratchetSyncP' bobId RSAgreed,
           serverUpP
         ]
-
       liftIO . getInAnyOrder bob2 $
         [ ratchetSyncP' aliceId RSAgreed,
           serverUpP
         ]
-
       get alice =##> ratchetSyncP bobId RSOk
-
       get bob2 =##> ratchetSyncP aliceId RSOk
-
       exchangeGreetingsMsgIds alice bobId 12 bob2 aliceId 9
+  disconnectAgentClient alice
+  disconnectAgentClient bob
+  disconnectAgentClient bob2
 
 makeConnection :: AgentClient -> AgentClient -> ExceptT AgentErrorType IO (ConnId, ConnId)
 makeConnection alice bob = makeConnectionForUsers alice 1 bob 1
@@ -1027,6 +1017,7 @@ testInactiveClientDisconnected t = do
     runRight_ $ do
       (connId, _cReq) <- createConnection alice 1 True SCMInvitation Nothing
       nGet alice ##> ("", "", DOWN testSMPServer [connId])
+    disconnectAgentClient alice
 
 testActiveClientNotDisconnected :: ATransport -> IO ()
 testActiveClientNotDisconnected t = do
@@ -1037,6 +1028,7 @@ testActiveClientNotDisconnected t = do
     runRight_ $ do
       (connId, _cReq) <- createConnection alice 1 True SCMInvitation Nothing
       keepSubscribing alice connId ts
+    disconnectAgentClient alice
   where
     keepSubscribing :: AgentClient -> ConnId -> SystemTime -> ExceptT AgentErrorType IO ()
     keepSubscribing alice connId ts = do
@@ -1056,10 +1048,8 @@ testActiveClientNotDisconnected t = do
     milliseconds ts = systemSeconds ts * 1000 + fromIntegral (systemNanoseconds ts `div` 1000000)
 
 testSuspendingAgent :: IO ()
-testSuspendingAgent = do
-  a <- getSMPAgentClient' agentCfg initAgentServers testDB
-  b <- getSMPAgentClient' agentCfg initAgentServers testDB2
-  runRight_ $ do
+testSuspendingAgent =
+  withAgentClients2 $ \a b -> runRight_ $ do
     (aId, bId) <- makeConnection a b
     4 <- sendMessage a bId SMP.noMsgFlags "hello"
     get a ##> ("", bId, SENT 4)
@@ -1074,9 +1064,7 @@ testSuspendingAgent = do
     get b =##> \case ("", c, Msg "hello 2") -> c == aId; _ -> False
 
 testSuspendingAgentCompleteSending :: ATransport -> IO ()
-testSuspendingAgentCompleteSending t = do
-  a <- getSMPAgentClient' agentCfg initAgentServers testDB
-  b <- getSMPAgentClient' agentCfg initAgentServers testDB2
+testSuspendingAgentCompleteSending t = withAgentClients2 $ \a b -> do
   (aId, bId) <- withSmpServerStoreLogOn t testPort $ \_ -> runRight $ do
     (aId, bId) <- makeConnection a b
     4 <- sendMessage a bId SMP.noMsgFlags "hello"
@@ -1106,9 +1094,7 @@ testSuspendingAgentCompleteSending t = do
     ackMessage a bId 6 Nothing
 
 testSuspendingAgentTimeout :: ATransport -> IO ()
-testSuspendingAgentTimeout t = do
-  a <- getSMPAgentClient' agentCfg initAgentServers testDB
-  b <- getSMPAgentClient' agentCfg initAgentServers testDB2
+testSuspendingAgentTimeout t = withAgentClients2 $ \a b -> do
   (aId, _) <- withSmpServer t . runRight $ do
     (aId, bId) <- makeConnection a b
     4 <- sendMessage a bId SMP.noMsgFlags "hello"
@@ -1161,6 +1147,8 @@ testBatchedSubscriptions nCreate nDel t = do
     delete b aIds'
     deleteFail a bIds'
     deleteFail b aIds'
+  disconnectAgentClient a
+  disconnectAgentClient b
   where
     subscribe :: AgentClient -> [ConnId] -> ExceptT AgentErrorType IO ()
     subscribe c cs = do
@@ -1191,10 +1179,8 @@ testBatchedSubscriptions nCreate nDel t = do
         pure res
 
 testAsyncCommands :: IO ()
-testAsyncCommands = do
-  alice <- getSMPAgentClient' agentCfg initAgentServers testDB
-  bob <- getSMPAgentClient' agentCfg initAgentServers testDB2
-  runRight_ $ do
+testAsyncCommands =
+  withAgentClients2 $ \alice bob -> runRight_ $ do
     bobId <- createConnectionAsync alice 1 "1" True SCMInvitation
     ("1", bobId', INV (ACR _ qInfo)) <- get alice
     liftIO $ bobId' `shouldBe` bobId
@@ -1248,12 +1234,11 @@ testAsyncCommandsRestore t = do
       subscribeConnection alice' bobId
       ("1", _, INV _) <- get alice'
       pure ()
+  disconnectAgentClient alice'
 
 testAcceptContactAsync :: IO ()
-testAcceptContactAsync = do
-  alice <- getSMPAgentClient' agentCfg initAgentServers testDB
-  bob <- getSMPAgentClient' agentCfg initAgentServers testDB2
-  runRight_ $ do
+testAcceptContactAsync =
+  withAgentClients2 $ \alice bob -> runRight_ $ do
     (_, qInfo) <- createConnection alice 1 True SCMContact Nothing
     aliceId <- joinConnection bob 1 True qInfo "bob's connInfo"
     ("", _, REQ invId _ "bob's connInfo") <- get alice
@@ -1308,6 +1293,7 @@ testDeleteConnectionAsync t = do
     get a =##> \case ("", c, DEL_CONN) -> c `elem` connIds; _ -> False
     get a =##> \case ("", c, DEL_CONN) -> c `elem` connIds; _ -> False
     liftIO $ noMessages a "nothing else should be delivered to alice"
+  disconnectAgentClient a
 
 testJoinConnectionAsyncReplyError :: HasCallStack => ATransport -> IO ()
 testJoinConnectionAsyncReplyError t = do
@@ -1348,12 +1334,12 @@ testJoinConnectionAsyncReplyError t = do
       get b ##> ("", aId, INFO "alice's connInfo")
       get b ##> ("", aId, CON)
       exchangeGreetings a bId b aId
+  disconnectAgentClient a
+  disconnectAgentClient b
 
 testUsers :: IO ()
-testUsers = do
-  a <- getSMPAgentClient' agentCfg initAgentServers testDB
-  b <- getSMPAgentClient' agentCfg initAgentServers testDB2
-  runRight_ $ do
+testUsers =
+  withAgentClients2 $ \a b -> runRight_ $ do
     (aId, bId) <- makeConnection a b
     exchangeGreetingsMsgId 4 a bId b aId
     auId <- createUser a [noAuthSrv testSMPServer] [noAuthSrv testXFTPServer]
@@ -1367,10 +1353,8 @@ testUsers = do
     liftIO $ noMessages a "nothing else should be delivered to alice"
 
 testDeleteUserQuietly :: IO ()
-testDeleteUserQuietly = do
-  a <- getSMPAgentClient' agentCfg initAgentServers testDB
-  b <- getSMPAgentClient' agentCfg initAgentServers testDB2
-  runRight_ $ do
+testDeleteUserQuietly =
+  withAgentClients2 $ \a b -> runRight_ $ do
     (aId, bId) <- makeConnection a b
     exchangeGreetingsMsgId 4 a bId b aId
     auId <- createUser a [noAuthSrv testSMPServer] [noAuthSrv testXFTPServer]
@@ -1381,9 +1365,7 @@ testDeleteUserQuietly = do
     liftIO $ noMessages a "nothing else should be delivered to alice"
 
 testUsersNoServer :: HasCallStack => ATransport -> IO ()
-testUsersNoServer t = do
-  a <- getSMPAgentClient' agentCfg {initialCleanupDelay = 10000, cleanupInterval = 10000, deleteErrorCount = 3} initAgentServers testDB
-  b <- getSMPAgentClient' agentCfg initAgentServers testDB2
+testUsersNoServer t = withAgentClientsCfg2 aCfg agentCfg $ \a b -> do
   (aId, bId, auId, _aId', bId') <- withSmpServerStoreLogOn t testPort $ \_ -> runRight $ do
     (aId, bId) <- makeConnection a b
     exchangeGreetingsMsgId 4 a bId b aId
@@ -1404,6 +1386,8 @@ testUsersNoServer t = do
     nGet a =##> \case ("", "", UP _ [c]) -> c == bId; _ -> False
     nGet b =##> \case ("", "", UP _ cs) -> length cs == 2; _ -> False
     exchangeGreetingsMsgId 6 a bId b aId
+  where
+    aCfg = agentCfg {initialCleanupDelay = 10000, cleanupInterval = 10000, deleteErrorCount = 3}
 
 testSwitchConnection :: InitialAgentServers -> IO ()
 testSwitchConnection servers = do
@@ -1414,6 +1398,8 @@ testSwitchConnection servers = do
     exchangeGreetingsMsgId 4 a bId b aId
     testFullSwitch a bId b aId 10
     testFullSwitch a bId b aId 16
+  disconnectAgentClient a
+  disconnectAgentClient b
 
 testFullSwitch :: AgentClient -> ByteString -> AgentClient -> ByteString -> Int64 -> ExceptT AgentErrorType IO ()
 testFullSwitch a bId b aId msgId = do
@@ -1521,6 +1507,8 @@ testSwitchDelete servers = do
     get a =##> \case ("", c, DEL_RCVQ _ _ Nothing) -> c == bId; _ -> False
     get a =##> \case ("", c, DEL_CONN) -> c == bId; _ -> False
     liftIO $ noMessages a "nothing else should be delivered to alice"
+  disconnectAgentClient a
+  disconnectAgentClient b
 
 testAbortSwitchStarted :: HasCallStack => InitialAgentServers -> IO ()
 testAbortSwitchStarted servers = do
@@ -1792,7 +1780,7 @@ testCreateQueueAuth :: HasCallStack => (Maybe BasicAuth, Version) -> (Maybe Basi
 testCreateQueueAuth clnt1 clnt2 = do
   a <- getClient clnt1
   b <- getClient clnt2
-  runRight $ do
+  r <- runRight $ do
     tryError (createConnection a 1 True SCMInvitation Nothing) >>= \case
       Left (SMP AUTH) -> pure 0
       Left e -> throwError e
@@ -1808,6 +1796,9 @@ testCreateQueueAuth clnt1 clnt2 = do
             get b ##> ("", aId, CON)
             exchangeGreetings a bId b aId
             pure 2
+  disconnectAgentClient a
+  disconnectAgentClient b
+  pure r
   where
     getClient (clntAuth, clntVersion) =
       let servers = initAgentServers {smp = userServers [ProtoServerWithAuth testSMPServer clntAuth]}
@@ -1821,20 +1812,16 @@ testSMPServerConnectionTest t newQueueBasicAuth srv =
     runRight $ testProtocolServer a 1 srv
 
 testRatchetAdHash :: HasCallStack => IO ()
-testRatchetAdHash = do
-  a <- getSMPAgentClient' agentCfg initAgentServers testDB
-  b <- getSMPAgentClient' agentCfg initAgentServers testDB2
-  runRight_ $ do
+testRatchetAdHash =
+  withAgentClients2 $ \a b -> runRight_ $ do
     (aId, bId) <- makeConnection a b
     ad1 <- getConnectionRatchetAdHash a bId
     ad2 <- getConnectionRatchetAdHash b aId
     liftIO $ ad1 `shouldBe` ad2
 
 testDeliveryReceipts :: HasCallStack => IO ()
-testDeliveryReceipts = do
-  a <- getSMPAgentClient' agentCfg initAgentServers testDB
-  b <- getSMPAgentClient' agentCfg initAgentServers testDB2
-  runRight_ $ do
+testDeliveryReceipts =
+  withAgentClients2 $ \a b -> runRight_ $ do
     (aId, bId) <- makeConnection a b
     -- a sends, b receives and sends delivery receipt
     4 <- sendMessage a bId SMP.noMsgFlags "hello"
@@ -1896,12 +1883,12 @@ testDeliveryReceiptsVersion t = do
       ackMessage a' bId 10 $ Just ""
       get b' =##> \case ("", c, Rcvd 10) -> c == aId; _ -> False
       ackMessage b' aId 11 Nothing
+    disconnectAgentClient a'
+    disconnectAgentClient b'
 
 testTwoUsers :: HasCallStack => IO ()
-testTwoUsers = do
+testTwoUsers = withAgentClients2 $ \a b -> do
   let nc = netCfg initAgentServers
-  a <- getSMPAgentClient' agentCfg initAgentServers testDB
-  b <- getSMPAgentClient' agentCfg initAgentServers testDB2
   sessionMode nc `shouldBe` TSMUser
   runRight_ $ do
     (aId1, bId1) <- makeConnectionForUsers a 1 b 1
@@ -1971,10 +1958,8 @@ getSMPAgentClient' cfg' initServers dbPath = do
   getSMPAgentClient cfg' initServers st
 
 testServerMultipleIdentities :: HasCallStack => IO ()
-testServerMultipleIdentities = do
-  alice <- getSMPAgentClient' agentCfg initAgentServers testDB
-  bob <- getSMPAgentClient' agentCfg initAgentServers testDB2
-  runRight_ $ do
+testServerMultipleIdentities =
+  withAgentClients2 $ \alice bob -> runRight_ $ do
     (bobId, cReq) <- createConnection alice 1 True SCMInvitation Nothing
     aliceId <- joinConnection bob 1 True cReq "bob's connInfo"
     ("", _, CONF confId _ "bob's connInfo") <- get alice

tests/AgentTests/NotificationTests.hs

@@ -111,7 +111,7 @@ testNotificationToken APNSMockServer {apnsQ} = do
     deleteNtfToken a tkn
     -- agent deleted this token
     Left (CMD PROHIBITED) <- tryE $ checkNtfToken a tkn
-    pure ()
+    disconnectAgentClient a
 
 (.->) :: J.Value -> J.Key -> ExceptT AgentErrorType IO ByteString
 v .-> key = do
@@ -143,7 +143,7 @@ testNtfTokenRepeatRegistration APNSMockServer {apnsQ} = do
     -- can still use the first verification code, it is the same after decryption
     verifyNtfToken a tkn nonce verification
     NTActive <- checkNtfToken a tkn
-    pure ()
+    disconnectAgentClient a
 
 testNtfTokenSecondRegistration :: APNSMockServer -> IO ()
 testNtfTokenSecondRegistration APNSMockServer {apnsQ} = do
@@ -179,7 +179,8 @@ testNtfTokenSecondRegistration APNSMockServer {apnsQ} = do
     Left (NTF AUTH) <- tryE $ checkNtfToken a tkn
     -- and the second is active
     NTActive <- checkNtfToken a' tkn
-    pure ()
+    disconnectAgentClient a
+    disconnectAgentClient a'
 
 testNtfTokenServerRestart :: ATransport -> APNSMockServer -> IO ()
 testNtfTokenServerRestart t APNSMockServer {apnsQ} = do
@@ -208,7 +209,7 @@ testNtfTokenServerRestart t APNSMockServer {apnsQ} = do
     liftIO $ sendApnsResponse' APNSRespOk
     verifyNtfToken a' tkn nonce' verification'
     NTActive <- checkNtfToken a' tkn
-    pure ()
+    disconnectAgentClient a'
 
 testNotificationSubscriptionExistingConnection :: APNSMockServer -> IO ()
 testNotificationSubscriptionExistingConnection APNSMockServer {apnsQ} = do
@@ -262,6 +263,8 @@ testNotificationSubscriptionExistingConnection APNSMockServer {apnsQ} = do
     get bob ##> ("", aliceId, SENT $ baseId + 2)
     -- no notifications should follow
     noNotification apnsQ
+  disconnectAgentClient alice
+  disconnectAgentClient bob
   where
     baseId = 3
     msgId = subtract baseId
@@ -305,6 +308,8 @@ testNotificationSubscriptionNewConnection APNSMockServer {apnsQ} = do
     ackMessage bob aliceId (baseId + 2) Nothing
     -- no unexpected notifications should follow
     noNotification apnsQ
+  disconnectAgentClient alice
+  disconnectAgentClient bob
   where
     baseId = 3
     msgId = subtract baseId
@@ -382,6 +387,8 @@ testChangeNotificationsMode APNSMockServer {apnsQ} = do
     ackMessage alice bobId (baseId + 5) Nothing
     -- no notifications should follow
     noNotification apnsQ
+  disconnectAgentClient alice
+  disconnectAgentClient bob
   where
     baseId = 3
     msgId = subtract baseId
@@ -425,6 +432,8 @@ testChangeToken APNSMockServer {apnsQ} = do
     ackMessage alice1 bobId (baseId + 2) Nothing
     -- no notifications should follow
     noNotification apnsQ
+  disconnectAgentClient alice1
+  disconnectAgentClient bob
   where
     baseId = 3
     msgId = subtract baseId
@@ -454,6 +463,8 @@ testNotificationsStoreLog t APNSMockServer {apnsQ} = do
     void $ messageNotification apnsQ
     get alice =##> \case ("", c, Msg "hello again") -> c == bobId; _ -> False
     liftIO $ killThread threadId
+  disconnectAgentClient alice
+  disconnectAgentClient bob
 
 testNotificationsSMPRestart :: ATransport -> APNSMockServer -> IO ()
 testNotificationsSMPRestart t APNSMockServer {apnsQ} = do
@@ -484,6 +495,8 @@ testNotificationsSMPRestart t APNSMockServer {apnsQ} = do
     _ <- messageNotificationData alice apnsQ
     get alice =##> \case ("", c, Msg "hello again") -> c == bobId; _ -> False
     liftIO $ killThread threadId
+  disconnectAgentClient alice
+  disconnectAgentClient bob
 
 testNotificationsSMPRestartBatch :: Int -> ATransport -> APNSMockServer -> IO ()
 testNotificationsSMPRestartBatch n t APNSMockServer {apnsQ} = do
@@ -522,6 +535,8 @@ testNotificationsSMPRestartBatch n t APNSMockServer {apnsQ} = do
       get b ##> ("", aliceId, SENT msgId)
       _ <- messageNotificationData a apnsQ
       get a =##> \case ("", c, Msg "hello again") -> c == bobId; _ -> False
+  disconnectAgentClient a
+  disconnectAgentClient b
   where
     runServers :: ExceptT AgentErrorType IO a -> IO a
     runServers a = do
@@ -551,6 +566,8 @@ testSwitchNotifications servers APNSMockServer {apnsQ} = do
     switchComplete a bId b aId
     liftIO $ threadDelay 500000
     testMessage "hello again"
+  disconnectAgentClient a
+  disconnectAgentClient b
 
 messageNotification :: TBQueue APNSMockRequest -> ExceptT AgentErrorType IO (C.CbNonce, ByteString)
 messageNotification apnsQ = do

tests/SMPClient.hs

@@ -193,3 +193,6 @@ smpTest4 _ test' = smpTestN 4 _test
     _test :: HasCallStack => [THandle c] -> IO ()
     _test [h1, h2, h3, h4] = test' h1 h2 h3 h4
     _test _ = error "expected 4 handles"
+
+unexpected :: (HasCallStack, Show a) => a -> Expectation
+unexpected r = expectationFailure $ "unexpected response " <> show r

tests/ServerTests.hs

@@ -551,7 +551,9 @@ testWithStoreLog at@(ATransport t) =
         writeTVar senderId1 sId1
         writeTVar notifierId nId
       Resp "dabc" _ OK <- signSendRecv h1 nKey ("dabc", nId, NSUB)
-      Resp "bcda" _ OK <- signSendRecv h sKey1 ("bcda", sId1, _SEND' "hello")
+      signSendRecv h sKey1 ("bcda", sId1, _SEND' "hello") >>= \case
+        Resp "bcda" _ OK -> pure ()
+        r -> unexpected r
       Resp "" _ (Msg mId1 msg1) <- tGet1 h
       (decryptMsgV3 dhShared mId1 msg1, Right "hello") #== "delivered from queue 1"
       Resp "" _ (NMSG _ _) <- tGet1 h1
@@ -560,7 +562,7 @@ testWithStoreLog at@(ATransport t) =
       atomically $ writeTVar senderId2 sId2
       signSendRecv h sKey2 ("cdab", sId2, _SEND "hello too") >>= \case
         Resp "cdab" _ OK -> pure ()
-        r -> print $ "unexpected response " <> show r
+        r -> unexpected r
       Resp "" _ (Msg mId2 msg2) <- tGet1 h
       (decryptMsgV3 dhShared2 mId2 msg2, Right "hello too") #== "delivered from queue 2"
 
@@ -973,7 +975,7 @@ testMsgExpireOnInterval t =
         testSMPClient @c $ \rh -> do
           signSendRecv rh rKey ("2", rId, SUB) >>= \case
             Resp "2" _ OK -> pure ()
-            r -> error $ "unexpected response: " <> show r
+            r -> unexpected r
           1000 `timeout` tGet @ErrorType @BrokerMsg rh >>= \case
             Nothing -> return ()
             Just _ -> error "nothing should be delivered"