Сделал не большой рефакторинг
Немного дописал тесты.