improved unit test