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

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