Fix some unit tests