Fixed bug in unittest.
Refactored.