Doc fix
Fixing unit tests