Fix test failures

Merged in the changes from @viveksjain and fixed any test failures that came up.
maybe I missed it, but did you write a test for this?