Implementing Simulation Testing
Software engineers are familiar with the concept of testing. For those who don’t know, there are quite a few varieties of testing that you can encounter, the most common being unit, integration, and end to end (E2E) testing. In my career, I’ve come across more specialized areas of testing like mutation testing and simulation testing. Simulation testing is the subject of today’s article.
Recently, my main focus at work has been porting legacy server scaling software into a modern framework called Prefect. Prefect is a pure-Python orchestration engine similar to Apache Airflow. We use it to manage our server fleets and create scaling tools for our site reliability engineers (SREs). There were some major headaches that came with developing these tools:
- Testing a single workflow (even a small, 3-part orchestration of reads and writes) took at least 2 minutes. More complex workflows took ~30 minutes or more because of latency between data centers and the cloud where the orchestration was running. Even when no writes were being executed, the executions were dreadfully slow, which made testing painful.
- We couldn’t test for failures and how the orchestrator would handle our failure logic without hardcoding the failure into the run, or without modifying data processing in production. So we could test the happy path as much as we wanted, but in a workflow with sometimes hundreds of verifications and steps to go through, the branching paths were well into the thousands — testing at that volume was essentially impossible.
- If we were to run a full test scenario and previous unit/integration tests hadn’t caught an error or didn’t properly handle a rollback due to an environment failure, we might be stuck manually rolling back resources which can take hours of time — something I had been burned by more than once while developing these tools.
- E2E tests were difficult because we didn’t want to affect live resources and dry-running the workflow would still take hours for just a handful of scenarios.
- Regular tests that would mock out the entire workflow were obtuse and unwieldy. We didn’t want to write hundreds of lines of mocks and provide test data for every single scenario that we wanted to test — that would require an enormous number of failure injections. We also weren’t interested as much in testing the individual stages (which could be done relatively easily), but in the overall outcome of a workflow. If a run failed at step 8, was the proper cleanup executed? Was an SRE notified? Are rollbacks executing correctly such that all resources are unchanged if something fails? This was difficult to catch with unit tests.
As I was thinking of a way to address these problems, I remembered something that I learned about from the company TigerBeetle (of which I am a big fan) — they use a deterministic simulation engine to run their database engine across hundreds of years of data every single night. I learned more about simulation testing, and I thought, “This might be exactly (ish) what we need to better test our scaling tools.”
For testing systems, you can imagine risk and thoroughness as trade-offs — a real system cannot be tested exhaustively because it can cause a variety of issues which is expensive and risky; on the other hand, performing shallow tests is cheap but could leave risks unaddressed and cause issues when systems are already deployed. Simulations break that tradeoff by making the cost-per-test almost zero and eliminating the risk-per-test entirely.
It should be pointed out that simulation testing is not the same as chaos testing, though they sound similar. Chaos testing puts strain on infrastructure and ensures that it can survive various scenarios; simulation testing verifies that the workflow logic correctly handles failures and focuses on the orchestration layer.
First approach — Classic Dependency Injection
My first approach to enabling simulation testing for our existing code was to try a classic dependency injection rewrite. It was feasible and simple to do, only in the sense that the steps to enable dependency injection were not complex (inject the API client, provide a response for the task within the flow, etc.).
The dependency injection was triggered using a ContextVar and a context manager. The simulated branches were run if the ContextVar was set, bypassing any calls to a staging or production environment and generating an in-memory environment. The in-memory environment randomly generated all of the resources that the flow would typically interact with based on a provided seed value, which allowed me to replicate environments as many times as needed when I encountered an error as long as I recorded the seed that failed. Once generated, the environment would be frozen and then duplicated to a mutable copy for evaluation at the end of the run. The simulation engine would compare the two states at the end against a set of rules that were defined for each workflow. The evaluation looked for three "success" states:
- If it ran successfully (flow exited cleanly), did the actual simulated state match the expected state based on the input parameters?
- If it failed, did the actual simulated environment return to its original state?
- If it failed and the environment was not reverted correctly, were the correct instructions generated to guide an SRE in a manual recovery process?
This was the most important feature of this system. When testing complex workflows, it’s tempting to test just to make sure that a workflow runs to completion without crashing; however, the real test is to verify that the state of the system is either the same as you started (in the case of a rollback), or has implemented the correct changes based on the workflow and its parameters. Comparing the final state should never be treated as only a ceremony, because that comparison is the entire point of running the test. This kind of testing is valuable because it’s difficult to conduct these tests in a live system — a live system is constantly changing, whereas a simulated environment can be “frozen” and evaluated at any time.
The biggest downside to this was that it involved a massive amount of refactoring. Given that I didn’t own every workflow in our toolkit, I was hesitant to go about rewriting entire sections of our codebase, even if AI could assist with it. There was a massive amount of implicit context that had been built up over 20+ years of operation, and missing even a single detail could be catastrophic. I also wanted to make it easy for each individual contributor to this toolkit to drop in the simulation features without deeply modifying their existing work, and this approach did not accomplish that goal.
But, the results were promising — I could now run a single workflow over a million times in the same amount of time it once took to execute it once in dry-run, with no data moving between hosts.
Second approach — Python Decorators
Since the biggest issue was rewriting the existing code to route for dependency injection, I tried to skirt the issue by using decorators to intercept the logic execution. With the same ContextVar usage, I was able to decorate all of the tasks and flows from Prefect with another layer of logic that would route to a simulation layer if the ContextVar was present. If it wasn't, then it would perform the operation and interact with live resources.
Get Tate Wilks’s stories in your inbox
Join Medium for free to get updates from this writer.
This worked amazingly well in simulations and in local development. I was able to achieve the same results for simulation throughput and with much less code being written. However, when I deployed the solution to our Prefect cloud, the decorators interfered with the Prefect runtime detecting the flows and tasks. This would have broken everything in production because when Prefect loads the flows and tasks, it does a type check on the task to make sure it’s an instance of prefect.Task. The design that I introduced as the top-layer decorator wrapped the prefect.Task in a different class. This would cause every flow that implemented the simulation testing to crash if it reached production.
So the decorators were a winning idea because of the ease of implementation, but it would break production. At this point, I had invested a decent amount of time researching and prototyping this (AI was decent enough at this point that the actual implementation didn’t suck a lot of time up). I was starting to think I might just need to do long-running manual testing before official releases.
Third approach — Subclass the Prefect Decorators
At first I thought I was at a dead end, mostly because altering an SDK class felt illegal. Surely that isn’t an acceptable way to write software! But at the same time… if it can be done, why not do it? “Why not?” isn’t a satisfying answer, so I browsed around StackOverflow and read some articles and came to the conclusion that subclassing a stable, public class is not only pretty common, but also an acceptable implementation pattern. There are always concerns of breaking changes introduced by the vendor, but the likelihood of that happening for such a core piece of the Prefect API is very low.
So, I subclassed the task and flow decorators that Prefect provides. Within those classes I had logic to intercept the process if the ContextVar was present; if not, then the class would cede to the parent class with a call to super() to access the original method. This allowed me to inject dependencies without rewriting tons of code.
This solved the engine loading problem as well. By subclassing the prefect.Task class, the Prefect engine could still register my dual-mode simulation tasks as regular Prefect tasks.
So the overall implementation went from this:
from prefect.tasks import task
@task(name="module.my_task:execution", ...)
def execution(*args, **kwargs):
...To this:
# subclassed task (details omitted for brevity)
import prefect
from contextvars import ContextVar
from typing import Generic, ParamSpec, TypeVar
P = ParamSpec("P")
R = TypeVar("R")
RANDOM_STATE_CONTEXT: ContextVar = ContextVar("engine")
class Task(prefect.Task, Generic[P, R]):
def __call__(self, ...):
if (sse := RANDOM_STATE_CONTEXT.get(None)) is None:
return super().__call__(...)
return self._simulated_call(sse, ...)
def task(...) -> Task:
# this is effectively the same as the library task function
# but it has the additional argument 'simulation_config'
...# function file with task decorator applied
from my_module.tasks import task
@task(
name="module.my_task:execution",
simulation_config=SimulationConfig(...),
)
def execution(*args, **kwargs):
...This felt almost like a drop-in replacement that could then easily be configured with a referenced simulation function that would supersede the logic defined in the original function. So instead of a db write being performed in this execution, the simulated in-memory state was modified by injecting the simulation function, which would look like this:
def sim_execution(sse, *args, **kwargs):
sse.write_option(...)Where sse is the Simulation State Engine that is accessed through the ContextVar and contains the in-memory simulated environment.
There was the downside that the simulated logic still had to be written, but that tradeoff was much more palatable than rewriting all of the original code and still needing to write the simulation logic (writing simulation logic unfortunately was not something that I could get around in any of the three scenarios that I experimented with).
Outcome
The outcome of this implementation allowed me to find obscure bugs in the branching logic by running millions of simulations within minutes instead of hours, days, weeks, or even months of testing, effectively removing the cost constraint of thorough testing. I was also able to run all these scenarios safely without affecting any live resources in stage or production, which removed any risk constraint. I could also re-run the same simulations with a seed value. If seed 9734208140238 failed, I could re-run it by specifying that seed and step through line-by-line to find where the sneakiest of bugs had snuck into my code.
Later on I implemented levers to tune certain probabilities higher or lower (pod crashes, environmental failures, cancellations, etc.) so that I could test antagonistic scenarios as well as happy paths where all things functioned as they should. To accomplish this, I added failure modes as factors in the simulation engine (a regular failed state, a crash, or cancellation). Each one has an amplification factor of 1 (i.e., no change to its base probability), and when loading the simulation I could toggle those values to go to zero for a perfect run, or increase them dramatically to ensure the execution failed.
An example in code:
def __call__(self, ...):
...
if (p := sse.rng.random()) < (self.__config.crash_probability * sse.crash_failure_factor):
# handle simulated crashed call
...sse is the Simulation State Engine that is stored in the Context variable. The __config is a task-level configuration for simulation that determines the probability of a crash (defaulted to 0.03 in this example). rng is a seeded instance of Random, and the crash_failure_factor is the amplifying factor.
I also added options to increase the size of the environments so that I could simulate jobs of varying blast radii e.g. moving two accounts from one host to another versus distributing a thousand accounts across 4 different destinations. I could now run rare and dangerous scenarios that I would never attempt in a live environment.
There are some drawbacks to this approach, the main two being drift and maintenance. I don’t consider the drift between the simulation and the actual environment to be a strong argument against this approach, because the product I’m working on has been running for over two decades. It is incredibly stable and breaking changes are essentially unheard of. The maintenance cost is a much more compelling argument, because Prefect is relatively new and has introduced breaking changes that have given me and my team headaches in the past. However, the Task interface is very stable, and the maintenance cost is one that I will happily pay for this level of testing.
Despite the drawbacks and the future required maintenance, the simulation logic provides a way to thoroughly test the control flow of complex, multi-step chains of logic in a way that dramatically reduces the execution time, completely eliminates the risk of altering resources when encountering failures, and allows me to cover more scenarios at an unprecedented scale. I think that the price of maintenance and possible drift is well worth the speed and insight that this harness provides.
