Saved a couple lines and updated the unit test