Add some test code