More unit tests and minor bug fixes!