Fix tests and another bug or two