diff --git a/mock_payload_state.h b/mock_payload_state.h
index 8959ca1..02db28b 100644
--- a/mock_payload_state.h
+++ b/mock_payload_state.h
@@ -31,7 +31,8 @@
 
   // Getters.
   MOCK_METHOD0(GetResponseSignature, std::string());
-  MOCK_METHOD0(GetPayloadAttemptNumber, uint32_t());
+  MOCK_METHOD0(GetPayloadAttemptNumber, int());
+  MOCK_METHOD0(GetFullPayloadAttemptNumber, int());
   MOCK_METHOD0(GetCurrentUrl, std::string());
   MOCK_METHOD0(GetUrlFailureCount, uint32_t());
   MOCK_METHOD0(GetUrlSwitchCount, uint32_t());