[TASK] Added more unit tests