Some changes on unit tests

Minor correction on unit test