Fix tests and more cleanup.