Made reviewed changes and add some test code as well