Preface
Let me start with some background. Our team’s service needs Python scripts to implement integration tests with relatively complex business logic, so we maintain a repository to manage this test code. This test suite is heavily coupled to an extremely complex internal pytest framework at Amazon. The framework’s role is to create and clean up AWS resources needed during a test run (such as EC2 instances, VPCs, etc.). However, for our use case it not only fails to solve any real problems, but instead introduces a lot of hassle. Still, because it exists to support the KPI of the team building the framework, it is an internal political requirement and we have no choice but to use it.
This internal framework seriously drags down those of us under GA pressure: since it manages resources, every test run must go through a series of complex steps before it actually starts executing the test code. This process can take from tens of seconds to over ten minutes, depending on whether the resources in the AWS account used for the test run are already ready. What makes the experience even worse is that Amazon’s repositories standardize on a Java-based build system called brazil. When building Python, it also performs some standard artifact packaging steps to produce a runtime package that can be deployed directly, so during this process it copies all .py source files into a new directory. As a result, if you manually execute certain test scripts in the repository, brazil will invoke them from the packaged environment rather than running the Python scripts directly from the repo.
We all know that writing tests in a dynamic language like Python can be very concise, but a major pain point is that: with only a small amount of context, it’s hard to figure out what type an object really is and what attributes and methods it has (this is especially bad for AWS Python clients). So you often end up tweaking things repeatedly. Taken together, these factors force you to endure an extremely inefficient feedback loop when writing integration tests: a one-line change in a test requires two minutes of build plus anywhere from a few minutes to over ten minutes of test framework initialization. It’s frustrating to the extreme. So, as an engineer who believes efficiency is life, I couldn’t help thinking—is there a way to pause the program when a test case doesn’t behave as expected, manually edit and reload the code, and then re-run from that point?
From a technical perspective, thanks to Python’s dynamic nature and pytest’s plugin mechanism, we can absolutely stop execution when a test case behaves unexpectedly, manually make adjustments, then directly reload the failing test code and re-run that case. And as long as we modify the search paths in sys.path correctly, we can make the Python runtime ignore the sources inside build artifacts and instead load the code we just modified. These two capabilities can be implemented with a pytest plugin, greatly improving development efficiency.
fix import path
First we need to update the import path. The basic idea is to take the absolute path of the currently running script (i.e., the build artifacts directory), then transform it into the actual path of the code repo according to convention. This way, the next time we try to import a module, Python will read the updated test case we modified:
def _redirect_import_search_path_to_brazil_source():
global _path_redirected
if _path_redirected:
return
def _redirect_path(path):
"""
redirect build artifacts path like `/local/home/<username>/workplace/MercuryControlPlaneIntegrationTestsEnv/build/MercuryControlPlaneIntegrationTests/MercuryControlPlaneIntegrationTests-1.0/AL2_x86_64/DEV.STD.PTHREAD/build/test-integ/MercuryAPI`
to its corresponding source dir for module import.
"""
basedir, file = os.path.split(os.path.abspath(__file__))
repo_root = os.path.abspath(os.path.join(basedir, ".."))
repo_name = os.path.basename(repo_root)
if repo_name in path:
path_elements = path.split(os.sep)
for i in range(len(path_elements) - 1, 0, -1):
new_path = os.path.join(repo_root, os.sep.join(path_elements[i:]))
if os.path.exists(new_path) and os.path.isdir(new_path) and new_path != path:
return new_path
return path
else:
return path
sys.path[:] = [
_redirect_path(p) for p in sys.path
]
_path_redirected = True
Pytest Hook
From ChatGPT:
In
pytest, the hook mechanism is a plugin system that allows users to extend and modifypytest’s behavior by defining custom hook functions.pytesttriggers these hooks at different stages of a test run, and you can insert custom logic by implementing or overriding them.Common hook functions:
pytest_configure(config): called during configuration initialization.pytest_collection_modifyitems(session, config, items): modifies collected test cases.pytest_runtest_setup(item): called before each test case runs.pytest_runtest_call(item): called when executing a test case.pytest_runtest_teardown(item): called after each test case runs.In pytest hook functions,
itemis apytest.Itemobject representing a test case. It encapsulates various information about the test case, such as its name, the function to execute, its path, the module it belongs to, parametrization, etc. You can accessitem’s attributes or methods to obtain or manipulate information related to the test case.In a
pytest.Item, bothfunctionandobjrelate to the Python function for the test case, but there is a subtle difference:
function: the wrapped test function object, typically used for execution; it may include pytest enhancements such aspytest.markhandling.obj: the original test function object, i.e., the Python function before pytest processing.
In this article, we use the hook pytest_runtest_protocol(item, nextitem). It controls the execution flow of a single test case, providing a very flexible mechanism that allows users to insert custom logic at each stage of test execution. It is commonly used to modify execution order, add condition checks, or satisfy other specific requirements. To fit our scenario, we can adjust the execution logic in this hook: when a test case fails, ask the user whether they want to modify the code and retry.
test code reload
Next we need to reload the code after making changes. The part below is the core logic for reloading a test case, and there are quite a few tricks in the overall process:
def load_method_from_classes(module, fname):
for name, obj in inspect.getmembers(module):
if inspect.isclass(obj) and hasattr(obj, fname):
return getattr(obj, fname)
raise RuntimeError("Error: [{}] not found, did you delete it?".format(fname))
def _reload_test_code(item):
# First, redirect the import path to the source repo rather than build artifacts
_redirect_import_search_path_to_brazil_source()
# Get the module where the currently failing test case resides; call it the old module
old_module = sys.modules[item.function.__module__]
globals_backup = {}
# Back up all global variables under the old module; the test execution needs them,
# but be careful to ignore those with special naming
for name, obj in inspect.getmembers(old_module):
if not inspect.isclass(obj) and \
not inspect.isfunction(obj) and \
not inspect.ismodule(obj) and \
not name.startswith("__") and \
not name.startswith("@"):
globals_backup[name] = obj
# Now the main part: reload the modified test code and get the new module
new_module = importlib.reload(old_module)
# Get the function name of the currently failing test case
name = item.function.__name__
# Note: a pytest test case can be either a function in a module or a method in a test class,
# so we need to handle them separately
if hasattr(new_module, name):
# If it's a function, it's easier: take the new test case function from the new module and override it
f = getattr(new_module, name)
item.obj = f
else:
# But if it's a class method, we have to traverse all classes in the new module and pick out the method with this name
f = load_method_from_classes(new_module, name)
# Then we need a somewhat tricky operation: override the `__code__` attribute,
# rather than directly overriding `obj`
item.obj.__func__.__code__ = f.__code__
# Finally, we need to copy global variables from the old module to the new module,
# because the logic in the test case may still depend on these data
for name, obj in globals_backup.items():
if _is_empty(getattr(new_module, name)):
setattr(new_module, name, obj)
print("\nReloaded test module from [{}]".format(new_module.__file__))
After fixing up various corner cases, the final complete implementation is here.
work with ipdb
ipdb is the IPython version of pdb. It provides enhancements such as syntax highlighting and auto-completion, making debugging more efficient. You can set a breakpoint by adding import ipdb; ipdb.set_trace() where needed. Combined with our pytest plugin, this debugger perfectly addresses the difficulty of not knowing dynamic types while writing test cases, because we can always bring up an IPython shell to inspect how to interact with various objects in the current scope, then modify the test logic until it runs through.