[JUnit] Minor functionality bug fixes and refactoring.
[JUnit] Added one more Unit test.