Fix pytest imports

This fixes a few ``pytest`` imports - what I'm confused about is why this doesn't raise any errors otherwise in Travis if we aren't installing ``pytest`` separately?