diff --git a/mock_payload_state.h b/mock_payload_state.h
index 2e8536a..4e21a0f 100644
--- a/mock_payload_state.h
+++ b/mock_payload_state.h
@@ -31,8 +31,9 @@
   MOCK_METHOD0(Rollback, void());
   MOCK_METHOD1(ExpectRebootInNewVersion,
                void(const std::string& target_version_uid));
-  MOCK_METHOD0(GetP2PNumAttempts, int());
-  MOCK_METHOD0(GetP2PFirstAttemptTimestamp, base::Time());
+  MOCK_METHOD0(P2PNewAttempt, void());
+  MOCK_METHOD0(P2PAttemptAllowed, bool());
+  MOCK_METHOD1(SetUsingP2PForDownloading, void(bool));
 
   // Getters.
   MOCK_METHOD0(GetResponseSignature, std::string());
@@ -49,8 +50,9 @@
   MOCK_METHOD1(GetTotalBytesDownloaded, uint64_t(DownloadSource source));
   MOCK_METHOD0(GetNumReboots, uint32_t());
   MOCK_METHOD0(GetRollbackVersion, std::string());
-  MOCK_METHOD0(P2PNewAttempt, void());
-  MOCK_METHOD0(P2PAttemptAllowed, bool());
+  MOCK_METHOD0(GetP2PNumAttempts, int());
+  MOCK_METHOD0(GetP2PFirstAttemptTimestamp, base::Time());
+  MOCK_METHOD0(GetUsingP2PForDownloading, bool());
 };
 
 }  // namespace chromeos_update_engine