add test and fix minor bugs