Track step nesting in StreamEngine.

recipe_engine/simulation_test.py

 # Copyright 2014 The LUCI Authors. All rights reserved.
 # Use of this source code is governed under the Apache License, Version 2.0
 # that can be found in the LICENSE file.
 
 """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.
 # We do this instead of passing because the threading system of 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
   from . import run
   from . import step_runner
   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'])
     api = loader.create_recipe_api(recipe_script.LOADED_DEPS, engine, test_data)
     result = engine.run(recipe_script, api)
 
     # Don't include tracebacks in expectations because they are too sensitive to
     # change.
     result.result.pop('traceback', None)
 
@@ skipping 71 matching lines @@
                   'TESTING_SLAVENAME']:
     if env_var in os.environ:
       logging.warn("Ignoring %s environment variable." % env_var)
       os.environ.pop(env_var)
 
   global _UNIVERSE
   _UNIVERSE = universe
 
   expect_tests.main('recipe_simulation_test', GenerateTests,
                     cover_omit=cover_omit(), args=args)