Track step nesting in StreamEngine.

recipe_engine/simulation_test.py

Index: recipe_engine/simulation_test.py
diff --git a/recipe_engine/simulation_test.py b/recipe_engine/simulation_test.py
index 09f1d8c0aff008e098bf4f1f51c074f1bef083cb..4927739eb62f2055d2b92667ac10937e320823b8 100644
--- a/recipe_engine/simulation_test.py
+++ b/recipe_engine/simulation_test.py
@@ -4,12 +4,15 @@
 
 """Provides simulator test coverage for individual recipes."""
 
+import StringIO
+import contextlib
 import logging
 import re
 import os
 import sys
 
 from . import env
+from . import stream
 import expect_tests
 
 # This variable must be set in the dynamic scope of the functions in this file.
@@ -17,6 +20,25 @@ import expect_tests
 # doesn't know how to serialize it.
 _UNIVERSE = None
 
+
+class SimulationAnnotatorStreamEngine(stream.AnnotatorStreamEngine):
+
+  def __init__(self):
+    self._step_buffer_map = {}
+    super(SimulationAnnotatorStreamEngine, self).__init__(
+        self.step_buffer(None))
+
+  def step_buffer(self, step_name):
+    return self._step_buffer_map.setdefault(step_name, StringIO.StringIO())
+
+  def _new_step_stream(self, step_name, allow_subannotations, nest_level):
+    return self._create_step_stream(
+        step_name,
+        self.step_buffer(step_name),
+        allow_subannotations,
+        nest_level)
+
+
 def RunRecipe(test_data):
   from . import config_types
   from . import loader
@@ -25,8 +47,14 @@ def RunRecipe(test_data):
   from . import stream
 
   config_types.ResetTostringFns()
-  with stream.StreamEngineInvariants() as stream_engine:
-    step_runner = step_runner.SimulationStepRunner(stream_engine, test_data)
+
+  annotator = SimulationAnnotatorStreamEngine()
+  stream_engine = stream.ProductStreamEngine(
+      stream.StreamEngineInvariants(),
+      annotator)
+  with stream_engine:
+    step_runner = step_runner.SimulationStepRunner(stream_engine, test_data,
+                                                   annotator)
 
     engine = run.RecipeEngine(step_runner, test_data.properties, _UNIVERSE)
     recipe_script = _UNIVERSE.load_recipe(test_data.properties['recipe'])