From 89b6c6eea67bc065ce2e35cafc8c8edf0ee4eb7c Mon Sep 17 00:00:00 2001
From: madcpf <madcpf@gmail.com>
Date: Fri, 6 Oct 2023 14:59:52 -0700
Subject: [PATCH] add more checks

---
 unitary/examples/quantum_chinese_chess/chess.py     |  1 +
 unitary/examples/quantum_chinese_chess/enums.py     |  6 +++---
 unitary/examples/quantum_chinese_chess/move.py      |  8 +++++++-
 unitary/examples/quantum_chinese_chess/move_test.py | 12 ++++++++++--
 4 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/unitary/examples/quantum_chinese_chess/chess.py b/unitary/examples/quantum_chinese_chess/chess.py
index eb924d3e..c4dc3499 100644
--- a/unitary/examples/quantum_chinese_chess/chess.py
+++ b/unitary/examples/quantum_chinese_chess/chess.py
@@ -94,6 +94,7 @@ def play(self) -> None:
             # If the game continues, switch the player.
             if game_over == -1:
                 self.current_player = 1 - self.current_player
+                self.board.current_player = self.current_player
                 continue
             elif game_over == 0:
                 print(f"{self.players_name[0]} wins! Game is over.")
diff --git a/unitary/examples/quantum_chinese_chess/enums.py b/unitary/examples/quantum_chinese_chess/enums.py
index 128f1843..0487dbc1 100644
--- a/unitary/examples/quantum_chinese_chess/enums.py
+++ b/unitary/examples/quantum_chinese_chess/enums.py
@@ -48,9 +48,9 @@ class MoveVariant(enum.Enum):
 
 
 class Color(enum.Enum):
-    NA = 0
-    RED = 1
-    BLACK = 2
+    NA = -1
+    RED = 0
+    BLACK = 1
 
 
 class Type(enum.Enum):
diff --git a/unitary/examples/quantum_chinese_chess/move.py b/unitary/examples/quantum_chinese_chess/move.py
index 4181ad12..8b1ad967 100644
--- a/unitary/examples/quantum_chinese_chess/move.py
+++ b/unitary/examples/quantum_chinese_chess/move.py
@@ -14,7 +14,7 @@
 from typing import Optional, List, Tuple
 from unitary.alpha.quantum_effect import QuantumEffect
 from unitary.examples.quantum_chinese_chess.board import Board
-from unitary.examples.quantum_chinese_chess.enums import MoveType, MoveVariant
+from unitary.examples.quantum_chinese_chess.enums import MoveType, MoveVariant, Type
 
 
 def parse_input_string(str_to_parse: str) -> Tuple[List[str], List[str]]:
@@ -72,6 +72,12 @@ def get_move_from_string(str_to_parse: str, board: Board) -> "Move":
         sources, targets = parse_input_string(str_to_parse)
     except ValueError as e:
         raise e
+    # Additional checks based on the current board.
+    for source in sources:
+        if board.board[source].type_ == Type.EMPTY:
+            raise ValueError("Could not move empty piece.")
+        if board.board[source].color.value != board.current_player:
+            raise ValueError("Could not move the other player's piece.")
     # TODO(): add analysis to determine move type and variant.
     move_type = MoveType.UNSPECIFIED_STANDARD
     move_variant = MoveVariant.UNSPECIFIED
diff --git a/unitary/examples/quantum_chinese_chess/move_test.py b/unitary/examples/quantum_chinese_chess/move_test.py
index 9bfffa60..9512a33b 100644
--- a/unitary/examples/quantum_chinese_chess/move_test.py
+++ b/unitary/examples/quantum_chinese_chess/move_test.py
@@ -21,13 +21,13 @@
 import pytest
 
 
-def test_parse_success():
+def test_parse_input_string_success():
     assert parse_input_string("a1b1") == (["a1"], ["b1"])
     assert parse_input_string("a1b1^c2") == (["a1", "b1"], ["c2"])
     assert parse_input_string("a1^b1c2") == (["a1"], ["b1", "c2"])
 
 
-def test_parse_fail():
+def test_parse_input_string_fail():
     with pytest.raises(ValueError, match="Invalid sources/targets string "):
         parse_input_string("a1^b1")
     with pytest.raises(ValueError, match="Invalid sources/targets string "):
@@ -44,6 +44,14 @@ def test_parse_fail():
         parse_input_string("a1n1")
 
 
+def test_get_move_from_string_fail():
+    board = Board.from_fen()
+    with pytest.raises(ValueError, match="Could not move empty piece."):
+        get_move_from_string("a1b1", board)
+    with pytest.raises(ValueError, match="Could not move the other player's piece."):
+        get_move_from_string("a0b1", board)
+
+
 def test_move_eq():
     board = Board.from_fen()
     move1 = Move(