From 1fbe2198d17673ae83ccc4ff75432fe03d0b65dd Mon Sep 17 00:00:00 2001
From: Bilal Amarni <bilal.amarni@gmail.com>
Date: Thu, 13 Jun 2019 17:32:17 +0200
Subject: [PATCH] Add heartbeat helper method

This adds a helper method for sending heartbeat control messages with
the `ATTACH_CONTAINER_INPUT` call.
---
 api/v1/lib/agent/calls/calls.go      | 16 ++++++++++++++++
 api/v1/lib/agent/calls/calls_test.go |  1 +
 2 files changed, 17 insertions(+)

diff --git a/api/v1/lib/agent/calls/calls.go b/api/v1/lib/agent/calls/calls.go
index 50060de2..2c2830d0 100644
--- a/api/v1/lib/agent/calls/calls.go
+++ b/api/v1/lib/agent/calls/calls.go
@@ -220,6 +220,22 @@ func AttachContainerInputTTY(t *mesos.TTYInfo) *agent.Call {
 	}
 }
 
+func AttachContainerInputHeartbeat(heartbeat *agent.ProcessIO_Control_Heartbeat) *agent.Call {
+	return &agent.Call{
+		Type: agent.Call_ATTACH_CONTAINER_INPUT,
+		AttachContainerInput: &agent.Call_AttachContainerInput{
+			Type: agent.Call_AttachContainerInput_PROCESS_IO,
+			ProcessIO: &agent.ProcessIO{
+				Type: agent.ProcessIO_CONTROL,
+				Control: &agent.ProcessIO_Control{
+					Type:      agent.ProcessIO_Control_HEARTBEAT,
+					Heartbeat: heartbeat,
+				},
+			},
+		},
+	}
+}
+
 func AddResourceProviderConfig(rpi mesos.ResourceProviderInfo) *agent.Call {
 	return &agent.Call{
 		Type: agent.Call_ADD_RESOURCE_PROVIDER_CONFIG,
diff --git a/api/v1/lib/agent/calls/calls_test.go b/api/v1/lib/agent/calls/calls_test.go
index 1f70d42d..c181c2bb 100644
--- a/api/v1/lib/agent/calls/calls_test.go
+++ b/api/v1/lib/agent/calls/calls_test.go
@@ -75,6 +75,7 @@ func Example() {
 		AttachContainerInput(mesos.ContainerID{}),
 		AttachContainerInputTTY(nil),
 		AttachContainerInputData(nil),
+		AttachContainerInputHeartbeat(nil),
 	)))
 
 	// Output: