Fixing broken tests.
cc @hrushikesh, all tests are passing right now