Improved test code