Fix sufficient material calculation (#296)

* fix sufficient material calculation

* tests

* remove prints

kaggle_environments/envs/chess/chess.py

@@ -100,11 +100,11 @@ def sufficient_material(pieces):
     """Checks if the given pieces are sufficient for checkmate."""
     if pieces['q'] > 0 or pieces['r'] > 0 or pieces['p'] > 0:
         return True
-    if pieces['n'] + pieces['b'] >= 3:
+    # we only have knights and bishops left on this team
+    knight_bishop_count = pieces['n'] + pieces['b']
+    if knight_bishop_count >= 3:
         return True
-    # TODO: they have to be opposite color bishops
-    # b/b or n/b can checkmate, n/n cannot
-    if pieces['n'] < 2:
+    if knight_bishop_count == 2 and pieces['b'] >= 1:
         return True
 
     return False
@@ -116,7 +116,7 @@ def is_insufficient_material(board):
 
     for square in range(64):
         piece = board.get_piece(square)
-        if piece:
+        if piece and piece != " ":
             if piece.isupper():
                 white_pieces[piece.lower()] += 1
             else:

kaggle_environments/envs/chess/test_chess.py

@@ -1,4 +1,6 @@
 from kaggle_environments import make
+from Chessnut import Game
+from chess import is_insufficient_material
 
 def test_chess_inits():
     env = make("chess", debug=True)
@@ -21,4 +23,28 @@ def test_chess_100_move_rule():
     json = env.toJSON()
     assert json["name"] == "chess"
     assert json["statuses"] == ["DONE", "DONE"]
-    assert json["rewards"] == [0, 0]
+    assert json["rewards"] == [0, 0]
+
+def test_sufficient_material():
+    game = Game()
+    assert not is_insufficient_material(game.board)
+
+def test_insufficient_material_with_two_kings():
+    game = Game('8/8/K7/8/8/3k4/8/8 w - - 58 282')
+    assert is_insufficient_material(game.board)
+
+def test_insufficient_material_with_two_kings_and_bishop():
+    game = Game('6k1/8/7B/8/8/8/8/2K5 b - - 90 250')
+    assert is_insufficient_material(game.board)
+
+def test_insufficient_material_with_two_kings_and_two_knights():
+    game = Game('6k1/8/6NN/8/8/8/8/2K5 b - - 90 250')
+    assert is_insufficient_material(game.board)
+
+def test_sufficient_material_with_king_knight_and_bishop():
+    game = Game('6k1/8/6NB/8/8/8/8/2K5 b - - 90 250')
+    assert not is_insufficient_material(game.board)
+
+def test_sufficient_material_with_king_bishop_and_bishop():
+    game = Game('6k1/8/6BB/8/8/8/8/2K5 b - - 90 250')
+    assert not is_insufficient_material(game.board)