4 files changed, 125 insertions, 112 deletions

diff --git a/bittorrent.cabal b/bittorrent.cabal
index 223584fc..da69e48d 100644
--- a/bittorrent.cabal
+++ b/bittorrent.cabal
@@ -114,10 +114,7 @@ test-suite info-hash
 test-suite properties
   type:                exitcode-stdio-1.0
   main-is:             Main.hs
-  other-modules:       Encoding
-
   hs-source-dirs:      tests
-
   build-depends:       base       == 4.*
                      , bytestring >= 0.10.2
                      , cereal     >= 0.3.5.2
@@ -130,7 +127,6 @@ test-suite properties
 
                      , bencoding
                      , bittorrent
-                     , intset
 
   ghc-options:         -Wall -fno-warn-orphans
 
diff --git a/src/Network/BitTorrent/Tracker.hs b/src/Network/BitTorrent/Tracker.hs
index 2319a551..2e599002 100644
--- a/src/Network/BitTorrent/Tracker.hs
+++ b/src/Network/BitTorrent/Tracker.hs
@@ -112,7 +112,8 @@ regularReq numWant ses pr = (genericReq ses pr) {
   , reqEvent      = Nothing
   }
 
--- | Must be sent to the tracker if the client is shutting down gracefully.
+-- | Must be sent to the tracker if the client is shutting down
+-- gracefully.
 --
 stoppedReq :: TConnection -> Progress -> TRequest
 stoppedReq ses pr = (genericReq ses pr) {
@@ -122,7 +123,8 @@ stoppedReq ses pr = (genericReq ses pr) {
   }
 
 -- | Must be sent to the tracker when the download completes.
---   However, must not be sent if the download was already 100% complete.
+-- However, must not be sent if the download was already 100%
+-- complete.
 --
 completedReq :: TConnection -> Progress -> TRequest
 completedReq ses pr = (genericReq ses pr) {
@@ -141,7 +143,9 @@ data TSession = TSession {
   }
 
 newSession :: Progress -> Int -> [PeerAddr] -> IO TSession
-newSession pr i ps = TSession <$> newTVarIO pr <*> newIORef i <*> newTVarIO ps
+newSession pr i ps = TSession <$> newTVarIO pr
+                              <*> newIORef i
+                              <*> newTVarIO psx
 
 getPeerList :: TSession -> IO [PeerAddr]
 getPeerList = readTVarIO . sePeers
diff --git a/tests/Encoding.hs b/tests/Encoding.hs
deleted file mode 100644
index 78f0dfc1..00000000
--- a/tests/Encoding.hs
+++ /dev/null
@@ -1,96 +0,0 @@
-{-# LANGUAGE StandaloneDeriving #-}
-{-# LANGUAGE GeneralizedNewtypeDeriving #-}
-{-# OPTIONS -fno-warn-orphans #-}
-module Encoding where
-
-import Control.Applicative
-import Data.Word
-import           Data.ByteString (ByteString)
-import qualified Data.ByteString as B
-import Data.Serialize
-import Test.Framework (Test)
-import Test.Framework.Providers.QuickCheck2 (testProperty)
-import Test.QuickCheck
-
-
-import Network.URI
-import Network
-
-
-import Data.Bitfield
-import Data.Torrent
-import Network.BitTorrent
-
-
-positive :: Gen Int
-positive = fromIntegral <$> (arbitrary :: Gen Word32)
-
-instance Arbitrary ByteString where
-  arbitrary = B.pack <$> arbitrary
-
-instance Arbitrary BlockIx where
-  arbitrary = BlockIx <$> positive <*> positive <*> positive
-
-instance Arbitrary Block where
-  arbitrary = Block <$> positive <*> positive <*> arbitrary
-
-instance Arbitrary Bitfield where
-  arbitrary = mkBitfield <$> (succ . min 1000 <$> positive)
-                         <*> arbitrary
-
-instance Arbitrary PortNumber where
-  arbitrary = fromIntegral <$> (arbitrary :: Gen Word16)
-
-instance Arbitrary Message where
-  arbitrary = oneof
-    [ pure KeepAlive
-    , pure Choke
-    , pure Unchoke
-    , pure Interested
-    , pure NotInterested
-    , Have     <$> positive
-    , Bitfield <$> arbitrary
-    , Request  <$> arbitrary
-    , Piece    <$> arbitrary
-    , Cancel   <$> arbitrary
-    , Port     <$> arbitrary
-    ]
-
-instance Arbitrary PeerID where
-  arbitrary = azureusStyle <$> pure defaultClientID
-                           <*> arbitrary
-                           <*> arbitrary
-
-instance Arbitrary InfoHash where
-  arbitrary = (hash . B.pack) <$> arbitrary
-
-instance Arbitrary Handshake where
-  arbitrary = defaultHandshake <$> arbitrary <*> arbitrary
-
-
-data T a = T
-
-prop_encoding :: (Serialize a, Eq a) => T a -> [a] -> Bool
-prop_encoding _ msgs = decode (encode msgs) == Right msgs
-
--- | Note that in 6 esample we intensionally do not agree with specification,
---   because taking in account '/' in query parameter seems to be meaningless.
---   (And thats because other clients do not chunk uri by parts)
---   Moreover in practice there should be no difference. (I hope)
---
-test_scrape_url :: [Test]
-test_scrape_url = zipWith mkTest [1 :: Int ..] (check `map` tests)
-  where
-    check (iu, ou) = (parseURI iu >>= (`scrapeURL` []) >>= return . show) == ou
-    tests =
-      [ ("http://example.com/announce"        , Just "http://example.com/scrape")
-      , ("http://example.com/x/announce"      , Just "http://example.com/x/scrape")
-      , ("http://example.com/announce.php"    , Just "http://example.com/scrape.php")
-      , ("http://example.com/a"               , Nothing)
-      , ("http://example.com/announce?x2%0644", Just "http://example.com/scrape?x2%0644")
-      , ("http://example.com/announce?x=2/4"  , Just "http://example.com/scrape?x=2/4")
---      , ("http://example.com/announce?x=2/4"  , Nothing) -- by specs
-      , ("http://example.com/x%064announce"   , Nothing)
-      ]
-
-    mkTest i = testProperty ("scrape test #" ++ show i)
diff --git a/tests/Main.hs b/tests/Main.hs
index ff571b6b..0e18a06b 100644
--- a/tests/Main.hs
+++ b/tests/Main.hs
@@ -1,28 +1,35 @@
-{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE StandaloneDeriving         #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE OverloadedStrings          #-}
 module Main (main) where
 
 import Control.Applicative
+import           Data.ByteString (ByteString)
+import qualified Data.ByteString as B
 import qualified Data.ByteString.Lazy as Lazy
-import Data.IntervalSet
 import Data.List as L
 import Data.Ord
 import Data.Maybe
 import Data.Word
+import Data.Serialize as S
 import Data.Text as T
+
+import Network
 import Network.URI
 
-import Test.Framework (defaultMain)
+import Test.Framework (Test, defaultMain)
 import Test.Framework.Providers.QuickCheck2 (testProperty)
 import Test.QuickCheck
 
-import Data.BEncode
+import Data.BEncode as BE
 import Data.Bitfield as BF
 import Data.Torrent
 import Network.BitTorrent as BT
 
-import Debug.Trace
-import Encoding
+-- import Debug.Trace
+
 
+data T a = T
 
 instance Arbitrary URI where
   arbitrary = pure $ fromJust
@@ -90,17 +97,119 @@ instance Arbitrary Torrent where
                  <*> arbitrary <*> arbitrary <*> arbitrary
                  <*> arbitrary <*> pure Nothing <*> arbitrary
 
+{-----------------------------------------------------------------------
+    Handshake
+-----------------------------------------------------------------------}
+
+instance Arbitrary PeerID where
+  arbitrary = azureusStyle <$> pure defaultClientID
+                           <*> arbitrary
+                           <*> arbitrary
+
+instance Arbitrary InfoHash where
+  arbitrary = (hash . B.pack) <$> arbitrary
+
+instance Arbitrary Handshake where
+  arbitrary = defaultHandshake <$> arbitrary <*> arbitrary
+
+prop_cerealEncoding :: (Serialize a, Eq a) => T a -> [a] -> Bool
+prop_cerealEncoding _ msgs = S.decode (S.encode msgs) == Right msgs
+
+{-----------------------------------------------------------------------
+    Tracker/Scrape
+-----------------------------------------------------------------------}
+
+-- | Note that in 6 esample we intensionally do not agree with
+-- specification, because taking in account '/' in query parameter
+-- seems to be meaningless.  (And thats because other clients do not
+-- chunk uri by parts) Moreover in practice there should be no
+-- difference. (I think so)
+--
+test_scrape_url :: [Test]
+test_scrape_url = L.zipWith mkTest [1 :: Int ..] (check `L.map` tests)
+  where
+    check (iu, ou) = (parseURI iu >>= (`scrapeURL` [])
+                      >>= return . show) == ou
+    tests =
+      [ (      "http://example.com/announce"
+        , Just "http://example.com/scrape")
+      , (      "http://example.com/x/announce"
+        , Just "http://example.com/x/scrape")
+      , (      "http://example.com/announce.php"
+        , Just "http://example.com/scrape.php")
+      , (      "http://example.com/a"    , Nothing)
+      , (      "http://example.com/announce?x2%0644"
+        , Just "http://example.com/scrape?x2%0644")
+      , (      "http://example.com/announce?x=2/4"
+        , Just "http://example.com/scrape?x=2/4")
+--      , ("http://example.com/announce?x=2/4"  , Nothing) -- by specs
+      , ("http://example.com/x%064announce"   , Nothing)
+      ]
+
+    mkTest i = testProperty ("scrape test #" ++ show i)
+
+{-----------------------------------------------------------------------
+    P2P/message
+-----------------------------------------------------------------------}
+
+positive :: Gen Int
+positive = fromIntegral <$> (arbitrary :: Gen Word32)
+
+instance Arbitrary ByteString where
+  arbitrary = B.pack <$> arbitrary
+
+instance Arbitrary BlockIx where
+  arbitrary = BlockIx <$> positive <*> positive <*> positive
+
+instance Arbitrary Block where
+  arbitrary = Block <$> positive <*> positive <*> arbitrary
+
+instance Arbitrary Bitfield where
+  arbitrary = mkBitfield <$> (succ . min 1000 <$> positive)
+                         <*> arbitrary
+
+instance Arbitrary PortNumber where
+  arbitrary = fromIntegral <$> (arbitrary :: Gen Word16)
+
+instance Arbitrary Message where
+  arbitrary = oneof
+    [ pure KeepAlive
+    , pure Choke
+    , pure Unchoke
+    , pure Interested
+    , pure NotInterested
+    , Have     <$> positive
+    , Bitfield <$> arbitrary
+    , Request  <$> arbitrary
+    , Piece    <$> arbitrary
+    , Cancel   <$> arbitrary
+    , Port     <$> arbitrary
+    ]
+
+{-----------------------------------------------------------------------
+    Main
+-----------------------------------------------------------------------}
+
 main :: IO ()
-main = defaultMain
-  [ testProperty "completeness range"      prop_completenessRange
+main = defaultMain $
+  [ -- bitfield module
+    testProperty "completeness range"      prop_completenessRange
   , testProperty "rarest in range"         prop_rarestInRange
   , testProperty "min less that max"       prop_minMax
   , testProperty "difference de morgan"    prop_differenceDeMorgan
 
+    -- torrent module
   , testProperty "file info encoding"      $
       prop_properBEncode (T :: T FileInfo)
   , testProperty "content info encoding"   $
       prop_properBEncode (T :: T ContentInfo)
   , testProperty "torrent encoding" $
       prop_properBEncode (T :: T Torrent)
-  ]
+
+    -- handshake module
+  , testProperty "handshake encoding" $
+      prop_cerealEncoding  (T :: T Handshake)
+  , testProperty "message encoding" $
+      prop_cerealEncoding (T :: T Message)
+
+  ] ++ test_scrape_url