Add test for mask update

wordle_env/test_wordle.py

@@ -290,3 +290,134 @@ def test_special_step_cases(wordleEnv):
                     1, 0, 0,
                     1, 0, 0)
     letter_test('P', new_state, letter_state)
+
+def test_mask_update(wordleEnv):
+    wordleEnv.reset()
+    wordleEnv.set_goal_encoded(0)
+
+    cur_state = wordleEnv.state
+    #"APPAA"
+    #"APPAB"
+    new_state = state.update_from_mask(cur_state, wordleEnv.words[1], [2, 2, 2, 2, 0])
+    # Expect B to be all 1,0,0
+    letter_test('B', new_state, tuple([1, 0, 0]*5))
+
+    # Expect A to be right in position 0 3, no in 1 2 and 0 otherwise
+    letter_state = (0, 0, 1,
+                    1, 0, 0,
+                    1, 0, 0,
+                    0, 0, 1,
+                    0, 0, 0)
+    letter_test('A', new_state, letter_state)
+
+    # Expect P to be right in position 1 2, no in 0 2 and 0 otherwise
+    letter_state = (1, 0, 0,
+                    0, 0, 1,
+                    0, 0, 1,
+                    1, 0, 0,
+                    0, 0, 0)
+    letter_test('P', new_state, letter_state)
+
+    # Expect C to be no everywhere except in 4
+    letter_state = (1, 0, 0,
+                    1, 0, 0,
+                    1, 0, 0,
+                    1, 0, 0,
+                    0, 0, 0)
+    letter_test('C', new_state, letter_state)
+    # "APPAA",
+    # "APPAB",
+    # "APAPD",
+    new_state = state.update_from_mask(new_state, wordleEnv.words[3], [2, 2, 1, 1, 0])
+    # Expect D to be all 1,0,0
+    letter_state = tuple([1, 0, 0]*5)
+    letter_test('D', new_state, letter_state)
+
+    # Expect A to be right in position 0 3 and Maybe in 4
+    letter_state = (0, 0, 1,
+                    1, 0, 0,
+                    1, 0, 0,
+                    0, 0, 1,
+                    0, 1, 0)
+    letter_test('A', new_state, letter_state)
+
+    # Expect P to be right in position 1 2 no in 3 and maybe in 4
+    letter_state = (1, 0, 0,
+                    0, 0, 1,
+                    0, 0, 1,
+                    1, 0, 0,
+                    0, 1, 0)
+    letter_test('P', new_state, letter_state)
+
+    wordleEnv.reset()
+    cur_state = wordleEnv.state
+    wordleEnv.set_goal_encoded(4)
+    # BPPAB - goal
+    # PPAPB - 1st guess
+    new_state = state.update_from_mask(cur_state, wordleEnv.words[5], [1, 2, 1, 0, 2])
+    # Expect A to be all maybe except 2, 1 and 4 that are no
+    letter_state = (0, 1, 0,
+                    1, 0, 0,
+                    1, 0, 0,
+                    0, 1, 0,
+                    1, 0, 0)
+    letter_test('A', new_state, letter_state)
+    # Expect B to be all 0 except 4 that is yes and 1 that is no
+    letter_state = (0, 0, 0,
+                    1, 0, 0,
+                    0, 0, 0,
+                    0, 0, 0,
+                    0, 0, 1)
+    letter_test('B', new_state, letter_state)
+    # Expect P to be yes at 1, maybe at 2 and 4, and no in 0, 3 and 4
+    letter_state = (1, 0, 0,
+                    0, 0, 1,
+                    0, 1, 0,
+                    1, 0, 0,
+                    1, 0, 0)
+    letter_test('P', new_state, letter_state)
+    # BPPAB - goal
+    # PPAPB - 1st guess
+    # PPBBA - 2nd guess
+    new_state = state.update_from_mask(new_state, wordleEnv.words[6], [1, 2, 1, 1, 1])
+    # Expect A to be all maybe except 2, 1 and 4 that are no
+    letter_state = (0, 1, 0,
+                    1, 0, 0,
+                    1, 0, 0,
+                    0, 1, 0,
+                    1, 0, 0)
+    letter_test('A', new_state, letter_state)
+    # Expect B to be maybe at 0, yes at 4, and 1 2 and 3 are no
+    letter_state = (0, 1, 0,
+                    1, 0, 0,
+                    1, 0, 0,
+                    1, 0, 0,
+                    0, 0, 1)
+    letter_test('B', new_state, letter_state)
+    # Expect P to be yes at 1, maybe at 2, and no in 0, 3 and 4
+    letter_state = (1, 0, 0,
+                    0, 0, 1,
+                    0, 1, 0,
+                    1, 0, 0,
+                    1, 0, 0)
+    letter_test('P', new_state, letter_state)
+    wordleEnv.reset()
+    wordleEnv.set_goal_encoded(7)
+    # BPABB - goal
+    # PPPAC - 1st guess
+    new_state = state.update_from_mask(new_state, wordleEnv.words[8], [0, 2, 0, 1, 0])
+    new_state, _, _, _ = wordleEnv.step(8)
+    # Expect A to be all maybe except 1 and 3 that is no
+    letter_state = (0, 1, 0,
+                    1, 0, 0,
+                    0, 1, 0,
+                    1, 0, 0,
+                    0, 1, 0)
+    letter_test('A', new_state, letter_state)
+    # Expect P to be yes at 1, maybe at 3 and 4 and no otherwise
+    letter_state = (1, 0, 0,
+                    0, 0, 1,
+                    1, 0, 0,
+                    1, 0, 0,
+                    1, 0, 0)
+    letter_test('P', new_state, letter_state)