Fix bugs + add more test cases