From 0f4cf413a3c527ddf2eab3d724ae758e39af1aa0 Mon Sep 17 00:00:00 2001
From: Trevor Tao <trevor.tao@arm.com>
Date: Thu, 21 Sep 2017 14:43:29 +0800
Subject: Add nodeSelector to enable selecting the desired Kubernetes running
 node when creating containers for Kubernetes context

For example, a yaml file may looks like:
 servers:
    host:
      image: xxx
      command: /bin/bash
      nodeSelector:
          xxx: yyy

Synchronously change the unit test for this function

Change-Id: If74c9dad9b1a70395bb79f34708a0fde04e7e650
Signed-off-by: Trevor Tao <trevor.tao@arm.com>
---
 tests/unit/orchestrator/test_kubernetes.py |  8 ++++++--
 yardstick/orchestrator/kubernetes.py       | 10 +++++++++-
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/tests/unit/orchestrator/test_kubernetes.py b/tests/unit/orchestrator/test_kubernetes.py
index 51718ab86..1a3291c89 100644
--- a/tests/unit/orchestrator/test_kubernetes.py
+++ b/tests/unit/orchestrator/test_kubernetes.py
@@ -62,7 +62,10 @@ service ssh restart;while true ; do sleep 10000; done"
                                 },
                                 "name": "k8s-86096c30-key"
                             }
-                        ]
+                        ],
+                        "nodeSelector": {
+                           "kubernetes.io/hostname": "node-01"
+                        }
                     }
                 }
             }
@@ -71,7 +74,8 @@ service ssh restart;while true ; do sleep 10000; done"
             'command': '/bin/bash',
             'args': ['-c', 'chmod 700 ~/.ssh; chmod 600 ~/.ssh/*; \
 service ssh restart;while true ; do sleep 10000; done'],
-            'ssh_key': 'k8s-86096c30-key'
+            'ssh_key': 'k8s-86096c30-key',
+            'nodeSelector': { 'kubernetes.io/hostname': 'node-01'}
         }
         name = 'host-k8s-86096c30'
         output_r = KubernetesObject(name, **input_s).get_template()
diff --git a/yardstick/orchestrator/kubernetes.py b/yardstick/orchestrator/kubernetes.py
index 6d7045f58..3d4548527 100644
--- a/yardstick/orchestrator/kubernetes.py
+++ b/yardstick/orchestrator/kubernetes.py
@@ -23,6 +23,7 @@ class KubernetesObject(object):
         self.command = [kwargs.get('command', '/bin/bash')]
         self.args = kwargs.get('args', [])
         self.ssh_key = kwargs.get('ssh_key', 'yardstick_key')
+        self.node_selector = kwargs.get('nodeSelector', {})
 
         self.volumes = []
 
@@ -42,7 +43,8 @@ class KubernetesObject(object):
                     },
                     "spec": {
                         "containers": [],
-                        "volumes": []
+                        "volumes": [],
+                        "nodeSelector": {}
                     }
                 }
             }
@@ -50,6 +52,7 @@ class KubernetesObject(object):
 
         self._change_value_according_name(name)
         self._add_containers()
+        self._add_node_selector()
         self._add_ssh_key_volume()
         self._add_volumes()
 
@@ -88,6 +91,11 @@ class KubernetesObject(object):
 
         return container
 
+    def _add_node_selector(self):
+        utils.set_dict_value(self.template,
+                             'spec.template.spec.nodeSelector',
+                             self.node_selector)
+
     def _add_volumes(self):
         utils.set_dict_value(self.template,
                              'spec.template.spec.volumes',
-- 
cgit 1.2.3-korg