Update unit test

Not familiar with unit test, tell me if it's good