Add constants for node roles

Signed-off-by: Michael Hanselmann <hansmi@google.com>
Reviewed-by: Iustin Pop <iustin@google.com>

diff --git a/lib/constants.py b/lib/constants.py
index 117b43c..63c4880 100644 (file)
--- a/lib/constants.py
+++ b/lib/constants.py
@@ -879,6 +879,20 @@ INSTST_ALL = frozenset([
   INSTST_ERRORDOWN,
   ])
 
+# Node roles
+NR_REGULAR = "R"
+NR_MASTER = "M"
+NR_MCANDIDATE = "C"
+NR_DRAINED = "D"
+NR_OFFLINE = "O"
+NR_ALL = frozenset([
+  NR_REGULAR,
+  NR_MASTER,
+  NR_MCANDIDATE,
+  NR_DRAINED,
+  NR_OFFLINE,
+  ])
+
 # SSL certificate check constants (in days)
 SSL_CERT_EXPIRATION_WARN = 30
 SSL_CERT_EXPIRATION_ERROR = 7

diff --git a/lib/query.py b/lib/query.py
index b185170..47a4d5c 100644 (file)
--- a/lib/query.py
+++ b/lib/query.py
@@ -381,15 +381,15 @@ def _GetNodeRole(node, master_name):
 
   """
   if node.name == master_name:
-    return "M"
+    return constants.NR_MASTER
   elif node.master_candidate:
-    return "C"
+    return constants.NR_MCANDIDATE
   elif node.drained:
-    return "D"
+    return constants.NR_DRAINED
   elif node.offline:
-    return "O"
+    return constants.NR_OFFLINE
   else:
-    return "R"
+    return constants.NR_REGULAR
 
 
 def _GetItemAttr(attr):

diff --git a/lib/rapi/rlib2.py b/lib/rapi/rlib2.py
index 6504ab2..70d8215 100644 (file)
--- a/lib/rapi/rlib2.py
+++ b/lib/rapi/rlib2.py
@@ -86,13 +86,15 @@ _NR_OFFLINE = "offline"
 _NR_REGULAR = "regular"
 
 _NR_MAP = {
-  "M": _NR_MASTER,
-  "C": _NR_MASTER_CANDIATE,
-  "D": _NR_DRAINED,
-  "O": _NR_OFFLINE,
-  "R": _NR_REGULAR,
+  constants.NR_MASTER: _NR_MASTER,
+  constants.NR_MCANDIDATE: _NR_MASTER_CANDIATE,
+  constants.NR_DRAINED: _NR_DRAINED,
+  constants.NR_OFFLINE: _NR_OFFLINE,
+  constants.NR_REGULAR: _NR_REGULAR,
   }
 
+assert frozenset(_NR_MAP.keys()) == constants.NR_ALL
+
 # Request data version field
 _REQ_DATA_VERSION = "__version__"
 

diff --git a/test/ganeti.query_unittest.py b/test/ganeti.query_unittest.py
index 8994425..8039435 100755 (executable)
--- a/test/ganeti.query_unittest.py
+++ b/test/ganeti.query_unittest.py
@@ -275,25 +275,26 @@ class TestQuery(unittest.TestCase):
 
 
 class TestGetNodeRole(unittest.TestCase):
-  def testMaster(self):
-    node = objects.Node(name="node1")
-    self.assertEqual(query._GetNodeRole(node, "node1"), "M")
-
-  def testMasterCandidate(self):
-    node = objects.Node(name="node1", master_candidate=True)
-    self.assertEqual(query._GetNodeRole(node, "master"), "C")
-
-  def testRegular(self):
-    node = objects.Node(name="node1")
-    self.assertEqual(query._GetNodeRole(node, "master"), "R")
+  def test(self):
+    tested_role = set()
+
+    checks = [
+      (constants.NR_MASTER, "node1", objects.Node(name="node1")),
+      (constants.NR_MCANDIDATE, "master",
+       objects.Node(name="node1", master_candidate=True)),
+      (constants.NR_REGULAR, "master", objects.Node(name="node1")),
+      (constants.NR_DRAINED, "master",
+       objects.Node(name="node1", drained=True)),
+      (constants.NR_OFFLINE,
+       "master", objects.Node(name="node1", offline=True)),
+      ]
 
-  def testDrained(self):
-    node = objects.Node(name="node1", drained=True)
-    self.assertEqual(query._GetNodeRole(node, "master"), "D")
+    for (role, master_name, node) in checks:
+      result = query._GetNodeRole(node, master_name)
+      self.assertEqual(result, role)
+      tested_role.add(result)
 
-  def testOffline(self):
-    node = objects.Node(name="node1", offline=True)
-    self.assertEqual(query._GetNodeRole(node, "master"), "O")
+    self.assertEqual(tested_role, constants.NR_ALL)
 
 
 class TestNodeQuery(unittest.TestCase):