Fix tests to correspond with the code

This should make Travis happy