1 files changed, 12 insertions, 4 deletions
diff --git a/test/Test.hs b/test/Test.hs
index e07fcc5..7155d84 100644
--- a/test/Test.hs
+++ b/test/Test.hs
@@ -8,6 +8,7 @@ import Data.ByteString.Char8 (pack)
 
 import Data.Either
 import Data.List
+import Data.Maybe
 
 import Control.Monad
 
@@ -77,12 +78,19 @@ makePGN event site date round white black result moves =
     tagPair :: String -> String -> String
     tagPair tag val = "[" ++ tag ++ " \"" ++ val ++ "\"]\r\n"
 
+-- From Chess.hs, with adaptations
+pieceCoords :: Maybe Color -> Maybe PieceType -> Board -> [((Int,Int),Piece)]
+pieceCoords clr piece brd = [(i,pc) | (i,Just pc) <- (assocs $ board brd), match pc]
+  where
+    match :: Piece -> Bool
+    match (Piece clr' piece') =
+      (isJust clr   && clr'   == fromJust clr   || isNothing clr) &&
+      (isJust piece && piece' == fromJust piece || isNothing piece)
+
 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)]
+  length (pieceCoords (Just Black) (Just King) b) == 1 &&
+  length (pieceCoords (Just White) (Just King) b) == 1)
 
 prop_number_of_pieces :: PGN -> Bool
 prop_number_of_pieces = atAnyTwoStates (\b1 b2 ->