Fixed imports and reinstantiated old unit test to make sure we actually test both scenarios. Original pull request: #497.