Improve unit test