Few more unit tests and bug fixes