Improve test code