2 files changed, 45 insertions, 0 deletions

diff --git a/bittorrent.cabal b/bittorrent.cabal
index a457a9a0..3ad1d7e4 100644
--- a/bittorrent.cabal
+++ b/bittorrent.cabal
@@ -163,6 +163,7 @@ test-suite spec
                        Network.BitTorrent.Tracker.MessageSpec
                        Network.BitTorrent.Tracker.RPC.HTTPSpec
                        Network.BitTorrent.Tracker.RPC.UDPSpec
+                       Network.BitTorrent.Exchange.MessageSpec
                        System.Torrent.FileMapSpec
   build-depends:       base == 4.*
                      , bytestring
@@ -172,6 +173,7 @@ test-suite spec
                      , convertible
                      , data-default
                      , monad-loops
+                     , containers
 
                      , aeson
                      , cereal
diff --git a/tests/Network/BitTorrent/Exchange/MessageSpec.hs b/tests/Network/BitTorrent/Exchange/MessageSpec.hs
new file mode 100644
index 00000000..8d1041dd
--- /dev/null
+++ b/tests/Network/BitTorrent/Exchange/MessageSpec.hs
@@ -0,0 +1,43 @@
+module Network.BitTorrent.Exchange.MessageSpec (spec) where
+import Control.Applicative
+import Data.ByteString as BS
+import Data.Default
+import Data.List as L
+import Data.Set as S
+import Data.Serialize as S
+import Test.Hspec
+import Test.QuickCheck
+
+import Data.Torrent.InfoHashSpec ()
+import Network.BitTorrent.CoreSpec ()
+import Network.BitTorrent.Core
+import Network.BitTorrent.Exchange.Message
+
+instance Arbitrary Extension where
+  arbitrary = elements [minBound .. maxBound]
+
+instance Arbitrary Caps where
+  arbitrary = toCaps <$> arbitrary
+
+instance Arbitrary Handshake where
+  arbitrary = Handshake <$> arbitrary <*> arbitrary
+                        <*> arbitrary <*> arbitrary
+
+spec :: Spec
+spec = do
+  describe "Caps" $ do
+    it "set-like container" $ property $ \ exts ->
+      L.all (`allowed` (toCaps exts :: Caps)) exts
+
+    it "preserve items" $ property $ \ extSet ->
+      S.fromList (fromCaps (toCaps (S.toList extSet) :: Caps))
+        `shouldBe` extSet
+
+  describe "Handshake" $ do
+    it "properly serialized" $ property $ \ hs ->
+      S.decode (S.encode hs ) `shouldBe` Right (hs :: Handshake)
+
+    it "fail if protocol string is too long" $ do
+      pid <- genPeerId
+      let hs = (defaultHandshake def pid) {hsProtocol = BS.replicate 256 0}
+      S.decode (S.encode hs) `shouldBe` Right hs
\ No newline at end of file