1 #!/usr/bin/python2.4

     2 #

     3 # Copyright 2008 Google Inc.

     4 #

     5 # Licensed under the Apache License, Version 2.0 (the "License");

     6 # you may not use this file except in compliance with the License.

     7 # You may obtain a copy of the License at

     8 #

     9 #      http://www.apache.org/licenses/LICENSE-2.0

    10 #

    11 # Unless required by applicable law or agreed to in writing, software

    12 # distributed under the License is distributed on an "AS IS" BASIS,

    13 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

    14 # See the License for the specific language governing permissions and

    15 # limitations under the License.

    16 

    17 import inspect

    18 

    19 class StubOutForTesting:

    20   """Sample Usage:

    21      You want os.path.exists() to always return true during testing.

    22 

    23      stubs = StubOutForTesting()

    24      stubs.Set(os.path, 'exists', lambda x: 1)

    25        ...

    26      stubs.UnsetAll()

    27 

    28      The above changes os.path.exists into a lambda that returns 1.  Once

    29      the ... part of the code finishes, the UnsetAll() looks up the old value

    30      of os.path.exists and restores it.

    31 

    32   """

    33   def __init__(self):

    34     self.cache = []

    35     self.stubs = []

    36 

    37   def __del__(self):

    38     self.SmartUnsetAll()

    39     self.UnsetAll()

    40 

    41   def SmartSet(self, obj, attr_name, new_attr):

    42     """Replace obj.attr_name with new_attr. This method is smart and works

    43        at the module, class, and instance level while preserving proper

    44        inheritance. It will not stub out C types however unless that has been

    45        explicitly allowed by the type.

    46 

    47        This method supports the case where attr_name is a staticmethod or a

    48        classmethod of obj.

    49 

    50        Notes:

    51       - If obj is an instance, then it is its class that will actually be

    52         stubbed. Note that the method Set() does not do that: if obj is

    53         an instance, it (and not its class) will be stubbed.

    54       - The stubbing is using the builtin getattr and setattr. So, the __get__

    55         and __set__ will be called when stubbing (TODO: A better idea would

    56         probably be to manipulate obj.__dict__ instead of getattr() and

    57         setattr()).

    58 

    59        Raises AttributeError if the attribute cannot be found.

    60     """

    61     if (inspect.ismodule(obj) or

    62         (not inspect.isclass(obj) and obj.__dict__.has_key(attr_name))):

    63       orig_obj = obj

    64       orig_attr = getattr(obj, attr_name)

    65 

    66     else:

    67       if not inspect.isclass(obj):

    68         mro = list(inspect.getmro(obj.__class__))

    69       else:

    70         mro = list(inspect.getmro(obj))

    71 

    72       mro.reverse()

    73 

    74       orig_attr = None

    75 

    76       for cls in mro:

    77         try:

    78           orig_obj = cls

    79           orig_attr = getattr(obj, attr_name)

    80         except AttributeError:

    81           continue

    82 

    83     if orig_attr is None:

    84       raise AttributeError("Attribute not found.")

    85 

    86     # Calling getattr() on a staticmethod transforms it to a 'normal' function.

    87     # We need to ensure that we put it back as a staticmethod.

    88     old_attribute = obj.__dict__.get(attr_name)

    89     if old_attribute is not None and isinstance(old_attribute, staticmethod):

    90       orig_attr = staticmethod(orig_attr)

    91 

    92     self.stubs.append((orig_obj, attr_name, orig_attr))

    93     setattr(orig_obj, attr_name, new_attr)

    94 

    95   def SmartUnsetAll(self):

    96     """Reverses all the SmartSet() calls, restoring things to their original

    97     definition.  Its okay to call SmartUnsetAll() repeatedly, as later calls

    98     have no effect if no SmartSet() calls have been made.

    99 

   100     """

   101     self.stubs.reverse()

   102 

   103     for args in self.stubs:

   104       setattr(*args)

   105 

   106     self.stubs = []

   107 

   108   def Set(self, parent, child_name, new_child):

   109     """Replace child_name's old definition with new_child, in the context

   110     of the given parent.  The parent could be a module when the child is a

   111     function at module scope.  Or the parent could be a class when a class'

   112     method is being replaced.  The named child is set to new_child, while

   113     the prior definition is saved away for later, when UnsetAll() is called.

   114 

   115     This method supports the case where child_name is a staticmethod or a

   116     classmethod of parent.

   117     """

   118     old_child = getattr(parent, child_name)

   119 

   120     old_attribute = parent.__dict__.get(child_name)

   121     if old_attribute is not None and isinstance(old_attribute, staticmethod):

   122       old_child = staticmethod(old_child)

   123 

   124     self.cache.append((parent, old_child, child_name))

   125     setattr(parent, child_name, new_child)

   126 

   127   def UnsetAll(self):

   128     """Reverses all the Set() calls, restoring things to their original

   129     definition.  Its okay to call UnsetAll() repeatedly, as later calls have

   130     no effect if no Set() calls have been made.

   131 

   132     """

   133     # Undo calls to Set() in reverse order, in case Set() was called on the

   134     # same arguments repeatedly (want the original call to be last one undone)

   135     self.cache.reverse()

   136 

   137     for (parent, old_child, child_name) in self.cache:

   138       setattr(parent, child_name, old_child)

   139     self.cache = []