Refactoring and adding unit tests.

[#49450933 #49448805]