Fixing all tests
Removed unnecessary import