Remove tests.py
Refactor imports