fixing the code review demo
minor refactor