more unit tests & code refactor