command_executer_test.py revision f81680c018729fd4499e1e200d04b48c4b90127c 
1#!/usr/bin/python2.6
2#
3# Copyright 2011 Google Inc. All Rights Reserved.
4#
5
6__author__ = 'kbaclawski@google.com (Krystian Baclawski)'
7
8import cStringIO
9import logging
10import os
11import signal
12import socket
13import sys
14import time
15import unittest
16
17
18def AddScriptDirToPath():
19  """Required for remote python script execution."""
20  path = os.path.abspath(__file__)
21
22  for _ in range(3):
23    path, _ = os.path.split(path)
24
25  if not path in sys.path:
26    sys.path.append(path)
27
28
29AddScriptDirToPath()
30
31from automation.common.command_executer import CommandExecuter
32
33
34class LoggerMock(object):
35  def LogCmd(self, cmd, machine='', user=''):
36    if machine:
37      logging.info('[%s] Executing: %s', machine, cmd)
38    else:
39      logging.info('Executing: %s', cmd)
40
41  def LogError(self, msg):
42    logging.error(msg)
43
44  def LogWarning(self, msg):
45    logging.warning(msg)
46
47  def LogOutput(self, msg):
48    logging.info(msg)
49
50
51class CommandExecuterUnderTest(CommandExecuter):
52  def __init__(self):
53    CommandExecuter.__init__(self, logger_to_set=LoggerMock())
54
55    # We will record stdout and stderr.
56    self._stderr = cStringIO.StringIO()
57    self._stdout = cStringIO.StringIO()
58
59  @property
60  def stdout(self):
61    return self._stdout.getvalue()
62
63  @property
64  def stderr(self):
65    return self._stderr.getvalue()
66
67  def DataReceivedOnOutput(self, data):
68    self._stdout.write(data)
69
70  def DataReceivedOnError(self, data):
71    self._stderr.write(data)
72
73
74class CommandExecuterLocalTests(unittest.TestCase):
75  HOSTNAME = None
76
77  def setUp(self):
78    self._executer = CommandExecuterUnderTest()
79
80  def tearDown(self):
81    pass
82
83  def RunCommand(self, method, **kwargs):
84    program = os.path.abspath(sys.argv[0])
85
86    return self._executer.RunCommand(
87        '%s runHelper %s' % (program, method), machine=self.HOSTNAME, **kwargs)
88
89  def testCommandTimeout(self):
90    exit_code = self.RunCommand('SleepForMinute', command_timeout=3)
91
92    self.assertTrue(-exit_code in [signal.SIGTERM, signal.SIGKILL],
93                    'Invalid exit code: %d' % exit_code)
94
95  def testCommandTimeoutIfSigTermIgnored(self):
96    exit_code = self.RunCommand('IgnoreSigTerm', command_timeout=3)
97
98    self.assertTrue(-exit_code in [signal.SIGTERM, signal.SIGKILL])
99
100  def testCommandSucceeded(self):
101    self.assertFalse(self.RunCommand('ReturnTrue'))
102
103  def testCommandFailed(self):
104    self.assertTrue(self.RunCommand('ReturnFalse'))
105
106  def testStringOnOutputStream(self):
107    self.assertFalse(self.RunCommand('EchoToOutputStream'))
108    self.assertEquals(self._executer.stderr, '')
109    self.assertEquals(self._executer.stdout, 'test')
110
111  def testStringOnErrorStream(self):
112    self.assertFalse(self.RunCommand('EchoToErrorStream'))
113    self.assertEquals(self._executer.stderr, 'test')
114    self.assertEquals(self._executer.stdout, '')
115
116  def testOutputStreamNonInteractive(self):
117    self.assertFalse(self.RunCommand('IsOutputStreamInteractive'),
118                     'stdout stream is a terminal!')
119
120  def testErrorStreamNonInteractive(self):
121    self.assertFalse(self.RunCommand('IsErrorStreamInteractive'),
122                     'stderr stream is a terminal!')
123
124  def testAttemptToRead(self):
125    self.assertFalse(self.RunCommand('WaitForInput', command_timeout=3))
126
127  def testInterruptedProcess(self):
128    self.assertEquals(self.RunCommand('TerminateBySigAbrt'), -signal.SIGABRT)
129
130
131class CommandExecuterRemoteTests(CommandExecuterLocalTests):
132  HOSTNAME = socket.gethostname()
133
134  def testCommandTimeoutIfSigTermIgnored(self):
135    exit_code = self.RunCommand('IgnoreSigTerm', command_timeout=6)
136
137    self.assertEquals(exit_code, 255)
138
139    lines = self._executer.stdout.splitlines()
140    pid = int(lines[0])
141
142    try:
143      with open('/proc/%d/cmdline' % pid) as f:
144        cmdline = f.read()
145    except IOError:
146      cmdline = ''
147
148    self.assertFalse('IgnoreSigTerm' in cmdline, 'Process is still alive.')
149
150
151class CommandExecuterTestHelpers(object):
152  def SleepForMinute(self):
153    time.sleep(60)
154    return 1
155
156  def ReturnTrue(self):
157    return 0
158
159  def ReturnFalse(self):
160    return 1
161
162  def EchoToOutputStream(self):
163    sys.stdout.write('test')
164    return 0
165
166  def EchoToErrorStream(self):
167    sys.stderr.write('test')
168    return 0
169
170  def IsOutputStreamInteractive(self):
171    return sys.stdout.isatty()
172
173  def IsErrorStreamInteractive(self):
174    return sys.stderr.isatty()
175
176  def IgnoreSigTerm(self):
177    os.write(1, '%d' % os.getpid())
178    signal.signal(signal.SIGTERM, signal.SIG_IGN)
179    time.sleep(30)
180    return 0
181
182  def WaitForInput(self):
183    try:
184      # can only read end-of-file marker
185      return os.read(0, 1) != ''
186    except OSError:
187      # that means that stdin descriptor is closed
188      return 0
189
190  def TerminateBySigAbrt(self):
191    os.kill(os.getpid(), signal.SIGABRT)
192    return 0
193
194
195if __name__ == '__main__':
196  FORMAT = '%(asctime)-15s %(levelname)s %(message)s'
197  logging.basicConfig(format=FORMAT, level=logging.DEBUG)
198
199  if len(sys.argv) > 1:
200    if sys.argv[1] == 'runHelper':
201      helpers = CommandExecuterTestHelpers()
202      sys.exit(getattr(helpers, sys.argv[2])())
203
204  unittest.main()
205