Track step nesting in StreamEngine.

recipe_engine/step_runner.py

 # Copyright 2016 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.
 
 import StringIO
 import collections
 import contextlib
 import datetime
 import json
 import os
@@ skipping 79 matching lines @@
     """Constructs an OpenStep object which can be used to actually run a step.
 
     step_dict parameters:
       name: name of the step, will appear in buildbots waterfall
       cmd: command to run, list of one or more strings
       cwd: absolute path to working directory for the command
       env: dict with overrides for environment variables
       allow_subannotations: if True, lets the step emit its own annotations
       trigger_specs: a list of trigger specifications, see also _trigger_builds.
       stdout: Path to a file to put step stdout into. If used, stdout won't
-              appear in annotator's stdout (and |allow_subannotations| is
-              ignored).
+          appear in annotator's stdout (and |allow_subannotations| is
+          ignored).
       stderr: Path to a file to put step stderr into. If used, stderr won't
-              appear in annotator's stderr.
+          appear in annotator's stderr.
       stdin: Path to a file to read step stdin from.
+      step_nest_level: the step's nesting level.
 
     Returns an OpenStep object.
     """
     raise NotImplementedError()
 
   def run_recipe(self, universe, recipe, properties):
     """Run the recipe named |recipe|.
 
     Args:
       universe: The RecipeUniverse where the recipe lives.
@@ skipping 42 matching lines @@
 
   def __init__(self, stream_engine):
     self._stream_engine = stream_engine
 
   @property
   def stream_engine(self):
     return self._stream_engine
 
   def open_step(self, step_dict):
     allow_subannotations = step_dict.get('allow_subannotations', False)
+    nest_level = step_dict.pop('step_nest_level', 0)
+
     step_stream = self._stream_engine.new_step_stream(
         step_dict['name'],
-        allow_subannotations=allow_subannotations)
+        allow_subannotations=allow_subannotations,
+        nest_level=nest_level)
     if not step_dict.get('cmd'):
       class EmptyOpenStep(OpenStep):
         def run(inner):
           if 'trigger_specs' in step_dict:
             self._trigger_builds(step_stream, step_dict['trigger_specs'])
           return types.StepData(step_dict, 0)
 
         def finalize(inner):
           step_stream.close()
 
         @property
         def stream(inner):
           return step_stream
 
       return EmptyOpenStep()
 
     step_dict, placeholders = render_step(
         step_dict, recipe_test_api.DisabledTestData())
     cmd = map(str, step_dict['cmd'])
     step_env = _merge_envs(os.environ, step_dict.get('env', {}))
-    if 'nest_level' in step_dict:
-      step_stream.step_nest_level(step_dict['nest_level'])
     self._print_step(step_stream, step_dict, step_env)
 
     class ReturnOpenStep(OpenStep):
       def run(inner):
         try:
           # Open file handles for IO redirection based on file names in
           # step_dict.
           handles = {
             'stdout': step_stream,
             'stderr': step_stream,
@@ skipping 190 matching lines @@
     return change
 
 
 class SimulationStepRunner(StepRunner):
   """Pretends to run steps, instead recording what would have been run.
 
   This is the main workhorse of recipes.py simulation_test.  Returns the log of
   steps that would have been run in steps_ran.  Uses test_data to mock return
   values.
   """
-  def __init__(self, stream_engine, test_data):
+  def __init__(self, stream_engine, test_data, annotator):
     self._test_data = test_data
     self._stream_engine = stream_engine
+    self._annotator = annotator
     self._step_history = collections.OrderedDict()
 
   @property
   def stream_engine(self):
     return self._stream_engine
 
   def open_step(self, step_dict):
     # We modify step_dict.  In particular, we add ~followup_annotations during
     # finalize, and depend on that side-effect being carried into what we
     # added to self._step_history, earlier.  So copy it here so at least we
     # keep the modifications local.
     step_dict = dict(step_dict)
+    nest_level = step_dict.pop('step_nest_level', 0)
 
     test_data_fn = step_dict.pop('step_test_data', recipe_test_api.StepTestData)
     step_test = self._test_data.pop_step_test_data(
         step_dict['name'], test_data_fn)
     step_dict, placeholders = render_step(step_dict, step_test)
-    outstream = StringIO.StringIO()
 
     # Layer the simulation step on top of the given stream engine.
-    step_stream = stream.ProductStreamEngine.StepStream(
-        self._stream_engine.new_step_stream(step_dict['name']),
-        stream.BareAnnotationStepStream(outstream))
+    step_stream = self._stream_engine.new_step_stream(
+        step_dict['name'],
+        nest_level=nest_level)
 
     class ReturnOpenStep(OpenStep):
       def run(inner):
         timeout = step_dict.get('timeout')
         if (timeout and step_test.times_out_after and
             step_test.times_out_after > timeout):
           raise recipe_api.StepTimeout(step_dict['name'], timeout)
 
         self._step_history[step_dict['name']] = step_dict
         return construct_step_result(step_dict, step_test.retcode, placeholders)
 
       def finalize(inner):
         # note that '~' sorts after 'z' so that this will be last on each
         # step. also use _step to get access to the mutable step
         # dictionary.
-        lines = filter(None, outstream.getvalue()).splitlines()
+        buf = self._annotator.step_buffer(step_dict['name'])
+        lines = filter(None, buf.getvalue()).splitlines()
         lines = [stream.encode_str(x) for x in lines]
         if lines:
           # This magically floats into step_history, which we have already
           # added step_dict to.
           step_dict['~followup_annotations'] = lines
+        step_stream.close()
 
       @property
       def stream(inner):
         return step_stream
 
     return ReturnOpenStep()
 
   def run_recipe(self, universe, recipe, properties):
     return self._test_data.depend_on_data.pop(types.freeze((recipe, properties),))
 
@@ skipping 197 matching lines @@
   supplied command, and only uses the |env| kwarg for modifying the environment
   of the child process.
   """
   saved_path = os.environ['PATH']
   try:
     if path is not None:
       os.environ['PATH'] = path
     yield
   finally:
     os.environ['PATH'] = saved_path