2 files changed, 25 insertions, 9 deletions

diff --git a/tests/encoding.hs b/tests/encoding.hs
index b820dc26..e1a8b63e 100644
--- a/tests/encoding.hs
+++ b/tests/encoding.hs
@@ -11,7 +11,9 @@ import Test.Framework (defaultMain)
 import Test.Framework.Providers.QuickCheck2 (testProperty)
 import Test.QuickCheck
 
+import Data.Torrent
 import Network.Torrent
+import Network.URI
 
 positive :: Gen Int
 positive = fromIntegral <$> (arbitrary :: Gen Word32)
@@ -45,16 +47,30 @@ instance Arbitrary PeerID where
                            <*> arbitrary
                            <*> arbitrary
 
+instance Arbitrary InfoHash where
+  arbitrary = (hash . B.pack) <$> vectorOf 20 arbitrary
+
 instance Arbitrary Handshake where
-  arbitrary = defaultHandshake
-                <$> (B.pack <$> (vectorOf 20 arbitrary))
-                <*> arbitrary
+  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
 
+test_scrape_url :: Bool
+test_scrape_url = check `all` 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"  , Nothing)
+      , ("http://example.com/x%064announce"   , Nothing)
+      ]
 
 main :: IO ()
 main = do
@@ -67,4 +83,7 @@ main = do
 
        , testProperty "Handshake encode <-> decode" $
             prop_encoding (T :: T Handshake)
+
+       , testProperty "Scrape URL" $
+            test_scrape_url
        ]
\ No newline at end of file
diff --git a/tests/info-hash.hs b/tests/info-hash.hs
index 07c01c9c..1638a205 100644
--- a/tests/info-hash.hs
+++ b/tests/info-hash.hs
@@ -1,23 +1,20 @@
 {-# LANGUAGE OverloadedStrings #-}
 module Main (main) where
 
-import Data.Foldable
 import           Data.ByteString (ByteString)
 import qualified Data.ByteString as B
 import qualified Data.ByteString.Char8 as BC
 import qualified Data.ByteString.Lazy as L
-import qualified Data.ByteString.Builder as B
-import qualified Data.ByteString.Builder.Prim as B
+
 import Data.Torrent
+import Data.Torrent.InfoHash
+
 import System.Environment
 import System.Exit
 
 checkInfo :: ByteString
 checkInfo = "0221caf96aa3cb94f0f58d458e78b0fc344ad8bf"
 
-ppHex :: B.ByteString -> B.ByteString
-ppHex = L.toStrict . B.toLazyByteString . foldMap (B.primFixed B.word8HexFixed) . B.unpack
-
 torrentFileName :: String
 torrentFileName = "tests/dapper-dvd-amd64.iso.torrent"