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