More unit tests

Added more unit tests and fixed small issues.