added unit tests / fixed bug