Revert "Update unit test"

This reverts commit 21e427599066728f80125674b2873f6b69dddbcd.