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