Refactor code. Fix tests.