More error handling and fixup

Signed-off-by: Mrunal Patel <mrunalp@gmail.com>