Handle decoding multiple ASCII armor messages from a single bytestream.

1 files changed, 11 insertions, 7 deletions

diff --git a/tests/suite.hs b/tests/suite.hs
index 272295a..45a2fb2 100644
--- a/tests/suite.hs
+++ b/tests/suite.hs
@@ -3,7 +3,7 @@ import Test.Framework.Providers.HUnit
 
 import Test.HUnit
 
-import Codec.Encryption.OpenPGP.ASCIIArmor (armor, decodeArmor)
+import Codec.Encryption.OpenPGP.ASCIIArmor (armor, decode)
 import Codec.Encryption.OpenPGP.ASCIIArmor.Types
 
 import Data.ByteString (ByteString)
@@ -15,13 +15,15 @@ import Data.Word (Word32)
 testCRC24 :: ByteString -> Word32 -> Assertion
 testCRC24 bs crc = assertEqual "crc24" crc (crc24 bs)
 
-testArmorDecode :: FilePath -> FilePath -> Assertion
-testArmorDecode fp target = do
+testArmorDecode :: FilePath -> [FilePath] -> Assertion
+testArmorDecode fp targets = do
     bs <- B.readFile $ "tests/data/" ++ fp
-    tbs <- B.readFile $ "tests/data/" ++ target
-    case decodeArmor bs of
+    tbss <- mapM (\target -> B.readFile $ "tests/data/" ++ target) targets
+    case decode bs of
         Left e -> assertFailure $ "Decode failed (" ++ e ++ ") on " ++ fp
-        Right (Armor at hdrs pl) -> do assertEqual ("for " ++ fp) tbs pl
+        Right as -> do assertEqual ("for " ++ fp) tbss (map getPayload as)
+    where
+        getPayload (Armor _ _ pl) = pl
 
 testArmorEncode :: FilePath -> FilePath -> Assertion
 testArmorEncode fp target = do
@@ -36,7 +38,9 @@ tests = [
     , testCase "CRC24: hOpenPGP and friends" (testCRC24 (BC8.pack "hOpenPGP and friends") 11940960)
     ]
  , testGroup "ASCII armor" [
-      testCase "Decode sample armor" (testArmorDecode "msg1.asc" "msg1.gpg")
+      testCase "Decode sample armor" (testArmorDecode "msg1.asc" ["msg1.gpg"])
+    , testCase "Decode sample armor with cruft" (testArmorDecode "msg1a.asc" ["msg1.gpg"])
+    , testCase "Decode multiple sample armors" (testArmorDecode "msg1b.asc" ["msg1.gpg","msg1.gpg","msg1.gpg"])
     , testCase "Encode sample armor" (testArmorEncode "msg1.gpg" "msg1.asc")
     ]
  ]