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..4c71435b10739d366e99483186875c3f1a9c55a1 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,23 @@ import expect_tests
 # doesn't know how to serialize it.
 _UNIVERSE = None
 
+
+class SimulationAnnotatorStreamEngine(stream.AnnotatorStreamEngine):

[dnj, 2016/08/15 17:33:53]

Now that the annotator stream engine tracks state, it's better to route the
simulation through the stream engine instead of having it have a separate weaker
version.

+
+  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=False,
+                      nested=False):
+    return self._create_step_stream(step_name, self.step_buffer(step_name),
+        allow_subannotations, nested)
+
+
 def RunRecipe(test_data):
   from . import config_types
   from . import loader
@@ -25,8 +45,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()

[dnj, 2016/08/15 17:33:53]

We pass the annotator directly to the stream engine because it needs to use it
to access the SimluationAnnotatorStreamEngine and, through it, the buffers for
each step.

+  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'])