Unit Test modified