small changes to code in testing