Move test code to correct package