Modified tests to match newest implementation