Refactor the way that TestApi works so that it is actually useful.

scripts/slave/recipe_test_api.py

Index: scripts/slave/recipe_test_api.py
diff --git a/scripts/slave/recipe_test_api.py b/scripts/slave/recipe_test_api.py
new file mode 100644
index 0000000000000000000000000000000000000000..2c24cd2abca4cba5313fe72ad75f67a0388d3558
--- /dev/null
+++ b/scripts/slave/recipe_test_api.py
@@ -0,0 +1,225 @@
+import collections
+import functools
+
+from .recipe_util import ModuleInjectionSite
+
+def combineify(name, dest, a, b):
+  """
+  Combines dictionary members in two objects into a third one using addition.
+
+  Args:
+    name - the name of the member
+    dest - the destination object
+    a - the first source object
+    b - the second source object
+  """
+  dest_dict = getattr(dest, name)
+  dest_dict.update(getattr(a, name))
+  for k, v in getattr(b, name).iteritems():
+    if k in dest_dict:
+      dest_dict[k] += v
+    else:
+      dest_dict[k] = v
+
+
+class PlaceholderTestData(object):
+  def __init__(self, data=None):
+    self.data = data
+    self.enabled = True
+
+  def __repr__(self):
+    return "PlaceholderTestData(%s)" % self.data
+
+
+class StepTestData(object):
+  """
+  Mutable container for per-step test data.
+
+  This data is consumed while running the recipe (during
+  annotated_run.run_steps).
+  """
+  def __init__(self):
+    # { (module, placeholder) -> [data] }
+    self.placeholder_data = collections.defaultdict(list)
+    self._retcode = None
+    self.enabled = True
+
+  def __add__(self, other):
+    assert isinstance(other, StepTestData)
+    ret = StepTestData()
+
+    # We don't use combineify to merge placeholder_data here because
+    # simply concatenating placeholder_data's list value is not meaningful to
+    # consumers of this object.
+    # Producers of this object should use the append() method.
+    ret.placeholder_data.update(self.placeholder_data)
+    for k, v in other.placeholder_data.iteritems():
+      assert k not in ret.placeholder_data
+      ret.placeholder_data[k] = v
+
+    # pylint: disable=W0212
+    ret._retcode = self._retcode
+    if other._retcode is not None:
+      assert ret._retcode is None
+      ret._retcode = other._retcode
+    return ret
+
+  def append(self, other):
+    self._retcode = self._retcode or other._retcode  # pylint: disable=W0212
+    combineify('placeholder_data', self, self, other)
+    return self
+
+  def pop_placeholder(self, name_pieces):
+    l = self.placeholder_data[name_pieces]
+    if l:
+      return l.pop(0)
+    else:
+      return PlaceholderTestData()
+
+  @property
+  def retcode(self):  # pylint: disable=E0202
+    return self._retcode or 0
+
+  @retcode.setter
+  def retcode(self, value):  # pylint: disable=E0202
+    self._retcode = value
+
+  def __repr__(self):
+    return "StepTestData(%s)" % str({
+      'placeholder_data': dict(self.placeholder_data.iteritems()),
+      'retcode': self._retcode,
+    })
+
+
+class ModuleTestData(dict):
+  """
+  Mutable container for test data for a specific module.
+
+  This test data is consumed at module load time (i.e. when CreateRecipeApi
+  runs).
+  """
+  def __init__(self):
+    super(ModuleTestData, self).__init__()
+    self.enabled = True
+
+  def __add__(self, other):
+    assert isinstance(other, ModuleTestData)
+    ret = ModuleTestData()
+    ret.update(self)
+    ret.update(other)
+    return ret
+
+  def __repr__(self):
+    return "ModuleTestData(%s)" % super(ModuleTestData, self).__repr__()
+
+
+class TestData(object):
+  def __init__(self, name=None):
+    self.name = name
+    self.properties = {}  # key -> val
+    self.mod_data = collections.defaultdict(ModuleTestData)
+    self.step_data = collections.defaultdict(StepTestData)
+    self.enabled = True
+
+  def __add__(self, other):
+    assert isinstance(other, TestData)
+    ret = TestData(self.name or other.name)
+
+    ret.properties.update(self.properties)
+    ret.properties.update(other.properties)
+
+    combineify('mod_data', ret, self, other)
+    combineify('step_data', ret, self, other)
+
+    return ret
+
+  def empty(self):
+    return not self.step_data
+
+  def __repr__(self):
+    return "TestData(%s)" % str({
+      'name': self.name,
+      'properties': self.properties,
+      'mod_data': dict(self.mod_data.iteritems()),
+      'step_data': dict(self.step_data.iteritems()),
+    })
+
+
+class DisabledTestData(object):
+  def __init__(self):
+    self.enabled = False
+
+  def __getattr__(self, name):
+    return self
+
+  def pop_placeholder(self, _name_pieces):
+    return self
+
+
+def static_wraps(func):
+  wrapped_fn = func
+  if isinstance(func, staticmethod):
+    wrapped_fn = func.__func__
+  return functools.wraps(wrapped_fn)
+
+
+def static_call(obj, func, *args, **kwargs):
+  if isinstance(func, staticmethod):
+    return func.__get__(obj)(*args, **kwargs)
+  else:
+    return func(obj, *args, **kwargs)
+
+
+def mod_test_data(func):
+  @static_wraps(func)
+  def inner(self, *args, **kwargs):
+    assert isinstance(self, RecipeTestApi)
+    mod_name = self._module.NAME  # pylint: disable=W0212
+    ret = TestData(None)
+    data = static_call(self, func, *args, **kwargs)
+    ret.mod_data[mod_name][inner.__name__] = data
+    return ret
+  return inner
+
+
+def placeholder_step_data(func):
+  @static_wraps(func)
+  def inner(self, *args, **kwargs):
+    assert isinstance(self, RecipeTestApi)
+    mod_name = self._module.NAME  # pylint: disable=W0212
+    ret = StepTestData()
+    placeholder_data, retcode = static_call(self, func, *args, **kwargs)
+    ret.placeholder_data[(mod_name, inner.__name__)].append(
+      PlaceholderTestData(placeholder_data))
+    ret.retcode = retcode
+    return ret
+  return inner
+
+
+class RecipeTestApi(object):
+  def __init__(self, module=None, test_data=DisabledTestData()):
+    """Note: Injected dependencies are NOT available in __init__()."""
+    # If we're the 'root' api, inject directly into 'self'.
+    # Otherwise inject into 'self.m'
+    self.m = self if module is None else ModuleInjectionSite()
+    self._module = module
+
+    assert isinstance(test_data, (ModuleTestData, DisabledTestData))
+    self._test_data = test_data
+
+  @staticmethod
+  def Test(name):
+    return TestData(name)
+
+  @staticmethod
+  def Properties(**kwargs):
+    ret = TestData(None)
+    ret.properties.update(kwargs)
+    return ret
+
+  @staticmethod
+  def StepData(name, *data):
+    assert all(isinstance(d, StepTestData) for d in data)
+    ret = TestData(None)
+    ret.step_data[name] = reduce(sum, data)
+    return ret