Code review fixes
Fixes after testing