More unit tests and some minor fixes