Refactor tests
Merge pull request #2 from ustasb/refactor_and_cleanup

Refactor tests