Add more tests. Refactor code.