Bug fix and add some tests.