Add missing import

Should fix tests on Travis.

I don't quite understand why the tests were failing
nondeterministically, though.