Refactor and use generics for some classes
Enable tests