From ec8b2da4af337248ffc05399fbf39d2ba6aa51e5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Christian=20Hu=CC=88ning?=
 <christian.huening@haw-hamburg.de>
Date: Tue, 9 Jan 2018 16:31:44 +0100
Subject: [PATCH] added agentDaemonsetTolerationKeyEnv Var to be able to set a
 Toleration Key

added new option to example rook-operator.yaml

updated tests for the AGENT_TOLERATION_KEY env var
---
 cluster/examples/kubernetes/rook-operator.yaml |  6 +++++-
 pkg/operator/agent/agent.go                    | 10 ++++++----
 pkg/operator/agent/agent_test.go               |  4 ++++
 3 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/cluster/examples/kubernetes/rook-operator.yaml b/cluster/examples/kubernetes/rook-operator.yaml
index 021f1e73c..d30267ec1 100644
--- a/cluster/examples/kubernetes/rook-operator.yaml
+++ b/cluster/examples/kubernetes/rook-operator.yaml
@@ -124,9 +124,13 @@ spec:
         # To disable RBAC, uncomment the following:
         # - name: RBAC_ENABLED
         #  value: "false"
-        # Rook Agent toleration. Choose between NoSchedule, PreferNoSchedule and NoExecute:
+        # Rook Agent toleration. Will tolerate all taints with all keys.
+        # Choose between NoSchedule, PreferNoSchedule and NoExecute:
         # - name: AGENT_TOLERATION
         #  value: "NoSchedule"
+        # (Optional) Rook Agent toleration key. Set this to the key of the taint you want to tolerate
+        # - name: AGENT_TOLERATION_KEY
+        #  value: "<KeyOfTheTaintToTolerate"
         # The interval to check if every mon is in the quorum.
         - name: ROOK_MON_HEALTHCHECK_INTERVAL
           value: "45s"
diff --git a/pkg/operator/agent/agent.go b/pkg/operator/agent/agent.go
index 9a6f4f902..4f550f1b2 100644
--- a/pkg/operator/agent/agent.go
+++ b/pkg/operator/agent/agent.go
@@ -36,10 +36,11 @@ import (
 )
 
 const (
-	agentDaemonsetName          = "rook-agent"
-	flexvolumePathDirEnv        = "FLEXVOLUME_DIR_PATH"
-	flexvolumeDefaultDirPath    = "/usr/libexec/kubernetes/kubelet-plugins/volume/exec/"
-	agentDaemonsetTolerationEnv = "AGENT_TOLERATION"
+	agentDaemonsetName             = "rook-agent"
+	flexvolumePathDirEnv           = "FLEXVOLUME_DIR_PATH"
+	flexvolumeDefaultDirPath       = "/usr/libexec/kubernetes/kubelet-plugins/volume/exec/"
+	agentDaemonsetTolerationEnv    = "AGENT_TOLERATION"
+	agentDaemonsetTolerationKeyEnv = "AGENT_TOLERATION_KEY"
 )
 
 var logger = capnslog.NewPackageLogger("github.com/rook/rook", "op-agent")
@@ -187,6 +188,7 @@ func (a *Agent) createAgentDaemonSet(namespace, agentImage string) error {
 			{
 				Effect:   v1.TaintEffect(tolerationValue),
 				Operator: v1.TolerationOpExists,
+				Key:      os.Getenv(agentDaemonsetTolerationKeyEnv),
 			},
 		}
 	}
diff --git a/pkg/operator/agent/agent_test.go b/pkg/operator/agent/agent_test.go
index fb2183268..5a81df50f 100644
--- a/pkg/operator/agent/agent_test.go
+++ b/pkg/operator/agent/agent_test.go
@@ -169,6 +169,9 @@ func TestStartAgentDaemonsetWithToleration(t *testing.T) {
 	os.Setenv(agentDaemonsetTolerationEnv, "NoSchedule")
 	defer os.Unsetenv(agentDaemonsetTolerationEnv)
 
+	os.Setenv(agentDaemonsetTolerationKeyEnv, "example")
+	defer os.Unsetenv(agentDaemonsetTolerationKeyEnv)
+
 	pod := v1.Pod{
 		ObjectMeta: metav1.ObjectMeta{
 			Name:      "rook-operator",
@@ -197,5 +200,6 @@ func TestStartAgentDaemonsetWithToleration(t *testing.T) {
 	assert.Nil(t, err)
 	assert.Equal(t, 1, len(agentDS.Spec.Template.Spec.Tolerations))
 	assert.Equal(t, "NoSchedule", string(agentDS.Spec.Template.Spec.Tolerations[0].Effect))
+	assert.Equal(t, "example", string(agentDS.Spec.Template.Spec.Tolerations[0].Key))
 	assert.Equal(t, "Exists", string(agentDS.Spec.Template.Spec.Tolerations[0].Operator))
 }