Add tests and fix code (not completed)