Added some test code and fix a bug in unit test.