Addressed pull request comments.

Fixed broken tests.