Update the tests to use the new refactored code
Document an assumption in the test better