Pair of linked transports for testing.

1 files changed, 24 insertions, 0 deletions

diff --git a/src/Network/QueryResponse.hs b/src/Network/QueryResponse.hs
index e4c73cd8..6df9ac5a 100644
--- a/src/Network/QueryResponse.hs
+++ b/src/Network/QueryResponse.hs
@@ -570,6 +570,30 @@ udpTransport' bind_address = do
     }
   return (tr, sock)
 
+chanTransport :: (addr -> TChan (x, addr)) -> addr -> TChan (x, addr) -> TVar Bool -> Transport err addr x
+chanTransport chanFromAddr self achan aclosed = Transport
+    { awaitMessage = \kont -> do
+        x <- atomically $ (Just <$> readTChan achan)
+                            `orElse`
+                          (readTVar aclosed >>= check >> return Nothing)
+        kont $ Right <$> x
+    , sendMessage = \them bs -> do
+        atomically $ writeTChan (chanFromAddr them) (bs,self)
+    , closeTransport = atomically $ writeTVar aclosed True
+    }
+
+-- | Returns a pair of transports linked together to simulate two computers talking to each other.
+testPairTransport :: IO (Transport err SockAddr ByteString, Transport err SockAddr ByteString)
+testPairTransport = do
+    achan <- atomically newTChan
+    bchan <- atomically newTChan
+    aclosed <- atomically $ newTVar False
+    bclosed <- atomically $ newTVar False
+    let a = SockAddrInet 1 1
+        b = SockAddrInet 2 2
+    return ( chanTransport (const bchan) a achan aclosed
+           , chanTransport (const achan) b bchan bclosed )
+
 serializeClient :: Client err meth tid addr x -> IO (Client err meth tid addr x)
 serializeClient c = do
     mvar <- newMVar ()