2 files changed, 17 insertions, 1 deletions
diff --git a/test/ArbitraryMove.hs b/test/ArbitraryMove.hs
index bcde156..df3d7fb 100644
--- a/test/ArbitraryMove.hs
+++ b/test/ArbitraryMove.hs
@@ -21,7 +21,7 @@ type Move = String
 arbitraryMove :: Gen Move
 arbitraryMove = frequency
   [ (2, probableArbitraryMove)
-  , (1, invalidMove)
+  --, (1, invalidMove)
   ]
   where
     invalidMove :: Gen Move
diff --git a/test/Test.hs b/test/Test.hs
index 1f0a02d..c5c694e 100644
--- a/test/Test.hs
+++ b/test/Test.hs
@@ -2,6 +2,7 @@
 -- vim: sw=2 ts=2 et ai:
 module Test where
 
+import Data.Array
 import Data.Attoparsec.ByteString.Char8
 import Data.ByteString.Char8 (pack)
 
@@ -18,6 +19,14 @@ import Chess.PGN
 
 import ArbitraryMove
 
+atAnyState :: (Board -> Bool) -> PGN -> Bool
+atAnyState ok pgn = and [ok b | Right b <- seqList [moveSAN m | m <- moves pgn] $ Right defaultBoard]
+  where
+    seqList :: [r -> Either e r] -> Either e r -> [Either e r]
+    seqList (f:fs) (Right io) = f io:seqList fs (f io)
+    seqList _      (Left _)   = []
+    seqList []     _          = []
+
 prop_checkPGN :: PGN -> Bool
 prop_checkPGN pgn
   | isLeft parsed       = False
@@ -59,5 +68,12 @@ makePGN event site date round white black result moves =
     tagPair :: String -> String -> String
     tagPair tag val = "[" ++ tag ++ " \"" ++ val ++ "\"]\r\n"
 
+prop_only_2_kings :: PGN -> Bool
+prop_only_2_kings = atAnyState (\b ->
+  length (pieceCoords Black b King) == 1 &&
+  length (pieceCoords White b King) == 1)
+  where
+    pieceCoords clr brd piece = [i | (i, pc) <- (assocs $ board brd), pc == Just (Piece clr piece)]
+
 return []
 main = $forAllProperties (quickCheckWithResult (stdArgs {maxSuccess=10000}))