Fix documentation
Refactored and added tests