Fixing unit tests (Shame on me!)