Code cleanup, unit tests, and minor bug fixes