Fixing some tests and removing unused classes