3 files changed, 50 insertions, 7 deletions

diff --git a/bittorrent.cabal b/bittorrent.cabal
index e987a4be..4cb9fce7 100644
--- a/bittorrent.cabal
+++ b/bittorrent.cabal
@@ -170,6 +170,7 @@ test-suite spec
                        Network.BitTorrent.Core.NodeSpec
                        Network.BitTorrent.Core.PeerAddrSpec
                        Network.BitTorrent.Core.PeerIdSpec
+                       Network.BitTorrent.DHT.MessageSpec
                        Network.BitTorrent.Tracker.MessageSpec
                        Network.BitTorrent.Tracker.RPCSpec
                        Network.BitTorrent.Tracker.RPC.HTTPSpec
diff --git a/src/Network/BitTorrent/DHT/Message.hs b/src/Network/BitTorrent/DHT/Message.hs
index 85abf019..9000a9be 100644
--- a/src/Network/BitTorrent/DHT/Message.hs
+++ b/src/Network/BitTorrent/DHT/Message.hs
@@ -55,7 +55,7 @@ node_id_key = "id"
 data Query a = Query
   { thisNodeId  :: NodeId
   , queryParams :: a
-  }
+  } deriving (Show, Eq)
 
 instance BEncode a => BEncode (Query a) where
   toBEncode Query {..} = toDict $
@@ -75,7 +75,7 @@ instance BEncode a => BEncode (Query a) where
 data Response a = Response
   { remoteNodeId :: NodeId
   , responseVals :: a
-  }
+  } deriving (Show, Eq)
 
 instance BEncode a => BEncode (Response a) where
   toBEncode = toBEncode . toQuery
@@ -91,9 +91,10 @@ instance BEncode a => BEncode (Response a) where
 -- ping method
 -----------------------------------------------------------------------}
 
--- | The most basic query is a ping.
+-- | The most basic query is a ping. Ping query is used to check if a
+-- quered node is still alive.
 data Ping = Ping
-  deriving Typeable
+  deriving (Show, Eq, Typeable)
 
 instance BEncode Ping where
   toBEncode Ping = toDict endDict
@@ -110,7 +111,7 @@ instance KRPC (Query Ping) [Ping] where
 -- | Find node is used to find the contact information for a node
 -- given its ID.
 newtype FindNode = FindNode NodeId
-  deriving Typeable
+  deriving (Show, Eq, Typeable)
 
 target_key :: BKey
 target_key = "target"
@@ -119,12 +120,12 @@ instance BEncode FindNode where
   toBEncode (FindNode nid) = toDict   $ target_key .=! nid .: endDict
   fromBEncode              = fromDict $ FindNode  <$>! target_key
 
--- | When a node receives a find_node query, it should respond with a
+-- | When a node receives a 'FindNode' query, it should respond with a
 -- the compact node info for the target node or the K (8) closest good
 -- nodes in its own routing table.
 --
 newtype NodeFound ip = NodeFound [NodeInfo ip]
-  deriving Typeable
+  deriving (Show, Eq, Typeable)
 
 nodes_key :: BKey
 nodes_key = "nodes"
diff --git a/tests/Network/BitTorrent/DHT/MessageSpec.hs b/tests/Network/BitTorrent/DHT/MessageSpec.hs
new file mode 100644
index 00000000..e8a76f14
--- /dev/null
+++ b/tests/Network/BitTorrent/DHT/MessageSpec.hs
@@ -0,0 +1,41 @@
+module Network.BitTorrent.DHT.MessageSpec (spec) where
+import Test.Hspec
+import Data.BEncode as BE
+import Network.BitTorrent.Core
+import Network.BitTorrent.DHT.Message
+
+
+spec :: Spec
+spec = do
+  describe "ping" $ do
+    it "properly bencoded" $ do
+      BE.decode "d2:id20:abcdefghij0123456789e"
+        `shouldBe` Right (Query "abcdefghij0123456789" Ping)
+
+      BE.encode (Query "abcdefghij0123456789" Ping)
+        `shouldBe` "d2:id20:abcdefghij0123456789e"
+
+      BE.decode "d2:id20:mnopqrstuvwxyz123456e"
+        `shouldBe` Right (Response "mnopqrstuvwxyz123456" Ping)
+
+      BE.encode (Response "mnopqrstuvwxyz123456" Ping)
+        `shouldBe` "d2:id20:mnopqrstuvwxyz123456e"
+
+  describe "find_node" $ do
+    it "properly bencoded" $ do
+      BE.decode "d2:id20:abcdefghij0123456789\
+                 \6:target20:mnopqrstuvwxyz123456e"
+        `shouldBe` Right (Query "abcdefghij0123456789"
+                      (FindNode "mnopqrstuvwxyz123456"))
+
+      BE.encode (Query "abcdefghij0123456789"
+                 (FindNode "mnopqrstuvwxyz123456"))
+        `shouldBe`
+         "d2:id20:abcdefghij01234567896:target20:mnopqrstuvwxyz123456e"
+
+      let addr = "127.0.0.1:256" :: NodeAddr IPv4
+      let nid  = "0123456789abcdefghij"
+      BE.decode "d2:id20:0123456789abcdefghij\
+                 \5:nodes26:mnopqrstuvwxyz123456\127\0\0\1\&56\
+                 \e"
+       `shouldBe` Right (Response nid (NodeFound [NodeInfo nid addr]))
\ No newline at end of file