1fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville#!/usr/bin/python2.4 2fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville# 3fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville# Copyright 2008 Google Inc. 4fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville# 5fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville# Licensed under the Apache License, Version 2.0 (the "License"); 6fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville# you may not use this file except in compliance with the License. 7fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville# You may obtain a copy of the License at 8fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville# 9fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville# http://www.apache.org/licenses/LICENSE-2.0 10fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville# 11fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville# Unless required by applicable law or agreed to in writing, software 12fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville# distributed under the License is distributed on an "AS IS" BASIS, 13fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville# See the License for the specific language governing permissions and 15fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville# limitations under the License. 16fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville 17fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville# This file is used for testing. The original is at: 18fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville# http://code.google.com/p/pymox/ 19fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville 20fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Savilleclass StubOutForTesting: 21fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville """Sample Usage: 22fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville You want os.path.exists() to always return true during testing. 23fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville 24fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville stubs = StubOutForTesting() 25fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville stubs.Set(os.path, 'exists', lambda x: 1) 26fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville ... 27fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville stubs.UnsetAll() 28fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville 29fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville The above changes os.path.exists into a lambda that returns 1. Once 30fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville the ... part of the code finishes, the UnsetAll() looks up the old value 31fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville of os.path.exists and restores it. 32fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville 33fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville """ 34fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville def __init__(self): 35fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville self.cache = [] 36fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville self.stubs = [] 37fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville 38fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville def __del__(self): 39fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville self.SmartUnsetAll() 40fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville self.UnsetAll() 41fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville 42fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville def SmartSet(self, obj, attr_name, new_attr): 43fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville """Replace obj.attr_name with new_attr. This method is smart and works 44fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville at the module, class, and instance level while preserving proper 45fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville inheritance. It will not stub out C types however unless that has been 46fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville explicitly allowed by the type. 47fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville 48fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville This method supports the case where attr_name is a staticmethod or a 49fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville classmethod of obj. 50fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville 51fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville Notes: 52fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville - If obj is an instance, then it is its class that will actually be 53fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville stubbed. Note that the method Set() does not do that: if obj is 54fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville an instance, it (and not its class) will be stubbed. 55fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville - The stubbing is using the builtin getattr and setattr. So, the __get__ 56fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville and __set__ will be called when stubbing (TODO: A better idea would 57fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville probably be to manipulate obj.__dict__ instead of getattr() and 58fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville setattr()). 59fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville 60fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville Raises AttributeError if the attribute cannot be found. 61fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville """ 62fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville if (inspect.ismodule(obj) or 63fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville (not inspect.isclass(obj) and obj.__dict__.has_key(attr_name))): 64fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville orig_obj = obj 65fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville orig_attr = getattr(obj, attr_name) 66fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville 67fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville else: 68fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville if not inspect.isclass(obj): 69fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville mro = list(inspect.getmro(obj.__class__)) 70fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville else: 71fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville mro = list(inspect.getmro(obj)) 72fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville 73fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville mro.reverse() 74fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville 75fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville orig_attr = None 76fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville 77fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville for cls in mro: 78fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville try: 79fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville orig_obj = cls 80fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville orig_attr = getattr(obj, attr_name) 81fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville except AttributeError: 82fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville continue 83fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville 84fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville if orig_attr is None: 85fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville raise AttributeError("Attribute not found.") 86fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville 87fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville # Calling getattr() on a staticmethod transforms it to a 'normal' function. 88fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville # We need to ensure that we put it back as a staticmethod. 89fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville old_attribute = obj.__dict__.get(attr_name) 90fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville if old_attribute is not None and isinstance(old_attribute, staticmethod): 91fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville orig_attr = staticmethod(orig_attr) 92fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville 93fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville self.stubs.append((orig_obj, attr_name, orig_attr)) 94fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville setattr(orig_obj, attr_name, new_attr) 95fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville 96fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville def SmartUnsetAll(self): 97fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville """Reverses all the SmartSet() calls, restoring things to their original 98fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville definition. Its okay to call SmartUnsetAll() repeatedly, as later calls 99fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville have no effect if no SmartSet() calls have been made. 100fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville 101fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville """ 102fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville self.stubs.reverse() 103fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville 104fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville for args in self.stubs: 105fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville setattr(*args) 106fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville 107fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville self.stubs = [] 108fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville 109fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville def Set(self, parent, child_name, new_child): 110fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville """Replace child_name's old definition with new_child, in the context 111fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville of the given parent. The parent could be a module when the child is a 112fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville function at module scope. Or the parent could be a class when a class' 113fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville method is being replaced. The named child is set to new_child, while 114fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville the prior definition is saved away for later, when UnsetAll() is called. 115fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville 116fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville This method supports the case where child_name is a staticmethod or a 117fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville classmethod of parent. 118fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville """ 119fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville old_child = getattr(parent, child_name) 120fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville 121fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville old_attribute = parent.__dict__.get(child_name) 122fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville if old_attribute is not None and isinstance(old_attribute, staticmethod): 123fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville old_child = staticmethod(old_child) 124fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville 125fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville self.cache.append((parent, old_child, child_name)) 126fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville setattr(parent, child_name, new_child) 127fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville 128fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville def UnsetAll(self): 129fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville """Reverses all the Set() calls, restoring things to their original 130fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville definition. Its okay to call UnsetAll() repeatedly, as later calls have 131fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville no effect if no Set() calls have been made. 132fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville 133fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville """ 134fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville # Undo calls to Set() in reverse order, in case Set() was called on the 135fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville # same arguments repeatedly (want the original call to be last one undone) 136fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville self.cache.reverse() 137fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville 138fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville for (parent, old_child, child_name) in self.cache: 139fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville setattr(parent, child_name, old_child) 140fbaaef999ba563838ebd00874ed8a1c01fbf286dWink Saville self.cache = [] 141