diff --git a/mock_payload_state.h b/mock_payload_state.h
index 6dccc64..4cfcac0 100644
--- a/mock_payload_state.h
+++ b/mock_payload_state.h
@@ -68,6 +68,8 @@
   MOCK_METHOD1(GetCurrentBytesDownloaded, uint64_t(DownloadSource source));
   MOCK_METHOD1(GetTotalBytesDownloaded, uint64_t(DownloadSource source));
   MOCK_METHOD0(GetNumReboots, uint32_t());
+  MOCK_METHOD0(GetRollbackHappened, bool());
+  MOCK_METHOD1(SetRollbackHappened, void(bool));
   MOCK_METHOD0(GetRollbackVersion, std::string());
   MOCK_METHOD0(GetP2PNumAttempts, int());
   MOCK_METHOD0(GetP2PFirstAttemptTimestamp, base::Time());