From ace2144bc18a3c88641e29ea3088fdca7f39f55c Mon Sep 17 00:00:00 2001
From: fjahja <fjahja@oath.com>
Date: Mon, 10 Sep 2018 14:29:37 -0700
Subject: [PATCH] feat(757): Pass disk annotation as toleration and node
 selector

---
 index.js           |  8 +++++-
 test/index.test.js | 64 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+), 1 deletion(-)

diff --git a/index.js b/index.js
index 0eefb32..12c6ac1 100644
--- a/index.js
+++ b/index.js
@@ -17,6 +17,7 @@ const MAXATTEMPTS = 5;
 const RETRYDELAY = 3000;
 const CPU_RESOURCE = 'beta.screwdriver.cd/cpu';
 const RAM_RESOURCE = 'beta.screwdriver.cd/ram';
+const DISK_RESOURCE = 'beta.screwdriver.cd/disk';
 const ANNOTATION_BUILD_TIMEOUT = 'beta.screwdriver.cd/timeout';
 const TOLERATIONS_PATH = 'spec.tolerations';
 const AFFINITY_NODE_SELECTOR_PATH = 'spec.affinity.nodeAffinity.' +
@@ -224,7 +225,12 @@ class K8sVMExecutor extends Executor {
 
         const podConfig = yaml.safeLoad(podTemplate);
 
-        setNodeSelector(podConfig, this.nodeSelectors);
+        const diskConfig = annotations[DISK_RESOURCE] || '';
+        const nodeSelectors = diskConfig.toUpperCase() === 'HIGH' ? { disk: 'high' } : {};
+
+        hoek.merge(nodeSelectors, this.nodeSelectors);
+
+        setNodeSelector(podConfig, nodeSelectors);
         setPreferredNodeSelector(podConfig, this.preferredNodeSelectors);
 
         const options = {
diff --git a/test/index.test.js b/test/index.test.js
index 0fa3299..b343550 100644
--- a/test/index.test.js
+++ b/test/index.test.js
@@ -71,6 +71,40 @@ describe('index', () => {
             }
         }
     };
+    const testSpecWithDisk = {
+        tolerations: [{
+            key: 'disk',
+            value: 'high',
+            effect: 'NoSchedule',
+            operator: 'Equal'
+        }, {
+            key: 'key',
+            value: 'value',
+            effect: 'NoSchedule',
+            operator: 'Equal'
+        }],
+        affinity: {
+            nodeAffinity: {
+                requiredDuringSchedulingIgnoredDuringExecution: {
+                    nodeSelectorTerms: [
+                        {
+                            matchExpressions: [{
+                                key: 'disk',
+                                operator: 'In',
+                                values: ['high']
+                            }]
+                        }, {
+                            matchExpressions: [{
+                                key: 'key',
+                                operator: 'In',
+                                values: ['value']
+                            }]
+                        }
+                    ]
+                }
+            }
+        }
+    };
     const testPreferredSpec = {
         affinity: {
             nodeAffinity: {
@@ -493,6 +527,36 @@ describe('index', () => {
             });
         });
 
+        it('sets disk toleration and node affinity when disk is HIGH', () => {
+            fakeStartConfig.annotations = { 'beta.screwdriver.cd/disk': 'HIGH' };
+            const spec = _.merge({}, testSpecWithDisk, testPodSpec);
+
+            postConfig.body.spec = spec;
+
+            executor = new Executor({
+                ecosystem: {
+                    api: testApiUri,
+                    store: testStoreUri
+                },
+                fusebox: { retry: { minTimeout: 1 } },
+                kubernetes: {
+                    nodeSelectors: { key: 'value' },
+                    token: 'api_key',
+                    host: 'kubernetes.default',
+                    baseImage: 'hyperctl'
+                },
+                prefix: 'beta_'
+            });
+
+            getConfig.retryStrategy = executor.podRetryStrategy;
+
+            return executor.start(fakeStartConfig).then(() => {
+                assert.calledWith(requestRetryMock.firstCall, postConfig);
+                assert.calledWith(requestRetryMock.secondCall,
+                    sinon.match(getConfig));
+            });
+        });
+
         it('sets preferred node affinity with appropriate node config', () => {
             const spec = _.merge({}, testPreferredSpec, testPodSpec);