Fix UT and change function to private