+ Add some tests.

4 files changed, 36 insertions, 6 deletions

diff --git a/bittorrent.cabal b/bittorrent.cabal
index d911c71f..ebb8b62f 100644
--- a/bittorrent.cabal
+++ b/bittorrent.cabal
@@ -119,6 +119,7 @@ test-suite properties
                      , QuickCheck
 
                      , bittorrent
+                     , intset
 
   ghc-options:         -Wall -fno-warn-orphans
 
diff --git a/src/Data/Bitfield.hs b/src/Data/Bitfield.hs
index 024f8f71..3424a0c6 100644
--- a/src/Data/Bitfield.hs
+++ b/src/Data/Bitfield.hs
@@ -9,6 +9,7 @@
 --   bitfields. Bitfields are used to keep track indices of complete
 --   pieces either peer have or client have.
 --
+{-# LANGUAGE CPP #-}
 {-# LANGUAGE BangPatterns #-}
 {-# LANGUAGE RecordWildCards #-}
 module Data.Bitfield
@@ -32,8 +33,10 @@ module Data.Bitfield
        , getBitfield, putBitfield
        , bitfieldByteCount
 
+-- #ifdef TESTING
        , -- * Debug
          mkBitfield
+-- #endif
        ) where
 
 import Control.Monad
@@ -143,11 +146,13 @@ frequencies xs = runST $ do
   where
     size = maximum (map bfSize xs)
 
+-- TODO it seems like this operation is veeery slow
+
 -- | Find least available piece index. If no piece available return 'Nothing'.
 rarest :: [Bitfield] -> Maybe PieceIx
 rarest xs
     | V.null freqMap = Nothing
-    |     otherwise  = Just $ fst $ V.ifoldr minIx (0, freqMap V.! 0) freqMap
+    |     otherwise  = Just $ fst $ V.ifoldr' minIx (0, freqMap V.! 0) freqMap
   where
     freqMap = frequencies xs
 
@@ -210,5 +215,5 @@ bitfieldByteCount = error "bitfieldByteCount"
 mkBitfield :: PieceCount -> [PieceIx] -> Bitfield
 mkBitfield s ixs = Bitfield {
     bfSize = s
-  , bfSet  = S.splitLT s $ S.fromList ixs
+  , bfSet  = S.splitGT (-1) $ S.splitLT s $ S.fromList ixs
   }
\ No newline at end of file
diff --git a/tests/Encoding.hs b/tests/Encoding.hs
index 0b678a25..8afd3625 100644
--- a/tests/Encoding.hs
+++ b/tests/Encoding.hs
@@ -35,7 +35,7 @@ instance Arbitrary Block where
   arbitrary = Block <$> positive <*> positive <*> arbitrary
 
 instance Arbitrary Bitfield where
-  arbitrary = mkBitfield <$> positive <*> arbitrary
+  arbitrary = mkBitfield <$> (succ <$> positive) <*> arbitrary
 
 instance Arbitrary PortNumber where
   arbitrary = fromIntegral <$> (arbitrary :: Gen Word16)
diff --git a/tests/Main.hs b/tests/Main.hs
index bc3f7809..ef4ab8a7 100644
--- a/tests/Main.hs
+++ b/tests/Main.hs
@@ -1,14 +1,38 @@
 {-# LANGUAGE OverloadedStrings #-}
 module Main (main) where
 
+import Control.Applicative
+import Data.IntervalSet
+import Data.List as L
+import Data.Ord
 import Data.Word
+
 import Test.Framework (defaultMain)
 import Test.Framework.Providers.QuickCheck2 (testProperty)
+import Test.QuickCheck
+
+import Data.Bitfield as BF
+import Network.BitTorrent as BT
 
 import Encoding
 
-import Data.Bitfield as BT
-import Network.BitTorrent as BT
+
+
+instance Arbitrary IntSet where
+  arbitrary = fromList <$> arbitrary
+
+prop_completenessRange :: Bitfield -> Bool
+prop_completenessRange bf = 0 <= c && c <= 1
+  where
+    c = completeness bf
+
+prop_rarestInRange :: [Bitfield] -> Bool
+prop_rarestInRange xs = case rarest xs of
+  Just r  -> 0 <= r && r < totalCount (maximumBy (comparing totalCount) xs)
+  Nothing -> True
 
 main :: IO ()
-main = defaultMain []
+main = defaultMain
+  [ testProperty "completeness range" prop_completenessRange
+  , testProperty "rarest in range"    prop_rarestInRange
+  ]