Fix some bugs. Adds tests