Unit test added, controller logic modified