Fix bugs and tests