Fix & Unit tests