Allow QA node acquisition and group creation using `with`

For acquiring nodes use `with AcquireManyNodes(num): ...`. The nodes
will be released automatically.
For creating a new group use `with NewGroupCtx() as group: ...`. The
group will be removed automatically.
Signed-off-by: Petr Pudlak <pudlak@google.com>
Reviewed-by: Hrvoje Ribicic <riba@google.com>

qa/qa_config.py

@@ -757,7 +757,25 @@ def AcquireNode(exclude=None, _cfg=None):
   return sorted(nodes, key=_NodeSortKey)[0].Use()
 
 
-def AcquireManyNodes(num, exclude=None):
+class AcquireManyNodesCtx(object):
+  """Returns the least used nodes for use with a `with` block
+
+  """
+  def __init__(self, num, exclude=None, cfg=None):
+    self._num = num
+    self._exclude = exclude
+    self._cfg = cfg
+
+  def __enter__(self):
+    self._nodes = AcquireManyNodes(self._num, exclude=self._exclude,
+                                   cfg=self._cfg)
+    return self._nodes
+
+  def __exit__(self, exc_type, exc_value, exc_tb):
+    ReleaseManyNodes(self._nodes)
+
+
+def AcquireManyNodes(num, exclude=None, cfg=None):
   """Return the least used nodes.
 
   @type num: int
@@ -779,7 +797,7 @@ def AcquireManyNodes(num, exclude=None):
 
   try:
     for _ in range(0, num):
-      n = AcquireNode(exclude=exclude)
+      n = AcquireNode(exclude=exclude, _cfg=cfg)
       nodes.append(n)
       exclude.append(n)
   except qa_error.OutOfNodesError:

qa/qa_group.py

@@ -127,6 +127,18 @@ def TestGroupAddWithOptions():
   AssertCommand(["gnt-group", "remove", group1])
 
 
+class NewGroupCtx(object):
+  """Creates a new group and disposes afterwards."""
+
+  def __enter__(self):
+    (self._group, ) = qa_utils.GetNonexistentGroups(1)
+    AssertCommand(["gnt-group", "add", self._group])
+    return self._group
+
+  def __exit__(self, exc_type, exc_val, exc_tb):
+    AssertCommand(["gnt-group", "remove", self._group])
+
+
 def _GetGroupIPolicy(groupname):
   """Return the run-time values of the cluster-level instance policy.