added unit tests

refactored and added more tests