Code review Fixes