wordle_env/test_wordle.py

import pytest

from . import wordle
from . import state

TESTWORDS = [
    "APPAA",
    "APPAB",
    "APPAC",
    "APAPD",

    "BPPAB",
    "BPPAC",
    "BPPAD",

    "CPPAB",
    "CPPAC",
    "CPPAD",
]


@pytest.fixture
def wordleEnv():
    env = wordle.WordleEnvBase(
        words=TESTWORDS,
        max_turns=6,
    )
    return env


def test_reset(wordleEnv):
    wordleEnv.reset(seed=13)


def test_guess_win(wordleEnv):
    wordleEnv.reset(seed=13)
    goal = wordleEnv.goal_word
    new_state, reward, done, _ = wordleEnv.step(goal)
    assert done
    assert wordleEnv.done
    assert reward == 0

    try:
        wordleEnv.step(goal)
        raise ValueError("Shouldn't reach here!")
    except ValueError:
        pass


def test_win_reward(wordleEnv):
    wordleEnv.reset(seed=13)
    goal = wordleEnv.goal_word
    word_index = (goal + 1) % len(wordleEnv.words)
    new_state, reward, done, _ = wordleEnv.step(word_index)
    assert state.remaining_steps(new_state) == wordleEnv.max_turns-1
    assert not done
    assert not wordleEnv.done
    assert reward < 1

    new_state, reward, done, _ = wordleEnv.step(goal)
    assert state.remaining_steps(new_state) == wordleEnv.max_turns-2
    assert done
    assert wordleEnv.done
    assert reward == wordle.REWARD

    try:
        wordleEnv.step(goal)
        raise ValueError("Shouldn't reach here!")
    except ValueError:
        pass


def test_win_reward_6(wordleEnv):
    wordleEnv.reset(seed=13)
    goal = wordleEnv.goal_word
    random_word = (goal + 1) % len(wordleEnv.words)

    for i in range(5):
        new_state, reward, done, _ = wordleEnv.step(random_word)

    new_state, reward, done, _ = wordleEnv.step(goal)

    assert wordleEnv.max_turns - state.remaining_steps(new_state) == 6
    assert done
    assert wordleEnv.done
    assert reward == wordle.REWARD


def test_lose_reward(wordleEnv):
    wordleEnv.reset(seed=13)
    goal = wordleEnv.goal_word
    random_word = (goal + 1) % len(wordleEnv.words)
    for i in range(1, wordleEnv.max_turns):
        new_state, reward, done, _ = wordleEnv.step(random_word)
        assert state.remaining_steps(new_state) == wordleEnv.max_turns-i
        assert not done
        assert not wordleEnv.done
        assert reward < 1

    word_index = (goal + wordleEnv.max_turns) % len(wordleEnv.words)
    random_word = word_index
    new_state, reward, done, _ = wordleEnv.step(random_word)
    assert state.remaining_steps(new_state) == 0
    assert done
    assert wordleEnv.done
    assert reward < -wordle.REWARD + 1

    try:
        wordleEnv.step(goal)
        raise ValueError("Shouldn't reach here!")
    except ValueError:
        pass

def letter_test(char, state, letter_state):
    offset = 1+3*5*(ord(char)-ord('A'))
    assert tuple(state[offset:offset+15]) == letter_state

def test_step(wordleEnv):
    wordleEnv.reset(seed=13)
    wordleEnv.set_goal_encoded(0)

    cur_state = wordleEnv.state
    new_state, reward, done, _ = wordleEnv.step(1)
    assert state.remaining_steps(cur_state) == wordleEnv.max_turns
    assert state.remaining_steps(new_state) == wordleEnv.max_turns-1
    # 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)
    cur_state = wordleEnv.state
    new_state, reward, done, _ = wordleEnv.step(3)
    assert state.remaining_steps(cur_state) == wordleEnv.max_turns-1
    assert state.remaining_steps(new_state) == wordleEnv.max_turns-2
    # 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)
    print(new_state, letter_state)
    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)

    new_state, reward, done, _ = wordleEnv.step(2)
    assert state.remaining_steps(new_state) == wordleEnv.max_turns-3
    # Expect B to be all 1,0,0
    letter_state = tuple([1, 0, 0]*5)
    letter_test('C', new_state, letter_state)

    # Expect C to be all 1,0,0
    letter_state = tuple([1, 0, 0]*5)
    letter_test('C', new_state, letter_state)

    # Expect A to be right in position 0 3 and, no 3 and maybe otherwise
    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 0 3 and maybe otherwise
    letter_state = (1, 0, 0,
                    0, 0, 1,
                    0, 0, 1,
                    1, 0, 0,
                    0, 1, 0)
    letter_test('P', new_state, letter_state)

    new_state, reward, done, _ = wordleEnv.step(0)
    # Expect A to be right in position 0 3 and 4
    letter_state = (0, 0, 1,
                    1, 0, 0,
                    1, 0, 0,
                    0, 0, 1,
                    0, 0, 1)
    letter_test('A', new_state, letter_state)

    # Expect P to be right in position 1 2 and not otherwise
    letter_state = (1, 0, 0,
                    0, 0, 1,
                    0, 0, 1,
                    1, 0, 0,
                    1, 0, 0)
    letter_test('P', new_state, letter_state)
    assert state.remaining_steps(new_state) == wordleEnv.max_turns-4
    assert done
    assert wordleEnv.done
    assert reward == wordle.REWARD