Bug fixes take3

Fix remaining unit tests