- [X] Fix unit tests