6 files changed, 40 insertions, 38 deletions

diff --git a/bittorrent.cabal b/bittorrent.cabal
index b82599b8..0fddf5d0 100644
--- a/bittorrent.cabal
+++ b/bittorrent.cabal
@@ -298,7 +298,6 @@ test-suite spec
 
                        -- * Bittorrent
                      , bencoding
-                     , krpc  >= 0.6.1
                      , bittorrent
                      , temporary
   if flag(network-uri)
@@ -335,7 +334,6 @@ benchmark bench
                      , mtl
                      , monad-logger
                      , criterion
-                     , krpc
   ghc-options:         -O2 -fforce-recomp
 
 -- Utility to work with torrent files.
diff --git a/tests/Network/BitTorrent/DHT/MessageSpec.hs b/tests/Network/BitTorrent/DHT/MessageSpec.hs
index ae9ab487..3b1f5acd 100644
--- a/tests/Network/BitTorrent/DHT/MessageSpec.hs
+++ b/tests/Network/BitTorrent/DHT/MessageSpec.hs
@@ -61,16 +61,16 @@ spec = do
          running <- retry 5 $ timeout (100 * 1000) $ do
                       nid <- genNodeId
                       Response _remoteAddr Ping <-
-                          rpc (query remoteAddr (Query nid Ping))
+                          rpc (query remoteAddr (Query nid False Ping))
                       return ()
          running `shouldSatisfy` isJust
 
   describe "ping" $ do
     it "properly bencoded" $ do
       BE.decode "d2:id20:abcdefghij0123456789e"
-        `shouldBe` Right (Query "abcdefghij0123456789" Ping)
+        `shouldBe` Right (Query "abcdefghij0123456789" False Ping)
 
-      BE.encode (Query "abcdefghij0123456789" Ping)
+      BE.encode (Query "abcdefghij0123456789" False Ping)
         `shouldBe` "d2:id20:abcdefghij0123456789e"
 
       BE.decode "d2:id20:mnopqrstuvwxyz123456e"
@@ -80,22 +80,22 @@ spec = do
         `shouldBe` "d2:id20:mnopqrstuvwxyz123456e"
 
     it "properly bencoded (iso)" $ property $ \ nid -> do
-      prop_bencode (Query    nid Ping)
+      prop_bencode (Query    nid False Ping)
       prop_bencode (Response nid Ping)
 
     it "does compatible with existing DHT" $ do
       nid <- genNodeId
-      Response _remoteAddr Ping <- rpc (query remoteAddr (Query nid Ping))
+      Response _remoteAddr Ping <- rpc (query remoteAddr (Query nid False Ping))
       return ()
 
   describe "find_node" $ do
     it "properly bencoded" $ do
       BE.decode "d2:id20:abcdefghij0123456789\
                  \6:target20:mnopqrstuvwxyz123456e"
-        `shouldBe` Right (Query "abcdefghij0123456789"
+        `shouldBe` Right (Query "abcdefghij0123456789" False
                       (FindNode "mnopqrstuvwxyz123456"))
 
-      BE.encode (Query "abcdefghij0123456789"
+      BE.encode (Query "abcdefghij0123456789" False
                  (FindNode "mnopqrstuvwxyz123456"))
         `shouldBe`
          "d2:id20:abcdefghij01234567896:target20:mnopqrstuvwxyz123456e"
@@ -109,13 +109,13 @@ spec = do
        `shouldBe` Right (Response nid (NodeFound [NodeInfo nid' naddr]))
 
     it "properly bencoded (iso)" $ property $ \ nid x xs -> do
-      prop_bencode (Query    nid (FindNode  x))
+      prop_bencode (Query    nid False (FindNode  x))
       prop_bencode (Response nid (NodeFound (xs :: [NodeInfo IPv4] )))
 
     it "does compatible with existing DHT" $ do
       nid <- genNodeId
       Response _remoteAddr (NodeFound xs) <- rpc $ do
-        query remoteAddr (Query nid (FindNode nid))
+        query remoteAddr (Query nid False (FindNode nid))
       L.length (xs :: [NodeInfo IPv4]) `shouldSatisfy` (> 0)
 
   describe "get_peers" $ do
@@ -123,7 +123,7 @@ spec = do
        BE.decode "d2:id20:abcdefghij0123456789\
                   \9:info_hash20:mnopqrstuvwxyz123456\
                   \e"
-         `shouldBe` Right (Query "abcdefghij0123456789"
+         `shouldBe` Right (Query "abcdefghij0123456789" False
                            (GetPeers "mnopqrstuvwxyz123456"))
 
        BE.decode "d2:id20:abcdefghij0123456789\
@@ -147,7 +147,7 @@ spec = do
              }))
 
     it "properly bencoded (iso)" $ property $ \ nid topic exs token -> do
-      prop_bencode (Query    nid (GetPeers topic))
+      prop_bencode (Query    nid False (GetPeers topic))
       let _ = exs :: Either [NodeInfo IPv4] [PeerAddr IPv4]
       let nullPeerId paddr = paddr {peerId = Nothing}
       let nullPeerIds      = either Left (Right . L.map nullPeerId)
@@ -156,7 +156,7 @@ spec = do
     it "does compatible with existing DHT" $ do
       nid <- genNodeId
       Response _remoteId (GotPeers {..})
-        <- rpc $ query remoteAddr (Query nid (GetPeers def))
+        <- rpc $ query remoteAddr (Query nid False (GetPeers def))
       let _ = peers :: Either [NodeInfo IPv4] [PeerAddr IPv4]
       either L.length L.length peers `shouldSatisfy` (> 0)
 
@@ -167,7 +167,7 @@ spec = do
                  \4:porti6881e\
                  \5:token8:aoeusnth\
                 \e" `shouldBe` Right
-        (Query "abcdefghij0123456789"
+        (Query "abcdefghij0123456789" False
                    (Announce False "mnopqrstuvwxyz123456" 6881 "aoeusnth"))
 
       BE.decode "d2:id20:abcdefghij0123456789\
@@ -176,7 +176,7 @@ spec = do
                  \4:porti6881e\
                  \5:token8:aoeusnth\
                 \e" `shouldBe` Right
-        (Query "abcdefghij0123456789"
+        (Query "abcdefghij0123456789" False
                    (Announce True "mnopqrstuvwxyz123456" 6881 "aoeusnth"))
 
 
@@ -185,27 +185,27 @@ spec = do
         (Response "mnopqrstuvwxyz123456" Announced)
 
     it "properly bencoded (iso)" $ property $ \ nid flag topic port token -> do
-      prop_bencode (Query    nid (Announce flag topic port token))
+      prop_bencode (Query    nid False (Announce flag topic port token))
       prop_bencode (Response nid (Announced))
 
 
     it "does compatible with existing DHT" $ do
       nid <- genNodeId
       Response _remoteId Announced <- rpc $ do
-        Response _ GotPeers {..} <- query remoteAddr (Query nid (GetPeers def))
+        Response _ GotPeers {..} <- query remoteAddr (Query nid False (GetPeers def))
         let _ = peers :: Either [NodeInfo IPv4] [PeerAddr IPv4]
-        query remoteAddr (Query nid (Announce False def thisPort grantedToken))
+        query remoteAddr (Query nid False (Announce False def thisPort grantedToken))
       return ()
 
     it "does fail on invalid token" $ do
       nid <- genNodeId
       (rpc $ do
-        Response _ GotPeers {..} <- query remoteAddr (Query nid (GetPeers def))
+        Response _ GotPeers {..} <- query remoteAddr (Query nid False (GetPeers def))
         let _ = peers :: Either [NodeInfo IPv4] [PeerAddr IPv4]
         let invalidToken = ""
         let q :: MonadKRPC h m => SockAddr -> Query Announce
               -> m (Response Announced)
             q = query
-        q remoteAddr (Query nid (Announce False def thisPort invalidToken)))
+        q remoteAddr (Query nid False (Announce False def thisPort invalidToken)))
           `shouldThrow` isQueryError
       return ()
diff --git a/tests/Network/BitTorrent/DHT/QuerySpec.hs b/tests/Network/BitTorrent/DHT/QuerySpec.hs
index 6b3b5485..93f78263 100644
--- a/tests/Network/BitTorrent/DHT/QuerySpec.hs
+++ b/tests/Network/BitTorrent/DHT/QuerySpec.hs
@@ -34,7 +34,7 @@ simpleDHT hs m =
 getBootInfo :: IO (NodeInfo IPv4)
 getBootInfo = do
   startAddr <- resolveHostName (L.head defaultBootstrapNodes)
-  simpleDHT [] $ pingQ startAddr
+  simpleDHT [] $ fmap fst (pingQ startAddr)
 
 spec :: Spec
 spec = parallel $ do
@@ -55,14 +55,14 @@ spec = parallel $ do
     it "findNode" $ do
       startInfo <- getBootInfo
       _ <- simpleDHT [] $ do
-        nid <- asks thisNodeId
+        nid <- myNodeIdAccordingTo (read "8.8.8.8:6881")
         findNodeQ nid startInfo
       return ()
 
     it "getPeers" $ do
       startInfo <- getBootInfo
       peers <- simpleDHT [] $ do
-        nid <- asks thisNodeId
+        nid <- myNodeIdAccordingTo (read "8.8.8.8:6881")
 
         -- we should not run getPeers query on boot node, because
         -- it may not support it
diff --git a/tests/Network/BitTorrent/DHT/RoutingSpec.hs b/tests/Network/BitTorrent/DHT/RoutingSpec.hs
index 3077a52a..6c0caad5 100644
--- a/tests/Network/BitTorrent/DHT/RoutingSpec.hs
+++ b/tests/Network/BitTorrent/DHT/RoutingSpec.hs
@@ -46,14 +46,17 @@ instance (Arbitrary ip, Eq ip) => Arbitrary (Table ip) where
 --    nodes   <- vector nodeN
 
     node    <- arbitrary
-    mt <- runSimulation  <$> arbitrary
-          <*> pure (T.insert node table)
+    mt <- do
+        env <- arbitrary
+        return $ runSimulation env $ do
+            (_,t') <- T.insert (currentTime env) (TryInsert node) table
+            return t' :: Routing ip (Table ip)
           --(foldM (flip fillTable) table nodes)
     return (fromJust mt)
-   where
-    fillTable x t = do
-      t' <- T.insert x t
-      return $ if T.full t' then t else t'
+--    where
+--     fillTable x t = do
+--       t' <- T.insert x t
+--       return $ if T.full t' then t else t'
 
 spec :: Spec
 spec = do
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
diff --git a/tests/Network/BitTorrent/DHTSpec.hs b/tests/Network/BitTorrent/DHTSpec.hs
index 30abc867..77160eb5 100644
--- a/tests/Network/BitTorrent/DHTSpec.hs
+++ b/tests/Network/BitTorrent/DHTSpec.hs
@@ -36,8 +36,8 @@ spec = do
     it "partial bootstrapping should finish in less than 10 seconds" $ do
       node <- resolveHostName (L.head defaultBootstrapNodes)
       res  <- timeout partialBootstrapTimeout $ do
-                dht opts def $ do
-                  bootstrap [node]
+                dht opts def fullLogging $ do
+                  bootstrap Nothing [node]
                   isBootstrapped
       res `shouldBe` Just True
 
@@ -57,4 +57,4 @@ spec = do
       pending
 
   describe "delete" $ do
-    return ()
\ No newline at end of file
+    return ()