REFACTOR Fix unit test
Fix unit tests