More unit tests on unity module
Code cleanup