Skip to content

Instantly share code, notes, and snippets.

@dangunter
Last active November 8, 2022 10:12
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 1 You must be signed in to fork a gist
  • Save dangunter/7c00171ee0d5dbf0689582e85bdbef53 to your computer and use it in GitHub Desktop.
Save dangunter/7c00171ee0d5dbf0689582e85bdbef53 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.
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