adb-remount-test: Refactor restore()

The redefining of cleanup hooks are making the script rather difficult
to read. Instead of redefining restore() just to skip some parts of it,
let restore() check flags and conditionally execute cleanup code.

Bug: 243116800
Test: adb-remount-test
Change-Id: If9d627618b54e215200455e8133492670737571d