Fix test cases and minor cleanups