Follow to D127355, converting more `assertEquals` to `assertState`. Differential Revision: https://reviews.llvm.org/D127378