code.grnet.gr Git - ganeti-local/blob - qa/qa_utils.py

   1 #
   2 #
   3
   4 # Copyright (C) 2007 Google Inc.
   5 #
   6 # This program is free software; you can redistribute it and/or modify
   7 # it under the terms of the GNU General Public License as published by
   8 # the Free Software Foundation; either version 2 of the License, or
   9 # (at your option) any later version.
  10 #
  11 # This program is distributed in the hope that it will be useful, but
  12 # WITHOUT ANY WARRANTY; without even the implied warranty of
  13 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  14 # General Public License for more details.
  15 #
  16 # You should have received a copy of the GNU General Public License
  17 # along with this program; if not, write to the Free Software
  18 # Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
  19 # 02110-1301, USA.
  20
  21
  22 """Utilities for QA tests.
  23
  24 """
  25
  26 import os
  27 import re
  28 import sys
  29 import subprocess
  30
  31 from ganeti import utils
  32
  33 import qa_config
  34 import qa_error
  35
  36
  37 _INFO_SEQ = None
  38 _WARNING_SEQ = None
  39 _ERROR_SEQ = None
  40 _RESET_SEQ = None
  41
  42
  43 def _SetupColours():
  44   """Initializes the colour constants.
  45
  46   """
  47   global _INFO_SEQ, _WARNING_SEQ, _ERROR_SEQ, _RESET_SEQ
  48
  49   # Don't use colours if stdout isn't a terminal
  50   if not sys.stdout.isatty():
  51     return
  52
  53   try:
  54     import curses
  55   except ImportError:
  56     # Don't use colours if curses module can't be imported
  57     return
  58
  59   curses.setupterm()
  60
  61   _RESET_SEQ = curses.tigetstr("op")
  62
  63   setaf = curses.tigetstr("setaf")
  64   _INFO_SEQ = curses.tparm(setaf, curses.COLOR_GREEN)
  65   _WARNING_SEQ = curses.tparm(setaf, curses.COLOR_YELLOW)
  66   _ERROR_SEQ = curses.tparm(setaf, curses.COLOR_RED)
  67
  68
  69 _SetupColours()
  70
  71
  72 def AssertIn(item, sequence):
  73   """Raises an error when item is not in sequence.
  74
  75   """
  76   if item not in sequence:
  77     raise qa_error.Error('%r not in %r' % (item, sequence))
  78
  79
  80 def AssertEqual(first, second):
  81   """Raises an error when values aren't equal.
  82
  83   """
  84   if not first == second:
  85     raise qa_error.Error('%r == %r' % (first, second))
  86
  87
  88 def AssertNotEqual(first, second):
  89   """Raises an error when values are equal.
  90
  91   """
  92   if not first != second:
  93     raise qa_error.Error('%r != %r' % (first, second))
  94
  95
  96 def AssertMatch(string, pattern):
  97   """Raises an error when string doesn't match regexp pattern.
  98
  99   """
 100   if not re.match(pattern, string):
 101     raise qa_error.Error("%r doesn't match /%r/" % (string, pattern))
 102
 103
 104 def GetSSHCommand(node, cmd, strict=True):
 105   """Builds SSH command to be executed.
 106
 107   Args:
 108   - node: Node the command should run on
 109   - cmd: Command to be executed as a list with all parameters
 110   - strict: Whether to enable strict host key checking
 111
 112   """
 113   args = [ 'ssh', '-oEscapeChar=none', '-oBatchMode=yes', '-l', 'root', '-t' ]
 114
 115   if strict:
 116     tmp = 'yes'
 117   else:
 118     tmp = 'no'
 119   args.append('-oStrictHostKeyChecking=%s' % tmp)
 120   args.append('-oClearAllForwardings=yes')
 121   args.append('-oForwardAgent=yes')
 122   args.append(node)
 123   args.append(cmd)
 124
 125   print 'SSH:', utils.ShellQuoteArgs(args)
 126
 127   return args
 128
 129
 130 def StartSSH(node, cmd, strict=True):
 131   """Starts SSH.
 132
 133   """
 134   return subprocess.Popen(GetSSHCommand(node, cmd, strict=strict),
 135                           shell=False)
 136
 137
 138 def GetCommandOutput(node, cmd):
 139   """Returns the output of a command executed on the given node.
 140
 141   """
 142   p = subprocess.Popen(GetSSHCommand(node, cmd),
 143                        shell=False, stdout=subprocess.PIPE)
 144   AssertEqual(p.wait(), 0)
 145   return p.stdout.read()
 146
 147
 148 def UploadFile(node, src):
 149   """Uploads a file to a node and returns the filename.
 150
 151   Caller needs to remove the returned file on the node when it's not needed
 152   anymore.
 153
 154   """
 155   # Make sure nobody else has access to it while preserving local permissions
 156   mode = os.stat(src).st_mode & 0700
 157
 158   cmd = ('tmp=$(tempfile --mode %o --prefix gnt) && '
 159          '[[ -f "${tmp}" ]] && '
 160          'cat > "${tmp}" && '
 161          'echo "${tmp}"') % mode
 162
 163   f = open(src, 'r')
 164   try:
 165     p = subprocess.Popen(GetSSHCommand(node, cmd), shell=False, stdin=f,
 166                          stdout=subprocess.PIPE)
 167     AssertEqual(p.wait(), 0)
 168
 169     # Return temporary filename
 170     return p.stdout.read().strip()
 171   finally:
 172     f.close()
 173
 174
 175 def BackupFile(node, path):
 176   """Creates a backup of a file on the node and returns the filename.
 177
 178   Caller needs to remove the returned file on the node when it's not needed
 179   anymore.
 180
 181   """
 182   cmd = ("tmp=$(tempfile --prefix .gnt --directory=$(dirname %s)) && "
 183          "[[ -f \"$tmp\" ]] && "
 184          "cp %s $tmp && "
 185          "echo $tmp") % (utils.ShellQuote(path), utils.ShellQuote(path))
 186
 187   # Return temporary filename
 188   return GetCommandOutput(node, cmd).strip()
 189
 190
 191 def _ResolveName(cmd, key):
 192   """Helper function.
 193
 194   """
 195   master = qa_config.GetMasterNode()
 196
 197   output = GetCommandOutput(master['primary'], utils.ShellQuoteArgs(cmd))
 198   for line in output.splitlines():
 199     (lkey, lvalue) = line.split(':', 1)
 200     if lkey == key:
 201       return lvalue.lstrip()
 202   raise KeyError("Key not found")
 203
 204
 205 def ResolveInstanceName(instance):
 206   """Gets the full name of an instance.
 207
 208   """
 209   return _ResolveName(['gnt-instance', 'info', instance['name']],
 210                       'Instance name')
 211
 212
 213 def ResolveNodeName(node):
 214   """Gets the full name of a node.
 215
 216   """
 217   return _ResolveName(['gnt-node', 'info', node['primary']],
 218                       'Node name')
 219
 220
 221 def GetNodeInstances(node, secondaries=False):
 222   """Gets a list of instances on a node.
 223
 224   """
 225   master = qa_config.GetMasterNode()
 226   node_name = ResolveNodeName(node)
 227
 228   # Get list of all instances
 229   cmd = ['gnt-instance', 'list', '--separator=:', '--no-headers',
 230          '--output=name,pnode,snodes']
 231   output = GetCommandOutput(master['primary'], utils.ShellQuoteArgs(cmd))
 232
 233   instances = []
 234   for line in output.splitlines():
 235     (name, pnode, snodes) = line.split(':', 2)
 236     if ((not secondaries and pnode == node_name) or
 237         (secondaries and node_name in snodes.split(','))):
 238       instances.append(name)
 239
 240   return instances
 241
 242
 243 def _FormatWithColor(text, seq):
 244   if not seq:
 245     return text
 246   return "%s%s%s" % (seq, text, _RESET_SEQ)
 247
 248
 249 FormatWarning = lambda text: _FormatWithColor(text, _WARNING_SEQ)
 250 FormatError = lambda text: _FormatWithColor(text, _ERROR_SEQ)
 251 FormatInfo = lambda text: _FormatWithColor(text, _INFO_SEQ)