updates based on richards comments
code review changes
Merge pull request #11 from TK-23/refactor_tests

Refactor tests