diff --git a/mock_payload_state.h b/mock_payload_state.h
index fe744de..88578b5 100644
--- a/mock_payload_state.h
+++ b/mock_payload_state.h
@@ -21,6 +21,7 @@
   MOCK_METHOD1(SetResponse, void(const OmahaResponse& response));
   MOCK_METHOD0(DownloadComplete, void());
   MOCK_METHOD1(DownloadProgress, void(size_t count));
+  MOCK_METHOD0(UpdateResumed, void());
   MOCK_METHOD0(UpdateRestarted, void());
   MOCK_METHOD0(UpdateSucceeded, void());
   MOCK_METHOD1(UpdateFailed, void(ActionExitCode error));
@@ -37,6 +38,7 @@
   MOCK_METHOD0(GetUpdateDurationUptime, base::TimeDelta());
   MOCK_METHOD1(GetCurrentBytesDownloaded, uint64_t(DownloadSource source));
   MOCK_METHOD1(GetTotalBytesDownloaded, uint64_t(DownloadSource source));
+  MOCK_METHOD0(GetNumReboots, uint32_t());
 };
 
 }  // namespace chromeos_update_engine