Revert "Revert "Update unit test""

This reverts commit 548f5eaff37892596643f192cc22aa715c8c4be4.