Restructuring and fix minor bugs in tests