Export backend.GetMasterInfo over the rpc layer
[ganeti-local] / 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 sys
28 import subprocess
29
30 from ganeti import utils
31
32 import qa_config
33 import qa_error
34
35
36 _INFO_SEQ = None
37 _WARNING_SEQ = None
38 _ERROR_SEQ = None
39 _RESET_SEQ = None
40
41
42 def _SetupColours():
43   """Initializes the colour constants.
44
45   """
46   global _INFO_SEQ, _WARNING_SEQ, _ERROR_SEQ, _RESET_SEQ
47
48   # Don't use colours if stdout isn't a terminal
49   if not sys.stdout.isatty():
50     return
51
52   try:
53     import curses
54   except ImportError:
55     # Don't use colours if curses module can't be imported
56     return
57
58   curses.setupterm()
59
60   _RESET_SEQ = curses.tigetstr("op")
61
62   setaf = curses.tigetstr("setaf")
63   _INFO_SEQ = curses.tparm(setaf, curses.COLOR_GREEN)
64   _WARNING_SEQ = curses.tparm(setaf, curses.COLOR_YELLOW)
65   _ERROR_SEQ = curses.tparm(setaf, curses.COLOR_RED)
66
67
68 _SetupColours()
69
70
71 def AssertIn(item, sequence):
72   """Raises an error when item is not in sequence.
73
74   """
75   if item not in sequence:
76     raise qa_error.Error('%r not in %r' % (item, sequence))
77
78
79 def AssertEqual(first, second):
80   """Raises an error when values aren't equal.
81
82   """
83   if not first == second:
84     raise qa_error.Error('%r == %r' % (first, second))
85
86
87 def AssertNotEqual(first, second):
88   """Raises an error when values are equal.
89
90   """
91   if not first != second:
92     raise qa_error.Error('%r != %r' % (first, second))
93
94
95 def GetSSHCommand(node, cmd, strict=True):
96   """Builds SSH command to be executed.
97
98   Args:
99   - node: Node the command should run on
100   - cmd: Command to be executed as a list with all parameters
101   - strict: Whether to enable strict host key checking
102
103   """
104   args = [ 'ssh', '-oEscapeChar=none', '-oBatchMode=yes', '-l', 'root', '-t' ]
105
106   if strict:
107     tmp = 'yes'
108   else:
109     tmp = 'no'
110   args.append('-oStrictHostKeyChecking=%s' % tmp)
111   args.append('-oClearAllForwardings=yes')
112   args.append('-oForwardAgent=yes')
113   args.append(node)
114   args.append(cmd)
115
116   print 'SSH:', utils.ShellQuoteArgs(args)
117
118   return args
119
120
121 def StartSSH(node, cmd, strict=True):
122   """Starts SSH.
123
124   """
125   return subprocess.Popen(GetSSHCommand(node, cmd, strict=strict),
126                           shell=False)
127
128
129 def GetCommandOutput(node, cmd):
130   """Returns the output of a command executed on the given node.
131
132   """
133   p = subprocess.Popen(GetSSHCommand(node, cmd),
134                        shell=False, stdout=subprocess.PIPE)
135   AssertEqual(p.wait(), 0)
136   return p.stdout.read()
137
138
139 def UploadFile(node, src):
140   """Uploads a file to a node and returns the filename.
141
142   Caller needs to remove the returned file on the node when it's not needed
143   anymore.
144   """
145   # Make sure nobody else has access to it while preserving local permissions
146   mode = os.stat(src).st_mode & 0700
147
148   cmd = ('tmp=$(tempfile --mode %o --prefix gnt) && '
149          '[[ -f "${tmp}" ]] && '
150          'cat > "${tmp}" && '
151          'echo "${tmp}"') % mode
152
153   f = open(src, 'r')
154   try:
155     p = subprocess.Popen(GetSSHCommand(node, cmd), shell=False, stdin=f,
156                          stdout=subprocess.PIPE)
157     AssertEqual(p.wait(), 0)
158
159     # Return temporary filename
160     return p.stdout.read().strip()
161   finally:
162     f.close()
163
164
165 def _ResolveName(cmd, key):
166   """Helper function.
167
168   """
169   master = qa_config.GetMasterNode()
170
171   output = GetCommandOutput(master['primary'], utils.ShellQuoteArgs(cmd))
172   for line in output.splitlines():
173     (lkey, lvalue) = line.split(':', 1)
174     if lkey == key:
175       return lvalue.lstrip()
176   raise KeyError("Key not found")
177
178
179 def ResolveInstanceName(instance):
180   """Gets the full name of an instance.
181
182   """
183   return _ResolveName(['gnt-instance', 'info', instance['name']],
184                       'Instance name')
185
186
187 def ResolveNodeName(node):
188   """Gets the full name of a node.
189
190   """
191   return _ResolveName(['gnt-node', 'info', node['primary']],
192                       'Node name')
193
194
195 def GetNodeInstances(node, secondaries=False):
196   """Gets a list of instances on a node.
197
198   """
199   master = qa_config.GetMasterNode()
200   node_name = ResolveNodeName(node)
201
202   # Get list of all instances
203   cmd = ['gnt-instance', 'list', '--separator=:', '--no-headers',
204          '--output=name,pnode,snodes']
205   output = GetCommandOutput(master['primary'], utils.ShellQuoteArgs(cmd))
206
207   instances = []
208   for line in output.splitlines():
209     (name, pnode, snodes) = line.split(':', 2)
210     if ((not secondaries and pnode == node_name) or
211         (secondaries and node_name in snodes.split(','))):
212       instances.append(name)
213
214   return instances
215
216
217 def _FormatWithColor(text, seq):
218   if not seq:
219     return text
220   return "%s%s%s" % (seq, text, _RESET_SEQ)
221
222
223 FormatWarning = lambda text: _FormatWithColor(text, _WARNING_SEQ)
224 FormatError = lambda text: _FormatWithColor(text, _ERROR_SEQ)
225 FormatInfo = lambda text: _FormatWithColor(text, _INFO_SEQ)