~ Fix bitfield encoding.

1 files changed, 10 insertions, 2 deletions

diff --git a/tests/Main.hs b/tests/Main.hs
index 0e18a06b..84870d66 100644
--- a/tests/Main.hs
+++ b/tests/Main.hs
@@ -185,6 +185,15 @@ instance Arbitrary Message where
     , Cancel   <$> arbitrary
     , Port     <$> arbitrary
     ]
+-- todo add all messages
+
+prop_messageEncoding :: Message -> Bool
+prop_messageEncoding msg @ (Bitfield bf)
+  = case S.decode (S.encode msg) of
+      Right (Bitfield bf') -> bf == adjustSize (totalCount bf) bf'
+      Left  _   -> False
+prop_messageEncoding msg
+  = S.decode (S.encode msg) == Right msg
 
 {-----------------------------------------------------------------------
     Main
@@ -209,7 +218,6 @@ main = defaultMain $
     -- handshake module
   , testProperty "handshake encoding" $
       prop_cerealEncoding  (T :: T Handshake)
-  , testProperty "message encoding" $
-      prop_cerealEncoding (T :: T Message)
+  , testProperty "message encoding" prop_messageEncoding
 
   ] ++ test_scrape_url