Modified unit test