1 files changed, 6 insertions, 5 deletions

diff --git a/tests/Network/BitTorrent/DHT/SessionSpec.hs b/tests/Network/BitTorrent/DHT/SessionSpec.hs
index ef844b31..32e4c158 100644
--- a/tests/Network/BitTorrent/DHT/SessionSpec.hs
+++ b/tests/Network/BitTorrent/DHT/SessionSpec.hs
@@ -15,6 +15,7 @@ import Network.BitTorrent.Address
 import Network.BitTorrent.DHT
 import Network.BitTorrent.DHT.Message
 import Network.BitTorrent.DHT.Session
+import Network.BitTorrent.DHT.Query
 
 import Data.TorrentSpec ()
 import Network.BitTorrent.CoreSpec ()
@@ -68,7 +69,7 @@ spec = do
       property $ \ (nid :: NodeId) -> do
         let info = NodeInfo nid myAddr
         closest <- simpleDHT $ do
-           _ <- insertNode info
+           _ <- insertNode info Nothing
            liftIO $ yield
            getClosest nid
         closest `shouldSatisfy` L.elem info
@@ -93,15 +94,15 @@ spec = do
       it "should always ping this node" $ do
         (rid, tid) <- simpleDHT $ do
           (remoteId, Ping) <- queryNode myAddr Ping
-          thisId    <-  asks thisNodeId
+          thisId <- myNodeIdAccordingTo (read "8.8.8.8:6881")
           return (remoteId, thisId)
         rid `shouldBe` tid
 
     describe "queryParallel" $ do
       it "should handle parallel requests" $ do
-        (nid, resps) <- simpleDHT $ (,)
-          <$> asks thisNodeId
-          <*> queryParallel (L.replicate 100 $ queryNode myAddr Ping)
+        (nid, resps) <- simpleDHT $ do
+          me <- myNodeIdAccordingTo (read "8.8.8.8:6881")
+          ( (,) me ) <$> queryParallel (L.replicate 100 $ queryNode myAddr Ping)
         resps `shouldSatisfy` L.all (== (nid, Ping))
 
     describe "(<@>) operator" $ do