Added unit tests + some refactoring

Change-Id: Ic58f68ecef03153a33d6c3ffcf3ba4bcd7c261ba