diff --git a/cloud/agent/e2e/main_test.go b/cloud/agent/e2e/main_test.go
index 84fd611..bcc38f3 100644
--- a/cloud/agent/e2e/main_test.go
+++ b/cloud/agent/e2e/main_test.go
@@ -67,8 +67,8 @@
 	installationReport  *bpb.OSInstallationReport
 }
 
-func (f *fakeServer) Heartbeat(ctx context.Context, req *bpb.AgentHeartbeatRequest) (*bpb.AgentHeartbeatResponse, error) {
-	var res bpb.AgentHeartbeatResponse
+func (f *fakeServer) Heartbeat(ctx context.Context, req *bpb.HeartbeatRequest) (*bpb.HeartbeatResponse, error) {
+	var res bpb.HeartbeatResponse
 	if req.HardwareReport != nil {
 		f.hardwareReport = req.HardwareReport
 	}