graphyte_utils: Guard current working directory

Currently in Graphyte, all plugins run on the same process. It is
possible that the current working directory is changed by other
plugins. One way to solve this is running each plugins on a different
process. However, it needs a lot of work. Another approach is using
decorator. This CL adds a decorater which can guard the current working
directory and uses it on dummy_dut and dummy_inst for example.

BUG=None
TEST=make test PYTHON=python2; manually test on devices

Change-Id: I16c7f25826d93cf77cbc64c41ee1941789f93523
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/graphyte/+/1958081
Reviewed-by: Yong Hong <yhong@chromium.org>
Commit-Queue: Cheng Yueh <cyueh@chromium.org>
Tested-by: Cheng Yueh <cyueh@chromium.org>
diff --git a/graphyte/dut/sample/dummy_dut.py b/graphyte/dut/sample/dummy_dut.py
index 2c673c2..372b420 100644
--- a/graphyte/dut/sample/dummy_dut.py
+++ b/graphyte/dut/sample/dummy_dut.py
@@ -5,18 +5,29 @@
 
 """The mock DUT class. """
 
+import os
+import tempfile
+
 import graphyte_common  # pylint: disable=unused-import
 from graphyte.dut import DUTBase
+from graphyte.default_setting import GRAPHYTE_DIR
 from graphyte.default_setting import logger
 from graphyte.utils.graphyte_utils import MakeMockPassResult
+from graphyte.utils import graphyte_utils
 from graphyte.utils import process_utils
 
+# DUT_JAIL makes the classes it decorates share the same working directory.
+# It protects working directory from being changed by other plugins and also
+# prevents working directory of other plugins being changed.
+DUT_JAIL = graphyte_utils.IsolateCWD()
 
+@DUT_JAIL.IsolateAllMethods
 class DUT(DUTBase):
   name = 'Mock DUT Plugin'
 
   def __init__(self, **kwargs):
     super(DUT, self).__init__(**kwargs)
+    self.temp_dir = None
     self.controllers = {
         'WLAN': self.WlanController(self),
         'BLUETOOTH': self.BluetoothController(self),
@@ -24,14 +35,22 @@
 
   def _Initialize(self):
     logger.info('DUT Initialize')
+    # Change to other working directory if you want.
+    self.temp_dir = tempfile.mkdtemp('_DUT')
+    logger.info('DUT Change diretory to %r', self.temp_dir)
+    os.chdir(self.temp_dir)
 
   def _Terminate(self):
-    logger.info('DUT Terminate')
+    logger.info('DUT Terminate at %r', os.getcwd())
+    os.chdir(GRAPHYTE_DIR)
+    if self.temp_dir:
+      os.rmdir(self.temp_dir)
 
   def GetVersion(self):
     logger.info('DUT GetVersion')
     return process_utils.CheckOutput(['uname', '-a']).strip()
 
+  @DUT_JAIL.IsolateAllMethods
   class WlanController(DUTBase.WlanControllerBase):
     def _Initialize(self):
       logger.info('DUT Wlan Initialize')
@@ -56,6 +75,7 @@
       logger.info('DUT Wlan RxGetResult')
       return MakeMockPassResult(result_limit)
 
+  @DUT_JAIL.IsolateAllMethods
   class BluetoothController(DUTBase.BluetoothControllerBase):
     def _Initialize(self):
       logger.info('DUT Bluetooth Initialize')
@@ -78,6 +98,7 @@
       logger.info('DUT Bluetooth RxGetResult')
       return MakeMockPassResult(result_limit)
 
+  @DUT_JAIL.IsolateAllMethods
   class ZigbeeController(DUTBase.ZigbeeControllerBase):
     def _Initialize(self):
       logger.info('DUT Zigbee Initialize')
diff --git a/graphyte/inst/sample/dummy_inst.py b/graphyte/inst/sample/dummy_inst.py
index e59136d..f6c2228 100644
--- a/graphyte/inst/sample/dummy_inst.py
+++ b/graphyte/inst/sample/dummy_inst.py
@@ -5,20 +5,31 @@
 
 """The mock instrument class. """
 
+import os
+import tempfile
+
 import graphyte_common  # pylint: disable=unused-import
 from graphyte.inst import InstBase
 from graphyte.testplan import ChainMaskToList
+from graphyte.default_setting import GRAPHYTE_DIR
 from graphyte.default_setting import logger
 from graphyte.utils.graphyte_utils import MakeMockPassResult
+from graphyte.utils import graphyte_utils
 from graphyte.utils import process_utils
 
+# INST_JAIL makes the classes it decorates share the same working directory.
+# It protects working directory from being changed by other plugins and also
+# prevents working directory of other plugins being changed.
+INST_JAIL = graphyte_utils.IsolateCWD()
 
+@INST_JAIL.IsolateAllMethods
 class Inst(InstBase):
   name = 'Mock inst Plugin'
   VALID_PORT_NAMES = ['dummy_port']
 
   def __init__(self, **kwargs):
     super(Inst, self).__init__(**kwargs)
+    self.temp_dir = None
     self.controllers = {
         'WLAN': self.WlanController(self),
         'BLUETOOTH': self.BluetoothController(self),
@@ -26,9 +37,16 @@
 
   def _Initialize(self):
     logger.info('Inst Initialize')
+    # Change to other working directory if you want.
+    self.temp_dir = tempfile.mkdtemp('_Inst')
+    logger.info('Inst Change diretory to %r', self.temp_dir)
+    os.chdir(self.temp_dir)
 
   def _Terminate(self):
-    logger.info('Inst Terminate')
+    logger.info('Inst Terminate at %r', os.getcwd())
+    os.chdir(GRAPHYTE_DIR)
+    if self.temp_dir:
+      os.rmdir(self.temp_dir)
 
   def _SetPortConfig(self, port_mapping, pathloss):
     logger.info('Inst _SetPortConfig')
@@ -43,6 +61,7 @@
     logger.info('Inst GetVersion')
     return process_utils.CheckOutput(['uname', '-a']).strip()
 
+  @INST_JAIL.IsolateAllMethods
   class WlanController(InstBase.WlanControllerBase):
     def __init__(self, inst):
       super(Inst.WlanController, self).__init__(inst)
@@ -79,6 +98,7 @@
     def _RxStop(self, **kwargs):
       logger.info('Wlan RxStop')
 
+  @INST_JAIL.IsolateAllMethods
   class BluetoothController(InstBase.BluetoothControllerBase):
     def _Initialize(self):
       logger.info('Inst Bluetooth Initialize')
@@ -104,6 +124,7 @@
     def _RxStop(self, **kwargs):
       logger.info('Bluetooth RxStop')
 
+  @INST_JAIL.IsolateAllMethods
   class ZigbeeController(InstBase.ZigbeeControllerBase):
     def _Initialize(self):
       logger.info('Inst Zigbee Initialize')
diff --git a/graphyte/utils/graphyte_utils.py b/graphyte/utils/graphyte_utils.py
index 07137db..fbcc7d0 100644
--- a/graphyte/utils/graphyte_utils.py
+++ b/graphyte/utils/graphyte_utils.py
@@ -208,6 +208,41 @@
   return cls
 
 
+class IsolateCWD(object):
+  """The decorator that isolates changes of current working directory.
+
+  The methods decorated with the same IsolateCWD share the same working
+  directory. It protects working directory from being changed by other plugins
+  and also prevents working directory of other plugins being changed.
+  """
+  def __init__(self):
+    self._outside_cwd = None
+    self._inside_cwd = os.getcwd()
+    self._depth = 0
+
+  def IsolateFunc(self, func):
+    def Wrapper(*args, **kwargs):
+      try:
+        # _depth equals to zero implies that its call from outside.
+        if self._depth == 0:
+          self._outside_cwd = os.getcwd()
+          os.chdir(self._inside_cwd)
+        self._depth += 1
+        return func(*args, **kwargs)
+      finally:
+        self._depth -= 1
+        # _depth equals to zero implies that the call ends to outside.
+        if self._depth == 0:
+          self._inside_cwd = os.getcwd()
+          os.chdir(self._outside_cwd)
+    return Wrapper
+
+  def IsolateAllMethods(self, cls):
+    for func_name, func in inspect.getmembers(cls, inspect.ismethod):
+      setattr(cls, func_name, self.IsolateFunc(func))
+    return cls
+
+
 def AssertRangeContain(condition_message, arg_name, outer, inner):
   """Check if the outer interval contains the inner interval.
 
diff --git a/graphyte/utils/graphyte_utils_unittest.py b/graphyte/utils/graphyte_utils_unittest.py
index faaccbc..e67e831 100755
--- a/graphyte/utils/graphyte_utils_unittest.py
+++ b/graphyte/utils/graphyte_utils_unittest.py
@@ -5,6 +5,8 @@
 
 import math
 import mock
+import os
+import tempfile
 import unittest
 
 import graphyte_common  # pylint: disable=unused-import
@@ -172,3 +174,39 @@
     mock_logger.debug.assert_not_called()
     obj.Foo(3)
     mock_logger.debug.assert_called_with('Calling %s(%s)', 'FakeClass.Foo', '3')
+
+
+class IsolateAllMethodsCurrentWorkingDirectoryTest(unittest.TestCase):
+  def setUp(self):
+    self._my_jail = graphyte_utils.IsolateCWD()
+    self._initial_dir = os.getcwd()
+    self._temp_dir = tempfile.mkdtemp()
+
+  def tearDown(self):
+    os.chdir(self._initial_dir)
+    if self._temp_dir:
+      os.rmdir(self._temp_dir)
+
+  def testIsolating(self):
+    @self._my_jail.IsolateAllMethods
+    class FakeClassOne(object):
+      def Foo(self, temp_dir):
+        os.chdir(temp_dir)
+      def Baz(self):
+        raise Exception('Test')
+    @self._my_jail.IsolateAllMethods
+    class FakeClassTwo(object):
+      def Bar(self):
+        return os.getcwd()
+    obj_1 = FakeClassOne()
+    obj_2 = FakeClassTwo()
+    obj_1.Foo(self._temp_dir)
+    # Test if changing inside affects outside.
+    self.assertEqual(os.getcwd(), self._initial_dir)
+    # Test the working directory inside.
+    self.assertEqual(obj_2.Bar(), self._temp_dir)
+    # Raise an exception in the decorated object.
+    self.assertRaisesRegexp(Exception, 'Test', obj_1.Baz)
+    # Redo test after an exception to check exception safety.
+    self.assertEqual(os.getcwd(), self._initial_dir)
+    self.assertEqual(obj_2.Bar(), self._temp_dir)