src/tools/swarming_client/tests/test_utils.py

   1 # Copyright 2014 The Swarming Authors. All rights reserved.
   2 # Use of this source code is governed under the Apache License, Version 2.0 that
   3 # can be found in the LICENSE file.
   4
   5 import logging
   6 import os
   7 import sys
   8 import unittest
   9
  10 _UMASK = None
  11
  12 class EnvVars(object):
  13   """Context manager for environment variables.
  14
  15   Passed a dict to the constructor it sets variables named with the key to the
  16   value.  Exiting the context causes all the variables named with the key to be
  17   restored to their value before entering the context.
  18   """
  19   def __init__(self, var_map):
  20     self.var_map = var_map
  21     self._backup = None
  22
  23   def __enter__(self):
  24     self._backup = os.environ
  25     os.environ = os.environ.copy()
  26     os.environ.update(self.var_map)
  27
  28   def __exit__(self, exc_type, exc_value, traceback):
  29     os.environ = self._backup
  30
  31
  32 def umask():
  33   """Returns current process umask without modifying it."""
  34   global _UMASK
  35   if _UMASK is None:
  36     _UMASK = os.umask(0777)
  37     os.umask(_UMASK)
  38   return _UMASK
  39
  40
  41 def main():
  42   logging.basicConfig(
  43       level=logging.DEBUG if '-v' in sys.argv else logging.ERROR,
  44       format='%(levelname)5s %(filename)15s(%(lineno)3d): %(message)s')
  45   if '-v' in sys.argv:
  46     unittest.TestCase.maxDiff = None
  47   # Use an unusual umask.
  48   os.umask(0070)
  49   unittest.main()