Updates tests to reflect the new code. One test still breaks