Revert "Revert "Revert "Update unit test"""

This reverts commit 7a64c2c3f0fd99699b682423674bcc75ebbb26e5.