Update test code.
Fixed replace and restore methods.