-
-
Save fbatroni/d5ba6bc9f05e985656063b035f4b0f19 to your computer and use it in GitHub Desktop.
For testing, use the mock module to wrap non-importable code. Also provide hooks for putting back in real functions/classes where needed.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import importlib | |
from mock import MagicMock, patch | |
def example(): | |
with patch_modules(): | |
import something.not.importable | |
print("yay! it worked!") | |
def patch_modules(): | |
sm = mock_import('not.importable', 'something') | |
# .. can add more to `sm` dict here.. | |
return patch.dict('sys.modules', sm) | |
# This won't work out of the box. It would work if | |
# you modified 'actually.importable' to be an importable | |
# module and modified 'Thing' to be the name of some member | |
# of that module. | |
def patch_modules2(): | |
sm = mock_import('not.importable', 'something', assign={ | |
'not.importable:SomeClass': 'actually.importable:Thing' | |
}) | |
return patch.dict('sys.modules', sm) | |
def mock_import(name, package, assign=None): | |
"""Set up mocking for imports of "<package>.<name>". | |
Optionally, assign real (importable) modules and/or module members | |
to the mocked ones (see `assign`). | |
Args: | |
name (str): Module path to mock up | |
package (str): Package name to mock up | |
assign (dict): Mapping where key is mock module and | |
value is real module path. Optionally, a ':<Class>' can | |
be appended to both key and value to assign classes (or really | |
any other module member) in the real module to the mock. | |
If None, nothing is done. | |
Returns: | |
(dict) New values for sys.modules dict to be patched. | |
If the given `package` imports without error, | |
nothing is done and return value is None. | |
""" | |
name, package = name.strip(), package.strip() | |
assert name, 'module must be non-empty' | |
assert package, 'package must be non-empty' | |
mlist = name.split('.') | |
try: | |
importlib.import_module(package, mlist[0]) | |
return None | |
except ImportError: | |
pass | |
pkg = MagicMock() | |
# add module path to mock | |
modpath, curmod = package, pkg | |
sys_modules = {modpath: curmod} | |
for m in mlist: | |
modpath = modpath + '.' + m | |
curmod = getattr(curmod, m) | |
sys_modules[modpath] = curmod | |
# assign real modules/classes to mock ones | |
if assign: | |
for mock_imp, real_imp in assign.items(): | |
# extract class, if present | |
if ':' in mock_imp: | |
mock_imp, mock_class = mock_imp.split(':') | |
real_imp, real_class = real_imp.split(':') | |
else: | |
mock_class, real_class = None, None | |
# import real module | |
real_pkg = real_imp.split('.')[0] | |
real_modules = '.' + real_imp[real_imp.index('.') + 1:] | |
real_mod = importlib.import_module(real_modules, real_pkg) | |
# assign real module, or class, to mock one | |
mock_mod, mock_imp_list = pkg, mock_imp.split('.') | |
if mock_class: | |
# set mock.module.path.Class = real.module.path.Class | |
for m in mock_imp_list: | |
mock_mod = getattr(mock_mod, m) | |
real_classobj = getattr(real_mod, real_class) | |
setattr(mock_mod, mock_class, real_classobj) | |
else: | |
# set mock.module.path = real.module.path | |
for m in mock_imp_list[:-1]: | |
mock_mod = getattr(mock_mod, m) | |
setattr(mock_mod, mock_imp_list[-1], real_mod) | |
return sys_modules |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment