Update and cleanup tests