Improved unit test