Fix test failures

Merged in the changes from @viveksjain and fixed any test failures that came up.
No good reason. I probably just did it because this was simpler.