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 81091a0a362f8a19f24cdd09110c20df6e2f3613..87c81ff8a0618d9e807033e40cd96390dd860bb3 100644
--- a/recipe_engine/simulation_test.py
+++ b/recipe_engine/simulation_test.py
@@ -4,6 +4,8 @@
 
 """Provides simulator test coverage for individual recipes."""
 
+import StringIO
+import contextlib
 import json
 import logging
 import os
@@ -11,6 +13,7 @@ import re
 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.
@@ -55,6 +58,24 @@ def RenderExpectation(test_data, raw_expectations):
   return expect_tests.Result(raw_expectations)
 
 
+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):
   """Actually runs the recipe given the GenTests-supplied test_data."""
   from . import config_types
@@ -64,8 +85,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'])