Add test functionality