Fix test failures

Merged in the changes from @viveksjain and fixed any test failures that came up.
Please use one style for this. So the } should be on the same line and the code split over multiple lines.