Fix bugs and increase test coverage