Fix broken import.

Seems like IntelliJ messes with my python path so I didn't
see this issue when executing the tests from it.