Fix tests and some cleanup.