src/tools/swarming_client/tests/subprocess42_test.py

   1 #!/usr/bin/env python
   2 # Copyright 2013 The Swarming Authors. All rights reserved.
   3 # Use of this source code is governed under the Apache License, Version 2.0 that
   4 # can be found in the LICENSE file.
   5
   6 import itertools
   7 import logging
   8 import os
   9 import sys
  10 import unittest
  11
  12 ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
  13 sys.path.insert(0, ROOT_DIR)
  14
  15 from utils import subprocess42
  16
  17 OUTPUT = os.path.join(ROOT_DIR, 'tests', 'subprocess42', 'output.py')
  18
  19
  20 # Disable pre-set unbuffered output to not interfere with the testing being done
  21 # here. Otherwise everything would test with unbuffered; which is fine but
  22 # that's not what we specifically want to test here.
  23 ENV = os.environ.copy()
  24 ENV.pop('PYTHONUNBUFFERED', None)
  25
  26
  27 def to_native_eol(string):
  28   if string is None:
  29     return string
  30   if sys.platform == 'win32':
  31     return string.replace('\n', '\r\n')
  32   return string
  33
  34
  35 def get_output_sleep_proc(flush, unbuffered, sleep_duration):
  36   """Returns process with universal_newlines=True that prints to stdout before
  37   after a sleep.
  38
  39   It also optionally sys.stdout.flush() before the sleep and optionally enable
  40   unbuffered output in python.
  41   """
  42   command = [
  43     'import sys,time',
  44     'print(\'A\')',
  45   ]
  46   if flush:
  47     # Sadly, this doesn't work otherwise in some combination.
  48     command.append('sys.stdout.flush()')
  49   command.extend((
  50     'time.sleep(%s)' % sleep_duration,
  51     'print(\'B\')',
  52   ))
  53   cmd = [sys.executable, '-c', ';'.join(command)]
  54   if unbuffered:
  55     cmd.append('-u')
  56   return subprocess42.Popen(
  57       cmd, env=ENV, stdout=subprocess42.PIPE, universal_newlines=True)
  58
  59
  60 def get_output_sleep_proc_err(sleep_duration):
  61   """Returns process with universal_newlines=True that prints to stderr before
  62   and after a sleep.
  63   """
  64   command = [
  65     'import sys,time',
  66     'sys.stderr.write(\'A\\n\')',
  67   ]
  68   command.extend((
  69     'time.sleep(%s)' % sleep_duration,
  70     'sys.stderr.write(\'B\\n\')',
  71   ))
  72   cmd = [sys.executable, '-c', ';'.join(command)]
  73   return subprocess42.Popen(
  74       cmd, env=ENV, stderr=subprocess42.PIPE, universal_newlines=True)
  75
  76
  77 class Subprocess42Test(unittest.TestCase):
  78   def test_call_with_timeout(self):
  79     timedout = 1 if sys.platform == 'win32' else -9
  80     # Format is:
  81     # ( (cmd, stderr_pipe, timeout), (stdout, stderr, returncode) ), ...
  82     test_data = [
  83       # 0 means no timeout, like None.
  84       (
  85         (['out_sleeping', '0.001', 'out_slept', 'err_print'], None, 0),
  86         ('Sleeping.\nSlept.\n', None, 0),
  87       ),
  88       (
  89         (['err_print'], subprocess42.STDOUT, 0),
  90         ('printing', None, 0),
  91       ),
  92       (
  93         (['err_print'], subprocess42.PIPE, 0),
  94         ('', 'printing', 0),
  95       ),
  96
  97       # On a loaded system, this can be tight.
  98       (
  99         (['out_sleeping', 'out_flush', '100', 'out_slept'], None, 0.5),
 100         ('Sleeping.\n', '', timedout),
 101       ),
 102       (
 103         (
 104           # Note that err_flush is necessary on Windows but not on the other
 105           # OSes. This means the likelihood of missing stderr output from a
 106           # killed child process on Windows is much higher than on other OSes.
 107           [
 108             'out_sleeping', 'out_flush', 'err_print', 'err_flush', '100',
 109             'out_slept',
 110           ],
 111           subprocess42.PIPE,
 112           0.5),
 113         ('Sleeping.\n', 'printing', timedout),
 114       ),
 115
 116       (
 117         (['out_sleeping', '0.001', 'out_slept'], None, 100),
 118         ('Sleeping.\nSlept.\n', '', 0),
 119       ),
 120     ]
 121     for i, (data, expected) in enumerate(test_data):
 122       stdout, stderr, code, duration = subprocess42.call_with_timeout(
 123           [sys.executable, OUTPUT] + data[0],
 124           env=ENV,
 125           stderr=data[1],
 126           timeout=data[2])
 127       self.assertTrue(duration > 0.0001, (data, duration))
 128       self.assertEqual(
 129           (i, stdout, stderr, code),
 130           (i,
 131             to_native_eol(expected[0]),
 132             to_native_eol(expected[1]),
 133             expected[2]))
 134
 135       # Try again with universal_newlines=True.
 136       stdout, stderr, code, duration = subprocess42.call_with_timeout(
 137           [sys.executable, OUTPUT] + data[0],
 138           env=ENV,
 139           stderr=data[1],
 140           timeout=data[2],
 141           universal_newlines=True)
 142       self.assertTrue(duration > 0.0001, (data, duration))
 143       self.assertEqual(
 144           (i, stdout, stderr, code),
 145           (i,) + expected)
 146
 147   def test_recv_any(self):
 148     # Test all pipe direction and output scenarios.
 149     combinations = [
 150       {
 151         'cmd': ['out_print', 'err_print'],
 152         'stdout': None,
 153         'stderr': None,
 154         'expected': {},
 155       },
 156       {
 157         'cmd': ['out_print', 'err_print'],
 158         'stdout': None,
 159         'stderr': subprocess42.STDOUT,
 160         'expected': {},
 161       },
 162
 163       {
 164         'cmd': ['out_print'],
 165         'stdout': subprocess42.PIPE,
 166         'stderr': subprocess42.PIPE,
 167         'expected': {'stdout': 'printing'},
 168       },
 169       {
 170         'cmd': ['out_print'],
 171         'stdout': subprocess42.PIPE,
 172         'stderr': None,
 173         'expected': {'stdout': 'printing'},
 174       },
 175       {
 176         'cmd': ['out_print'],
 177         'stdout': subprocess42.PIPE,
 178         'stderr': subprocess42.STDOUT,
 179         'expected': {'stdout': 'printing'},
 180       },
 181
 182       {
 183         'cmd': ['err_print'],
 184         'stdout': subprocess42.PIPE,
 185         'stderr': subprocess42.PIPE,
 186         'expected': {'stderr': 'printing'},
 187       },
 188       {
 189         'cmd': ['err_print'],
 190         'stdout': None,
 191         'stderr': subprocess42.PIPE,
 192         'expected': {'stderr': 'printing'},
 193       },
 194       {
 195         'cmd': ['err_print'],
 196         'stdout': subprocess42.PIPE,
 197         'stderr': subprocess42.STDOUT,
 198         'expected': {'stdout': 'printing'},
 199       },
 200
 201       {
 202         'cmd': ['out_print', 'err_print'],
 203         'stdout': subprocess42.PIPE,
 204         'stderr': subprocess42.PIPE,
 205         'expected': {'stderr': 'printing', 'stdout': 'printing'},
 206       },
 207       {
 208         'cmd': ['out_print', 'err_print'],
 209         'stdout': subprocess42.PIPE,
 210         'stderr': subprocess42.STDOUT,
 211         'expected': {'stdout': 'printingprinting'},
 212       },
 213     ]
 214     for i, testcase in enumerate(combinations):
 215       cmd = [sys.executable, OUTPUT] + testcase['cmd']
 216       p = subprocess42.Popen(
 217           cmd, env=ENV, stdout=testcase['stdout'], stderr=testcase['stderr'])
 218       actual = {}
 219       while p.poll() is None:
 220         pipe, data = p.recv_any()
 221         if data:
 222           actual.setdefault(pipe, '')
 223           actual[pipe] += data
 224
 225       # The process exited, read any remaining data in the pipes.
 226       while True:
 227         pipe, data = p.recv_any()
 228         if pipe is None:
 229           break
 230         actual.setdefault(pipe, '')
 231         actual[pipe] += data
 232       self.assertEqual(
 233           testcase['expected'],
 234           actual,
 235           (i, testcase['cmd'], testcase['expected'], actual))
 236       self.assertEqual((None, None), p.recv_any())
 237       self.assertEqual(0, p.returncode)
 238
 239   def test_recv_any_different_buffering(self):
 240     # Specifically test all buffering scenarios.
 241     for flush, unbuffered in itertools.product([True, False], [True, False]):
 242       actual = ''
 243       proc = get_output_sleep_proc(flush, unbuffered, 0.5)
 244       while True:
 245         p, data = proc.recv_any()
 246         if not p:
 247           break
 248         self.assertEqual('stdout', p)
 249         self.assertTrue(data, (p, data))
 250         actual += data
 251
 252       self.assertEqual('A\nB\n', actual)
 253       # Contrary to yield_any() or recv_any(0), wait() needs to be used here.
 254       proc.wait()
 255       self.assertEqual(0, proc.returncode)
 256
 257   def test_recv_any_timeout_0(self):
 258     # rec_any() is expected to timeout and return None with no data pending at
 259     # least once, due to the sleep of 'duration' and the use of timeout=0.
 260     for flush, unbuffered in itertools.product([True, False], [True, False]):
 261       for duration in (0.05, 0.1, 0.5, 2):
 262         try:
 263           actual = ''
 264           proc = get_output_sleep_proc(flush, unbuffered, duration)
 265           got_none = False
 266           while True:
 267             p, data = proc.recv_any(timeout=0)
 268             if not p:
 269               if proc.poll() is None:
 270                 got_none = True
 271                 continue
 272               break
 273             self.assertEqual('stdout', p)
 274             self.assertTrue(data, (p, data))
 275             actual += data
 276
 277           self.assertEqual('A\nB\n', actual)
 278           self.assertEqual(0, proc.returncode)
 279           self.assertEqual(True, got_none)
 280           break
 281         except AssertionError:
 282           if duration != 2:
 283             print('Sleeping rocks. Trying slower.')
 284             continue
 285           raise
 286
 287   def _test_recv_common(self, proc, is_err):
 288     actual = ''
 289     while True:
 290       if is_err:
 291         data = proc.recv_err()
 292       else:
 293         data = proc.recv_out()
 294       if not data:
 295         break
 296       self.assertTrue(data)
 297       actual += data
 298
 299     self.assertEqual('A\nB\n', actual)
 300     proc.wait()
 301     self.assertEqual(0, proc.returncode)
 302
 303   def test_yield_any_no_timeout(self):
 304     for duration in (0.05, 0.1, 0.5, 2):
 305       try:
 306         proc = get_output_sleep_proc(True, True, duration)
 307         expected = [
 308           'A\n',
 309           'B\n',
 310         ]
 311         for p, data in proc.yield_any():
 312           self.assertEqual('stdout', p)
 313           self.assertEqual(expected.pop(0), data)
 314         self.assertEqual(0, proc.returncode)
 315         self.assertEqual([], expected)
 316         break
 317       except AssertionError:
 318         if duration != 2:
 319           print('Sleeping rocks. Trying slower.')
 320           continue
 321         raise
 322
 323   def test_yield_any_hard_timeout(self):
 324     # Kill the process due to hard_timeout.
 325     proc = get_output_sleep_proc(True, True, 10)
 326     got_none = False
 327     actual = ''
 328     for p, data in proc.yield_any(hard_timeout=1):
 329       if not data:
 330         got_none = True
 331         continue
 332       self.assertEqual('stdout', p)
 333       actual += data
 334     if sys.platform == 'win32':
 335       self.assertEqual(1, proc.returncode)
 336     else:
 337       self.assertEqual(-9, proc.returncode)
 338     self.assertEqual('A\n', actual)
 339     # No None is returned, since it's not using soft_timeout.
 340     self.assertEqual(False, got_none)
 341
 342   def test_yield_any_soft_timeout_0(self):
 343     # rec_any() is expected to timeout and return None with no data pending at
 344     # least once, due to the sleep of 'duration' and the use of timeout=0.
 345     for duration in (0.05, 0.1, 0.5, 2):
 346       try:
 347         proc = get_output_sleep_proc(True, True, duration)
 348         expected = [
 349           'A\n',
 350           'B\n',
 351         ]
 352         got_none = False
 353         for p, data in proc.yield_any(soft_timeout=0):
 354           if not p:
 355             got_none = True
 356             continue
 357           self.assertEqual('stdout', p)
 358           self.assertEqual(expected.pop(0), data)
 359         self.assertEqual(0, proc.returncode)
 360         self.assertEqual([], expected)
 361         self.assertEqual(True, got_none)
 362         break
 363       except AssertionError:
 364         if duration != 2:
 365           print('Sleeping rocks. Trying slower.')
 366           continue
 367         raise
 368
 369
 370 if __name__ == '__main__':
 371   if '-v' in sys.argv:
 372     unittest.TestCase.maxDiff = None
 373   logging.basicConfig(
 374       level=logging.DEBUG if '-v' in sys.argv else logging.ERROR)
 375   unittest.main()