Fixed the Unit Tests