Unit tests and some refactor to have constructor DI