Improve test driver