Bug fix and add tests
test