refactored tests and added feature
tidy up