Removing unused code
Fixing failing tests