add tests

quantumlib · Oct 11, 2023 · 2dcb890 · 2dcb890
1 parent ee07708
commit 2dcb890
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 20 deletions.
diff --git a/unitary/examples/quantum_chinese_chess/chess.py b/unitary/examples/quantum_chinese_chess/chess.py
@@ -265,10 +265,10 @@ def classify_move(
                 raise ValueError(
                     "Both sources need to be in quantum state in order to merge."
                 )
+            # TODO(): Currently we don't support merge + excluded/capture, or cannon_merge_fire + capture. Maybe add support later.
+            if len(classical_path_pieces_0) > 0 or len(classical_path_pieces_1) > 0:
+                raise ValueError("Currently CANNON could not merge while fire.")
             if target.type_ != Type.EMPTY:
-                # TODO(): Currently we don't support merge + excluded/capture, or cannon_merge_fire + capture. Maybe add support later.
-                if len(classical_path_pieces_0) > 0 or len(classical_path_pieces_1) > 0:
-                    raise ValueError("Currently CANNON could not merge while fire.")
                 raise ValueError("Currently we could only merge into an empty piece.")
             if len(quantum_path_pieces_0) == 0 and len(quantum_path_pieces_1) == 0:
                 move_type = MoveType.MERGE_JUMP
@@ -278,10 +278,10 @@ def classify_move(
 
         elif len(targets) == 2:
             target_1 = self.board.board[targets[1]]
+            # TODO(): Currently we don't support split + excluded/capture, or cannon_split_fire + capture. Maybee add support later.
+            if len(classical_path_pieces_0) > 0 or len(classical_path_pieces_1) > 0:
+                raise ValueError("Currently CANNON could not split while fire.")
             if target.type_ != Type.EMPTY or target_1.type_ != Type.EMPTY:
-                # TODO(): Currently we don't support split + excluded/capture, or cannon_split_fire + capture. Maybee add support later.
-                if len(classical_path_pieces_0) > 0 or len(classical_path_pieces_1) > 0:
-                    raise ValueError("Currently CANNON could not split while fire.")
                 raise ValueError("Currently we could only split into empty pieces.")
             if source.type_ == Type.KING:
                 # TODO(): Currently we don't support KING split. Maybe add support later.
@@ -314,6 +314,7 @@ def apply_move(self, str_to_parse: str) -> None:
                 raise ValueError("Two sources need to be the same type.")
         if len(targets) == 2:
             target_1 = self.board.board[targets[1]]
+            # TODO(): handle the case where a piece is split into the current piece and another piece, in which case two targets are different.
             if target_0.type_ != target_1.type_:
                 raise ValueError("Two targets need to be the same type.")
             if target_0.color != target_1.color:
@@ -362,21 +363,18 @@ def apply_move(self, str_to_parse: str) -> None:
             raise e
 
         if move_type == MoveType.CLASSICAL:
-            if move_variant == MoveVariant.BASIC or move_variant == MoveVariant.CAPTURE:
-                if source_0.type_ == Type.KING:
-                    # Update the locations of KING.
-                    self.board.king_locations[self.current_player] = targets[0]
-                if target_0.type_ == Type.KING:
-                    # King is captured, then the game is over.
-                    self.game_state = GameState(self.current_player)
-                target_0.reset(source_0)
-                source_0.reset()
+            if source_0.type_ == Type.KING:
+                # Update the locations of KING.
+                self.board.king_locations[self.current_player] = targets[0]
                 # TODO(): only make such prints for a certain debug level.
-                print("Classical move.")
-            else:
-                raise ValueError(
-                    "Unspecifed/excluded classical move. This should never happen."
-                )
+                print(f"Updated king locations: {self.board.king_locations}.")
+            if target_0.type_ == Type.KING:
+                # King is captured, then the game is over.
+                self.game_state = GameState(self.current_player)
+            target_0.reset(source_0)
+            source_0.reset()
+            # TODO(): only make such prints for a certain debug level.
+            print("Classical move.")
 
     def next_move(self) -> bool:
         """Check if the player wants to exit or needs help message. Otherwise parse and apply the move.

diff --git a/unitary/examples/quantum_chinese_chess/chess_test.py b/unitary/examples/quantum_chinese_chess/chess_test.py
@@ -76,6 +76,12 @@ def test_apply_move_fail(monkeypatch):
         game.apply_move("a1b1")
     with pytest.raises(ValueError, match="Could not move the other player's piece."):
         game.apply_move("a0b1")
+    with pytest.raises(ValueError, match="Two sources need to be the same type."):
+        game.apply_move("a9a6^a5")
+    with pytest.raises(ValueError, match="Two targets need to be the same type."):
+        game.apply_move("b7^a7h7")
+    with pytest.raises(ValueError, match="Two targets need to be the same color."):
+        game.apply_move("b7^b2h7")
 
 
 def test_game_invalid_move(monkeypatch):