Add more test code