Add test coverage and fix minor bugs (DO NOT MERGE)