Minor refactoring

11 files changed, 242 insertions, 101 deletions

diff --git a/src/Network/BitTorrent/Client/Handle.hs b/src/Network/BitTorrent/Client/Handle.hs
index e7ac779c..e7884caa 100644
--- a/src/Network/BitTorrent/Client/Handle.hs
+++ b/src/Network/BitTorrent/Client/Handle.hs
@@ -110,12 +110,12 @@ start Handle {..} = do
   Client {..} <- getClient
   liftIO $ Tracker.notify trackerManager trackers Tracker.Started
   unless private $ do
-    liftDHT $ DHT.insert topic undefined
+    liftDHT $ DHT.insert topic (error "start")
   liftIO $ do
     peers <- askPeers trackerManager trackers
     print $ "got: " ++ show (L.length peers) ++ " peers"
     forM_ peers $ \ peer -> do
-      Exchange.insert peer exchange
+      Exchange.connect peer exchange
 
 -- | Stop downloading this torrent.
 pause :: Handle -> BitTorrent ()
@@ -126,7 +126,7 @@ stop :: Handle -> BitTorrent ()
 stop Handle {..} = do
   Client {..} <- getClient
   unless private $ do
-    liftDHT $ DHT.delete topic undefined
+    liftDHT $ DHT.delete topic (error "stop")
   liftIO  $ Tracker.notify trackerManager trackers Tracker.Stopped
 
 {-----------------------------------------------------------------------
diff --git a/src/Network/BitTorrent/DHT/ContactInfo.hs b/src/Network/BitTorrent/DHT/ContactInfo.hs
index 06d2dac0..028a4214 100644
--- a/src/Network/BitTorrent/DHT/ContactInfo.hs
+++ b/src/Network/BitTorrent/DHT/ContactInfo.hs
@@ -1,4 +1,6 @@
-module Network.BitTorrent.DHT.ContactInfo () where
+module Network.BitTorrent.DHT.ContactInfo
+       (        ) where
+{-
 import Data.HashMap.Strict as HM
 
 import Data.Torrent.InfoHash
@@ -8,6 +10,57 @@ import Network.BitTorrent.Core
 -- decrease prefix when table is too small
 -- filter outdated peers
 
+{-----------------------------------------------------------------------
+--  PeerSet
+-----------------------------------------------------------------------}
+
+type PeerSet a = [(PeerAddr a, NodeInfo a, Timestamp)]
+
+-- compare PSQueue vs Ordered list
+
+takeNewest :: PeerSet a -> [PeerAddr a]
+takeNewest = undefined
+
+dropOld :: Timestamp -> PeerSet a -> PeerSet a
+dropOld = undefined
+
+insert :: PeerAddr a -> Timestamp -> PeerSet a -> PeerSet a
+insert = undefined
+
+type Mask = Int
+type Size = Int
+type Timestamp = Int
+
+{-----------------------------------------------------------------------
+--  InfoHashMap
+-----------------------------------------------------------------------}
+
+-- compare handwritten prefix tree versus IntMap
+
+data Tree a
+  = Nil
+  | Tip !InfoHash !(PeerSet a)
+  | Bin !InfoHash !Mask !Size !Timestamp (Tree a) (Tree a)
+
+insertTree :: InfoHash -> a -> Tree a -> Tree a
+insertTree = undefined
+
+type Prio = Int
+
+--shrink :: ContactInfo ip -> Int
+shrink  Nil      = Nil
+shrink (Tip _ _) = undefined
+shrink (Bin _ _) = undefined
+
+{-----------------------------------------------------------------------
+-- InfoHashMap
+-----------------------------------------------------------------------}
+
+-- compare new design versus HashMap
+
+data IntMap k p a
+type ContactInfo = Map InfoHash Timestamp (Set (PeerAddr IP) Timestamp)
+
 data ContactInfo ip = PeerStore
   { maxSize    :: Int
   , prefixSize :: Int
@@ -16,3 +69,25 @@ data ContactInfo ip = PeerStore
   , count      :: Int  -- ^ Cached size of the 'peerSet'
   , peerSet    :: HashMap InfoHash [PeerAddr ip]
   }
+
+size :: ContactInfo ip -> Int
+size = undefined
+
+prefixSize :: ContactInfo ip -> Int
+prefixSize = undefined
+
+lookup :: InfoHash -> ContactInfo ip -> [PeerAddr ip]
+lookup = undefined
+
+insert :: InfoHash -> PeerAddr ip -> ContactInfo ip -> ContactInfo ip
+insert = undefined
+
+-- | Limit in size.
+prune :: NodeId -> Int -> ContactInfo ip -> ContactInfo ip
+prune pref targetSize  Nil      = Nil
+prune pref targetSize (Tip _ _) = undefined
+
+-- | Remove expired entries.
+splitGT :: Timestamp -> ContactInfo ip -> ContactInfo ip
+splitGT = undefined
+-}
\ No newline at end of file
diff --git a/src/Network/BitTorrent/DHT/Session.hs b/src/Network/BitTorrent/DHT/Session.hs
index 755985fc..e770b1d3 100644
--- a/src/Network/BitTorrent/DHT/Session.hs
+++ b/src/Network/BitTorrent/DHT/Session.hs
@@ -434,15 +434,20 @@ insertNode info = fork $ do
 
 -- TODO limit dht peer store in size (probably by removing oldest peers)
 
+refreshContacts :: DHT ip ()
+refreshContacts = undefined
+
 -- | Insert peer to peer store. Used to handle announce requests.
 insertPeer :: Eq ip => InfoHash -> PeerAddr ip -> DHT ip ()
 insertPeer ih addr = do
+  refreshContacts
   var <- asks contactInfo
   liftIO $ atomically $ modifyTVar' var (P.insert ih addr)
 
 -- | Get peer set for specific swarm.
 lookupPeers :: InfoHash -> DHT ip [PeerAddr ip]
 lookupPeers ih = do
+  refreshContacts
   var <- asks contactInfo
   liftIO $ P.lookup ih <$> readTVarIO var
 
diff --git a/src/Network/BitTorrent/Exchange.hs b/src/Network/BitTorrent/Exchange.hs
index 86e13d58..8dac3c84 100644
--- a/src/Network/BitTorrent/Exchange.hs
+++ b/src/Network/BitTorrent/Exchange.hs
@@ -6,25 +6,23 @@
 --   Portability :  portable
 --
 module Network.BitTorrent.Exchange
-       ( -- * Options
+       ( -- * Manager
          Options (..)
-       , Caps
-       , Extension
-       , toCaps
-
-         -- * Manager
        , Manager
        , Handler
        , newManager
        , closeManager
 
          -- * Session
+       , Caps
+       , Extension
+       , toCaps
        , Session
        , newSession
        , closeSession
 
-         -- * Session control
-       , insert
+         -- * Connections
+       , connect
        ) where
 
 import Network.BitTorrent.Exchange.Manager
diff --git a/src/Network/BitTorrent/Exchange/Connection.hs b/src/Network/BitTorrent/Exchange/Connection.hs
index b23eb08b..dde9a468 100644
--- a/src/Network/BitTorrent/Exchange/Connection.hs
+++ b/src/Network/BitTorrent/Exchange/Connection.hs
@@ -723,10 +723,10 @@ extendedHandshake caps = do
     _ -> protocolError HandshakeRefused
 
 rehandshake :: ExtendedCaps -> Wire s ()
-rehandshake caps = undefined
+rehandshake caps = error "rehandshake"
 
 reconnect :: Wire s ()
-reconnect = undefined
+reconnect = error "reconnect"
 
 data ConnectionId    = ConnectionId
   { topic      :: !InfoHash
@@ -751,7 +751,7 @@ instance Default ConnectionPrefs where
     }
 
 normalize :: ConnectionPrefs -> ConnectionPrefs
-normalize = undefined
+normalize = error "normalize"
 
 -- | Bridge between 'Connection' and 'Network.BitTorrent.Exchange.Session'.
 data SessionLink s = SessionLink
diff --git a/src/Network/BitTorrent/Exchange/Manager.hs b/src/Network/BitTorrent/Exchange/Manager.hs
index f7f3cea7..b9aaa818 100644
--- a/src/Network/BitTorrent/Exchange/Manager.hs
+++ b/src/Network/BitTorrent/Exchange/Manager.hs
@@ -39,7 +39,7 @@ handleNewConn :: Socket -> PeerAddr IP -> Handler -> IO ()
 handleNewConn sock addr handler = do
   conn <- newPendingConnection sock addr
   ses  <- handler (pendingTopic conn) `onException` closePending conn
-  attach conn ses
+  establish conn ses
 
 listenIncoming :: Options -> Handler -> IO ()
 listenIncoming Options {..} handler = do
diff --git a/src/Network/BitTorrent/Exchange/Session.hs b/src/Network/BitTorrent/Exchange/Session.hs
index 74d0cc87..8c3d5388 100644
--- a/src/Network/BitTorrent/Exchange/Session.hs
+++ b/src/Network/BitTorrent/Exchange/Session.hs
@@ -2,16 +2,19 @@
 {-# LANGUAGE TemplateHaskell      #-}
 {-# LANGUAGE DeriveDataTypeable   #-}
 module Network.BitTorrent.Exchange.Session
-       ( Session
+       ( -- * Session
+         Session
        , LogFun
+       , sessionLogger
+
+         -- * Construction
        , newSession
        , closeSession
+       , withSession
 
-         -- * Connections
-       , Network.BitTorrent.Exchange.Session.insert
-       , Network.BitTorrent.Exchange.Session.attach
-       , Network.BitTorrent.Exchange.Session.delete
-       , Network.BitTorrent.Exchange.Session.deleteAll
+         -- * Connection Set
+       , connect
+       , establish
 
          -- * Events
        , waitMetadata
@@ -70,6 +73,7 @@ packException f m = try m >>= either (throwIO . f) return
 {-----------------------------------------------------------------------
 --  Session
 -----------------------------------------------------------------------}
+-- TODO unmap storage on zero connections
 
 data Cached a = Cached
   { cachedValue :: !a
@@ -79,9 +83,13 @@ data Cached a = Cached
 cache :: BEncode a => a -> Cached a
 cache s = Cached s (BE.encode s)
 
+-- | Logger function.
+type LogFun = Loc -> LogSource -> LogLevel -> LogStr -> IO ()
+
 data Session = Session
   { sessionPeerId          :: !(PeerId)
   , sessionTopic           :: !(InfoHash)
+  , sessionLogger          :: !(LogFun)
 
   , metadata               :: !(MVar Metadata.Status)
   , infodict               :: !(MVar (Cached InfoDict))
@@ -90,16 +98,27 @@ data Session = Session
   , storage                :: !(Storage)
 
   , connectionsPrefs       :: !ConnectionPrefs
+
+    -- | Connections either waiting for TCP/uTP 'connect' or waiting
+    -- for BT handshake.
   , connectionsPending     :: !(TVar (Set (PeerAddr IP)))
+
+    -- | Connections successfully handshaked and data transfer can
+    -- take place.
   , connectionsEstablished :: !(TVar (Map (PeerAddr IP) (Connection Session)))
+
+    -- | TODO implement choking mechanism
   , connectionsUnchoked    ::  [PeerAddr IP]
-  , broadcast              :: !(Chan Message)
 
-  , logger                 :: !(LogFun)
+    -- | Messages written to this channel will be sent to the all
+    -- connections, including pending connections (but right after
+    -- handshake).
+  , connectionsBroadcast   :: !(Chan Message)
   }
 
--- | Logger function.
-type LogFun = Loc -> LogSource -> LogLevel -> LogStr -> IO ()
+{-----------------------------------------------------------------------
+--  Session construction
+-----------------------------------------------------------------------}
 
 newSession :: LogFun
            -> PeerAddr (Maybe IP) -- ^ /external/ address of this peer;
@@ -108,17 +127,20 @@ newSession :: LogFun
            -> IO Session          -- ^
 newSession logFun addr rootPath dict = do
   pid         <- maybe genPeerId return (peerId addr)
-  pconnVar    <- newTVarIO S.empty
-  econnVar    <- newTVarIO M.empty
   store       <- openInfoDict ReadWriteEx rootPath dict
   statusVar   <- newMVar $ sessionStatus (BF.haveNone (totalPieces store))
                                          (piPieceLength (idPieceInfo dict))
-  metadataVar <- newMVar undefined
+  metadataVar <- newMVar (error "sessionMetadata")
   infodictVar <- newMVar (cache dict)
+
+  pSetVar     <- newTVarIO S.empty
+  eSetVar     <- newTVarIO M.empty
   chan        <- newChan
+
   return Session
     { sessionPeerId          = pid
     , sessionTopic           = idInfoHash dict
+    , sessionLogger          = logFun
 
     , metadata               = metadataVar
     , infodict               = infodictVar
@@ -127,21 +149,25 @@ newSession logFun addr rootPath dict = do
     , storage                = store
 
     , connectionsPrefs       = def
-    , connectionsPending     = pconnVar
-    , connectionsEstablished = econnVar
+    , connectionsPending     = pSetVar
+    , connectionsEstablished = eSetVar
     , connectionsUnchoked    = []
-    , broadcast              = chan
-
-    , logger                 = logFun
+    , connectionsBroadcast   = chan
     }
 
 closeSession :: Session -> IO ()
-closeSession ses = do
-  deleteAll ses
-  error "closeSession"
-
-waitMetadata :: Session -> IO InfoDict
-waitMetadata Session {..} = cachedValue <$> readMVar infodict
+closeSession Session {..} = do
+  close storage
+{-
+  hSet <- atomically $ do
+    pSet <- swapTVar connectionsPending     S.empty
+    eSet <- swapTVar connectionsEstablished S.empty
+    return pSet
+  mapM_ kill hSet
+-}
+
+withSession :: ()
+withSession = error "withSession"
 
 {-----------------------------------------------------------------------
 --  Logging
@@ -153,7 +179,7 @@ instance MonadLogger (Connected Session) where
     ses  <- asks connSession
     addr <- asks connRemoteAddr
     let addrSrc = src <> " @ " <> T.pack (render (pretty addr))
-    liftIO $ logger ses loc addrSrc lvl (toLogStr msg)
+    liftIO $ sessionLogger ses loc addrSrc lvl (toLogStr msg)
 
 logMessage :: MonadLogger m => Message -> m ()
 logMessage msg = logDebugN $ T.pack (render (pretty msg))
@@ -162,14 +188,22 @@ logEvent   :: MonadLogger m => Text    -> m ()
 logEvent       = logInfoN
 
 {-----------------------------------------------------------------------
---  Connection slots
+--  Connection set
 -----------------------------------------------------------------------}
---- pending -> established -> closed
----    |                        /|\
----    \-------------------------|
-
-pendingConnection :: PeerAddr IP -> Session -> IO Bool
-pendingConnection addr Session {..} = atomically $ do
+--- Connection status transition:
+---
+--- pending -> established -> finished -> closed
+---    |           \|/                      /|\
+---    \-------------------------------------|
+---
+--- Purpose of slots:
+---   1) to avoid duplicates
+---   2) connect concurrently
+---
+
+-- | Add connection to the pending set.
+pendingConnection :: PeerAddr IP -> Session -> STM Bool
+pendingConnection addr Session {..} = do
   pSet <- readTVar connectionsPending
   eSet <- readTVar connectionsEstablished
   if (addr `S.member` pSet) || (addr `M.member` eSet)
@@ -178,38 +212,37 @@ pendingConnection addr Session {..} = atomically $ do
       modifyTVar' connectionsPending (S.insert addr)
       return True
 
+-- | Pending connection successfully established, add it to the
+-- established set.
 establishedConnection :: Connected Session ()
-establishedConnection = undefined --atomically $ do
---  pSet <- readTVar pendingConnections
---  eSet <- readTVar
-  undefined
+establishedConnection = do
+  conn         <- ask
+  addr         <- asks connRemoteAddr
+  Session {..} <- asks connSession
+  liftIO $ atomically $ do
+    modifyTVar connectionsPending     (S.delete addr)
+    modifyTVar connectionsEstablished (M.insert addr conn)
 
+-- | Either this or remote peer decided to finish conversation
+-- (conversation is alread /established/ connection), remote it from
+-- the established set.
 finishedConnection :: Connected Session ()
-finishedConnection = return ()
+finishedConnection = do
+  Session {..} <- asks connSession
+  addr         <- asks connRemoteAddr
+  liftIO $ atomically $ do
+    modifyTVar connectionsEstablished $ M.delete addr
 
--- | There are no state for this connection, remove it.
-closedConnection :: PeerAddr IP -> Session -> IO ()
-closedConnection addr Session {..} = atomically $ do
+-- | There are no state for this connection, remove it from the all
+-- sets.
+closedConnection :: PeerAddr IP -> Session -> STM ()
+closedConnection addr Session {..} = do
   modifyTVar connectionsPending     $ S.delete addr
   modifyTVar connectionsEstablished $ M.delete addr
 
-{-----------------------------------------------------------------------
---  Connections
------------------------------------------------------------------------}
--- TODO unmap storage on zero connections
-
-mainWire :: Wire Session ()
-mainWire = do
-  lift establishedConnection
-  Session {..} <- asks connSession
-  lift $ resizeBitfield (totalPieces storage)
-  logEvent "Connection established"
-  iterM logMessage =$= exchange =$= iterM logMessage
-  lift finishedConnection
-
 getConnectionConfig :: Session -> IO (ConnectionConfig Session)
 getConnectionConfig s @ Session {..} = do
-  chan <- dupChan broadcast
+  chan <- dupChan connectionsBroadcast
   let sessionLink = SessionLink {
       linkTopic        = sessionTopic
     , linkPeerId       = sessionPeerId
@@ -223,28 +256,46 @@ getConnectionConfig s @ Session {..} = do
     , cfgWire    = mainWire
     }
 
-insert :: PeerAddr IP -> Session -> IO ()
-insert addr ses @ Session {..} = void $ forkIO (action `finally` cleanup)
+type Finalizer      = IO ()
+type Runner = (ConnectionConfig Session -> IO ())
+
+runConnection :: Runner -> Finalizer -> PeerAddr IP -> Session -> IO ()
+runConnection runner finalize addr set @ Session {..} = do
+    _ <- forkIO (action `finally` cleanup)
+    return ()
   where
     action = do
-      pendingConnection addr ses
-      cfg <- getConnectionConfig ses
-      connectWire addr cfg
+      notExist <- atomically $ pendingConnection addr set
+      when notExist $ do
+        cfg <- getConnectionConfig set
+        runner cfg
 
     cleanup = do
-      runStatusUpdates status (SS.resetPending addr)
+      finalize
+--      runStatusUpdates status (SS.resetPending addr)
       -- TODO Metata.resetPending addr
-      closedConnection addr ses
+      atomically $ closedConnection addr set
+
+-- | Establish connection from scratch. If this endpoint is already
+-- connected, no new connections is created. This function do not block.
+connect :: PeerAddr IP -> Session -> IO ()
+connect addr = runConnection (connectWire addr) (return ()) addr
 
--- TODO closePending on error
-attach :: PendingConnection -> Session -> IO ()
-attach = undefined
+-- | Establish connection with already pre-connected endpoint. If this
+-- endpoint is already connected, no new connections is created. This
+-- function do not block.
+--
+--  'PendingConnection' will be closed automatically, you do not need
+--  to call 'closePending'.
+establish :: PendingConnection -> Session -> IO ()
+establish conn = runConnection (acceptWire conn) (closePending conn)
+                 (pendingPeer conn)
 
-delete :: PeerAddr IP -> Session -> IO ()
-delete = undefined
+-- | Why do we need this message?
+type BroadcastMessage = ExtendedCaps -> Message
 
-deleteAll :: Session -> IO ()
-deleteAll = undefined
+broadcast :: BroadcastMessage -> Session -> IO ()
+broadcast = error "broadcast"
 
 {-----------------------------------------------------------------------
 --  Helpers
@@ -292,14 +343,17 @@ tryReadMetadataBlock pix = do
   Session {..} <- asks connSession
   mcached      <- liftIO (tryReadMVar infodict)
   case mcached of
-    Nothing            -> undefined
-    Just (Cached {..}) -> undefined
+    Nothing            -> error "tryReadMetadataBlock"
+    Just (Cached {..}) -> error "tryReadMetadataBlock"
 
 sendBroadcast :: PeerMessage msg => msg -> Wire Session ()
 sendBroadcast msg = do
   Session {..} <- asks connSession
-  ecaps <- use connExtCaps
-  liftIO $ writeChan broadcast (envelop ecaps msg)
+  error "sendBroadcast"
+--  liftIO $ msg `broadcast` sessionConnections
+
+waitMetadata :: Session -> IO InfoDict
+waitMetadata Session {..} = cachedValue <$> readMVar infodict
 
 {-----------------------------------------------------------------------
 --  Triggers
@@ -406,7 +460,7 @@ tryRequestMetadataBlock :: Trigger
 tryRequestMetadataBlock = do
   mpix <- lift $ withMetadataUpdates Metadata.scheduleBlock
   case mpix of
-    Nothing  -> undefined
+    Nothing  -> error "tryRequestMetadataBlock"
     Just pix -> sendMessage (MetadataRequest pix)
 
 metadataCompleted :: InfoDict -> Trigger
@@ -439,7 +493,7 @@ handleMetadata (MetadataUnknown _  ) = do
 -----------------------------------------------------------------------}
 
 acceptRehandshake :: ExtendedHandshake -> Trigger
-acceptRehandshake ehs = undefined
+acceptRehandshake ehs = error "acceptRehandshake"
 
 handleExtended :: Handler ExtendedMessage
 handleExtended (EHandshake     ehs) = acceptRehandshake ehs
@@ -451,8 +505,8 @@ handleMessage KeepAlive       = return ()
 handleMessage (Status s)      = handleStatus s
 handleMessage (Available msg) = handleAvailable msg
 handleMessage (Transfer  msg) = handleTransfer msg
-handleMessage (Port      n)   = undefined
-handleMessage (Fast      _)   = undefined
+handleMessage (Port      n)   = error "handleMessage"
+handleMessage (Fast      _)   = error "handleMessage"
 handleMessage (Extended  msg) = handleExtended msg
 
 exchange :: Wire Session ()
@@ -462,13 +516,22 @@ exchange = do
   sendMessage (Bitfield bf)
   awaitForever handleMessage
 
+mainWire :: Wire Session ()
+mainWire = do
+  lift establishedConnection
+  Session {..} <- asks connSession
+  lift $ resizeBitfield (totalPieces storage)
+  logEvent "Connection established"
+  iterM logMessage =$= exchange =$= iterM logMessage
+  lift finishedConnection
+
 data Event = NewMessage (PeerAddr IP) Message
            | Timeout -- for scheduling
 
 type Exchange a = Wire Session a
 
 awaitEvent :: Exchange Event
-awaitEvent = undefined
+awaitEvent = error "awaitEvent"
 
 yieldEvent :: Exchange Event
-yieldEvent = undefined
+yieldEvent = error "yieldEvent"
diff --git a/tests/Config.hs b/tests/Config.hs
index b37b727b..09e838cc 100644
--- a/tests/Config.hs
+++ b/tests/Config.hs
@@ -40,7 +40,7 @@ type ClientName = String
 
 
 instance Read PortNumber where
-  readsPrec = undefined
+  readsPrec = error "readsPrec"
 
 data ClientOpts = ClientOpts
   { peerPort   :: PortNumber  -- tcp port
diff --git a/tests/Network/BitTorrent/DHT/QuerySpec.hs b/tests/Network/BitTorrent/DHT/QuerySpec.hs
index 33114b40..d25bd120 100644
--- a/tests/Network/BitTorrent/DHT/QuerySpec.hs
+++ b/tests/Network/BitTorrent/DHT/QuerySpec.hs
@@ -18,7 +18,7 @@ import Network.BitTorrent.DHT.TestData
 
 
 myAddr :: NodeAddr IPv4
-myAddr = "0.0.0.0:8000"
+myAddr = "0.0.0.0:0"
 
 nullLogger :: LogFun
 nullLogger _ _ _ _ = return ()
@@ -37,7 +37,7 @@ getBootInfo = do
   simpleDHT [] $ pingQ startAddr
 
 spec :: Spec
-spec = do
+spec = parallel $ do
   describe "environment" $ do
     describe "test node" $ do
       it "is alive" $ do
diff --git a/tests/Network/BitTorrent/DHT/RoutingSpec.hs b/tests/Network/BitTorrent/DHT/RoutingSpec.hs
index 788c4341..c4a33357 100644
--- a/tests/Network/BitTorrent/DHT/RoutingSpec.hs
+++ b/tests/Network/BitTorrent/DHT/RoutingSpec.hs
@@ -27,7 +27,7 @@ runSimulation :: Eq ip => Env ip -> Routing ip a -> Maybe a
 runSimulation e m = evalState (runRouting ping closest timestamp m) e
   where
     ping    addr = gets (L.elem addr . network)
-    closest nid  = undefined
+    closest nid  = error "runSimulation"
     timestamp    = gets currentTime
 
 instance Arbitrary ip => Arbitrary (Env ip) where
diff --git a/tests/Network/BitTorrent/Exchange/SessionSpec.hs b/tests/Network/BitTorrent/Exchange/SessionSpec.hs
index 745649fe..f80fc864 100644
--- a/tests/Network/BitTorrent/Exchange/SessionSpec.hs
+++ b/tests/Network/BitTorrent/Exchange/SessionSpec.hs
@@ -4,7 +4,7 @@ import Test.Hspec
 
 import Data.Torrent
 import Network.BitTorrent.Core
-import Network.BitTorrent.Exchange.Session as S
+import Network.BitTorrent.Exchange.Session
 
 import Config
 
@@ -20,7 +20,7 @@ spec = do
         Torrent {..} <- getTestTorrent
         myAddr <- getMyAddr
         ses  <- newSession nullLogger myAddr "" tInfoDict
-        S.insert addr ses
+        connect addr ses
         dict <- waitMetadata ses
         closeSession ses
         dict `shouldBe` tInfoDict