Fixed imports of pytest

Can be merged once tests pass