1 files changed, 15 insertions, 5 deletions

diff --git a/tests/Network/BitTorrent/DHT/MessageSpec.hs b/tests/Network/BitTorrent/DHT/MessageSpec.hs
index 3b1f5acd..6f3c7489 100644
--- a/tests/Network/BitTorrent/DHT/MessageSpec.hs
+++ b/tests/Network/BitTorrent/DHT/MessageSpec.hs
@@ -21,6 +21,16 @@ import Data.TorrentSpec                 ()
 import Network.BitTorrent.CoreSpec      ()
 import Network.BitTorrent.DHT.TokenSpec ()
 
+-- Arbitrary queries and responses.
+instance Arbitrary Ping where arbitrary = pure Ping
+instance Arbitrary FindNode where arbitrary = FindNode <$> arbitrary
+instance Arbitrary ip => Arbitrary (NodeFound ip) where arbitrary = NodeFound <$> arbitrary
+instance Arbitrary GetPeers where arbitrary = GetPeers <$> arbitrary
+instance Arbitrary ip => Arbitrary (GotPeers ip) where arbitrary =  GotPeers <$> arbitrary <*> arbitrary
+instance Arbitrary Announce where arbitrary = Announce <$> arbitrary <*> arbitrary <*> arbitrary <*> arbitrary <*> arbitrary
+instance Arbitrary Announced where arbitrary = pure Announced
+instance Arbitrary x => Arbitrary (Query x) where arbitrary = Query <$> arbitrary <*> arbitrary <*> arbitrary
+instance Arbitrary x => Arbitrary (Response x) where arbitrary = Response <$> arbitrary <*> arbitrary
 
 instance MonadLogger IO where
   monadLoggerLog _ _ _ _ = return ()
@@ -168,7 +178,7 @@ spec = do
                  \5:token8:aoeusnth\
                 \e" `shouldBe` Right
         (Query "abcdefghij0123456789" False
-                   (Announce False "mnopqrstuvwxyz123456" 6881 "aoeusnth"))
+                   (Announce False "mnopqrstuvwxyz123456" Nothing 6881 "aoeusnth"))
 
       BE.decode "d2:id20:abcdefghij0123456789\
                  \12:implied_porti1e\
@@ -177,7 +187,7 @@ spec = do
                  \5:token8:aoeusnth\
                 \e" `shouldBe` Right
         (Query "abcdefghij0123456789" False
-                   (Announce True "mnopqrstuvwxyz123456" 6881 "aoeusnth"))
+                   (Announce True "mnopqrstuvwxyz123456" Nothing 6881 "aoeusnth"))
 
 
       BE.decode "d2:id20:mnopqrstuvwxyz123456e"
@@ -185,7 +195,7 @@ spec = do
         (Response "mnopqrstuvwxyz123456" Announced)
 
     it "properly bencoded (iso)" $ property $ \ nid flag topic port token -> do
-      prop_bencode (Query    nid False (Announce flag topic port token))
+      prop_bencode (Query    nid False (Announce flag topic Nothing port token))
       prop_bencode (Response nid (Announced))
 
 
@@ -194,7 +204,7 @@ spec = do
       Response _remoteId Announced <- rpc $ do
         Response _ GotPeers {..} <- query remoteAddr (Query nid False (GetPeers def))
         let _ = peers :: Either [NodeInfo IPv4] [PeerAddr IPv4]
-        query remoteAddr (Query nid False (Announce False def thisPort grantedToken))
+        query remoteAddr (Query nid False (Announce False def Nothing thisPort grantedToken))
       return ()
 
     it "does fail on invalid token" $ do
@@ -206,6 +216,6 @@ spec = do
         let q :: MonadKRPC h m => SockAddr -> Query Announce
               -> m (Response Announced)
             q = query
-        q remoteAddr (Query nid False (Announce False def thisPort invalidToken)))
+        q remoteAddr (Query nid False (Announce False def Nothing thisPort invalidToken)))
           `shouldThrow` isQueryError
       return ()