Bug 1333255: and then there was only one Task class; r=jonasfj
authorDustin J. Mitchell <dustin@mozilla.com>
Thu, 09 Mar 2017 23:14:30 -0500
changeset 398026 438ada631812058d59fca89218d7a472570f6711
parent 398025 90bee55f9797d5be2cdfafa56f97b320eb977088
child 398027 4fa17de0a52fd1ee7434cab1e610d78a6c85dc98
push id1490
push usermtabara@mozilla.com
push dateMon, 31 Jul 2017 14:08:16 +0000
treeherdermozilla-release@70e32e6bf15e [default view] [failures only]
perfherder[talos] [build metrics] [platform microbench] (compared to previous push)
reviewersjonasfj
bugs1333255
milestone55.0a1
first release with
nightly linux32
nightly linux64
nightly mac
nightly win32
nightly win64
last release without
nightly linux32
nightly linux64
nightly mac
nightly win32
nightly win64
Bug 1333255: and then there was only one Task class; r=jonasfj Note that the to_json method prefers the taskgraph's dependencies information (edges) to that from the task.dependencies entries. At a few points in task-graph generation, these values differ, although that is expected (for example, the full task set contains no edges, but that information is still in task.dependencies). Unifying that representation leads to some difficulty with task transforms that reach into the dependency tree (beetmover), so the different representations are left as-is. MozReview-Commit-ID: GeW8HNwFA9Z
taskcluster/ci/docker-image/kind.yml
taskcluster/docs/taskgraph.rst
taskcluster/taskgraph/task/base.py
taskcluster/taskgraph/task/docker_image.py
taskcluster/taskgraph/task/test.py
taskcluster/taskgraph/task/transform.py
taskcluster/taskgraph/taskgraph.py
taskcluster/taskgraph/test/test_create.py
taskcluster/taskgraph/test/test_decision.py
taskcluster/taskgraph/test/test_optimize.py
taskcluster/taskgraph/test/test_target_tasks.py
taskcluster/taskgraph/test/test_taskgraph.py
taskcluster/taskgraph/test/test_try_option_syntax.py
taskcluster/taskgraph/test/util.py
taskcluster/taskgraph/transforms/docker_image.py
--- a/taskcluster/ci/docker-image/kind.yml
+++ b/taskcluster/ci/docker-image/kind.yml
@@ -1,13 +1,13 @@
 # This Source Code Form is subject to the terms of the Mozilla Public
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
-loader: taskgraph.task.docker_image:load_tasks
+loader: taskgraph.task.transform:load_tasks
 
 transforms:
   - taskgraph.transforms.docker_image:transforms
   - taskgraph.transforms.task:transforms
 
 # make a task for each docker-image we might want.  For the moment, since we
 # write artifacts for each, these are whitelisted, but ideally that will change
 # (to use subdirectory clones of the proper directory), at which point we can
--- a/taskcluster/docs/taskgraph.rst
+++ b/taskcluster/docs/taskgraph.rst
@@ -203,36 +203,38 @@ Task graphs -- both the graph artifacts 
 output by the ``--json`` option to the ``mach taskgraph`` commands -- are JSON
 objects, keyed by label, or for optimized task graphs, by taskId.  For
 convenience, the decision task also writes out ``label-to-taskid.json``
 containing a mapping from label to taskId.  Each task in the graph is
 represented as a JSON object.
 
 Each task has the following properties:
 
+``kind``
+   The name of this task's kind
+
 ``task_id``
    The task's taskId (only for optimized task graphs)
 
 ``label``
    The task's label
 
 ``attributes``
    The task's attributes
 
 ``dependencies``
    The task's in-graph dependencies, represented as an object mapping
    dependency name to label (or to taskId for optimized task graphs)
 
+``optimizations``
+   The optimizations to be applied to this task
+
 ``task``
    The task's TaskCluster task definition.
 
-``kind_implementation``
-   The module and the class name which was used to implement this particular task.
-   It is always of the form ``<module-path>:<object-path>``
-
 The results from each command are in the same format, but with some differences
 in the content:
 
 * The ``tasks`` and ``target`` subcommands both return graphs with no edges.
   That is, just collections of tasks without any dependencies indicated.
 
 * The ``optimized`` subcommand returns tasks that have been assigned taskIds.
   The dependencies array, too, contains taskIds instead of labels, with
--- a/taskcluster/taskgraph/task/base.py
+++ b/taskcluster/taskgraph/task/base.py
@@ -1,16 +1,14 @@
 # This Source Code Form is subject to the terms of the Mozilla Public
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
 from __future__ import absolute_import, print_function, unicode_literals
 
-import abc
-
 
 class Task(object):
     """
     Representation of a task in a TaskGraph.  Each Task has, at creation:
 
     - kind: the name of the task kind
     - label; the label for this task
     - attributes: a dictionary of attributes for this task (used for filtering)
@@ -19,23 +17,19 @@ class Task(object):
     - dependencies: tasks this one depends on, in the form {name: label}, for example
       {'build': 'build-linux64/opt', 'docker-image': 'build-docker-image-desktop-test'}
 
     And later, as the task-graph processing proceeds:
 
     - task_id -- TaskCluster taskId under which this task will be created
     - optimized -- true if this task need not be performed
 
-    A kind represents a collection of tasks that share common characteristics.
-    For example, all build jobs.  Each instance of a kind is intialized with a
-    path from which it draws its task configuration.  The instance is free to
-    store as much local state as it needs.
+    This class is just a convenience wraper for the data type and managing
+    display, comparison, serialization, etc. It has no functionality of its own.
     """
-    __metaclass__ = abc.ABCMeta
-
     def __init__(self, kind, label, attributes, task,
                  optimizations=None, dependencies=None):
         self.kind = kind
         self.label = label
         self.attributes = attributes
         self.task = task
 
         self.task_id = None
@@ -50,20 +44,43 @@ class Task(object):
         return self.kind == other.kind and \
             self.label == other.label and \
             self.attributes == other.attributes and \
             self.task == other.task and \
             self.task_id == other.task_id and \
             self.optimizations == other.optimizations and \
             self.dependencies == other.dependencies
 
+    def __repr__(self):
+        return ('Task({kind!r}, {label!r}, {attributes!r}, {task!r}, '
+                'optimizations={optimizations!r}, '
+                'dependencies={dependencies!r})'.format(**self.__dict__))
+
+    def to_json(self):
+        rv = {
+            'kind': self.kind,
+            'label': self.label,
+            'attributes': self.attributes,
+            'dependencies': self.dependencies,
+            'optimizations': self.optimizations,
+            'task': self.task,
+        }
+        if self.task_id:
+            rv['task_id'] = self.task_id
+        return rv
+
     @classmethod
     def from_json(cls, task_dict):
         """
         Given a data structure as produced by taskgraph.to_json, re-construct
         the original Task object.  This is used to "resume" the task-graph
         generation process, for example in Action tasks.
         """
-        return cls(
-            kind=task_dict['attributes']['kind'],
+        rv = cls(
+            kind=task_dict['kind'],
             label=task_dict['label'],
             attributes=task_dict['attributes'],
-            task=task_dict['task'])
+            task=task_dict['task'],
+            optimizations=task_dict['optimizations'],
+            dependencies=task_dict.get('dependencies'))
+        if 'task_id' in task_dict:
+            rv.task_id = task_dict['task_id']
+        return rv
deleted file mode 100644
--- a/taskcluster/taskgraph/task/docker_image.py
+++ /dev/null
@@ -1,43 +0,0 @@
-# This Source Code Form is subject to the terms of the Mozilla Public
-# License, v. 2.0. If a copy of the MPL was not distributed with this
-# file, You can obtain one at http://mozilla.org/MPL/2.0/.
-
-from __future__ import absolute_import, print_function, unicode_literals
-
-import logging
-
-from . import transform
-from taskgraph.transforms.base import TransformSequence, TransformConfig
-from taskgraph.util.python_path import find_object
-
-logger = logging.getLogger(__name__)
-
-
-def transform_inputs(inputs, kind, path, config, params, loaded_tasks):
-    """
-    Transform a sequence of inputs according to the transform configuration.
-    """
-    transforms = TransformSequence()
-    for xform_path in config['transforms']:
-        transform = find_object(xform_path)
-        transforms.add(transform)
-
-    # perform the transformations
-    trans_config = TransformConfig(kind, path, config, params)
-    tasks = [DockerImageTask(kind, t)
-             for t in transforms(trans_config, inputs)]
-    return tasks
-
-
-def load_tasks(kind, path, config, params, loaded_tasks):
-    return transform_inputs(
-        transform.get_inputs(kind, path, config, params, loaded_tasks),
-        kind, path, config, params, loaded_tasks)
-
-
-class DockerImageTask(transform.TransformTask):
-
-    @classmethod
-    def from_json(cls, task_dict):
-        docker_image_task = cls(kind='docker-image', task=task_dict)
-        return docker_image_task
--- a/taskcluster/taskgraph/task/test.py
+++ b/taskcluster/taskgraph/task/test.py
@@ -16,17 +16,17 @@ logger = logging.getLogger(__name__)
 def get_inputs(kind, path, config, params, loaded_tasks):
     """
     Generate tasks implementing Gecko tests.
     """
 
     # the kind on which this one depends
     if len(config.get('kind-dependencies', [])) != 1:
         raise Exception(
-            "TestTask kinds must have exactly one item in kind-dependencies")
+            "Test kinds must have exactly one item in kind-dependencies")
     dep_kind = config['kind-dependencies'][0]
 
     # get build tasks, keyed by build platform
     builds_by_platform = get_builds_by_platform(dep_kind, loaded_tasks)
 
     # get the test platforms for those build tasks
     test_platforms_cfg = load_yaml(path, 'test-platforms.yml')
     test_platforms = get_test_platforms(test_platforms_cfg, builds_by_platform)
--- a/taskcluster/taskgraph/task/transform.py
+++ b/taskcluster/taskgraph/task/transform.py
@@ -58,35 +58,22 @@ def transform_inputs(inputs, kind, path,
     """
     transforms = TransformSequence()
     for xform_path in config['transforms']:
         transform = find_object(xform_path)
         transforms.add(transform)
 
     # perform the transformations
     trans_config = TransformConfig(kind, path, config, params)
-    tasks = [TransformTask(kind, t) for t in transforms(trans_config, inputs)]
+    tasks = [base.Task(kind,
+                       label=task_dict['label'],
+                       attributes=task_dict['attributes'],
+                       task=task_dict['task'],
+                       optimizations=task_dict.get('optimizations'),
+                       dependencies=task_dict.get('dependencies'))
+             for task_dict in transforms(trans_config, inputs)]
     return tasks
 
 
 def load_tasks(kind, path, config, params, loaded_tasks):
     return transform_inputs(
         get_inputs(kind, path, config, params, loaded_tasks),
         kind, path, config, params, loaded_tasks)
-
-
-class TransformTask(base.Task):
-    """
-    Tasks of this class are generated by applying transformations to a sequence
-    of input entities.  By default, it gets those inputs from YAML data in the
-    kind directory, but subclasses may override `get_inputs` to produce them in
-    some other way.
-    """
-
-    def __init__(self, kind, task):
-        super(TransformTask, self).__init__(kind, task['label'],
-                                            task['attributes'], task['task'],
-                                            optimizations=task.get('optimizations'),
-                                            dependencies=task.get('dependencies'))
-
-    @classmethod
-    def from_json(cls, task_dict):
-        return cls(task_dict['attributes']['kind'], task_dict)
--- a/taskcluster/taskgraph/taskgraph.py
+++ b/taskcluster/taskgraph/taskgraph.py
@@ -1,51 +1,31 @@
 # This Source Code Form is subject to the terms of the Mozilla Public
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
 from __future__ import absolute_import, print_function, unicode_literals
 
 from .graph import Graph
-from .util.python_path import find_object
+from .task.base import Task
 
 
 class TaskGraph(object):
     """
     Representation of a task graph.
 
     A task graph is a combination of a Graph and a dictionary of tasks indexed
     by label.  TaskGraph instances should be treated as immutable.
     """
 
     def __init__(self, tasks, graph):
         assert set(tasks) == graph.nodes
         self.tasks = tasks
         self.graph = graph
 
-    def to_json(self):
-        "Return a JSON-able object representing the task graph, as documented"
-        named_links_dict = self.graph.named_links_dict()
-        # this dictionary may be keyed by label or by taskid, so let's just call it 'key'
-        tasks = {}
-        for key in self.graph.visit_postorder():
-            task = self.tasks[key]
-            implementation = task.__class__.__module__ + ":" + task.__class__.__name__
-            task_json = {
-                'label': task.label,
-                'attributes': task.attributes,
-                'dependencies': named_links_dict.get(key, {}),
-                'task': task.task,
-                'kind_implementation': implementation
-            }
-            if task.task_id:
-                task_json['task_id'] = task.task_id
-            tasks[key] = task_json
-        return tasks
-
     def for_each_task(self, f, *args, **kwargs):
         for task_label in self.graph.visit_postorder():
             task = self.tasks[task_label]
             f(task, self, *args, **kwargs)
 
     def __getitem__(self, label):
         "Get a task by label"
         return self.tasks[label]
@@ -55,28 +35,35 @@ class TaskGraph(object):
         return self.tasks.itervalues()
 
     def __repr__(self):
         return "<TaskGraph graph={!r} tasks={!r}>".format(self.graph, self.tasks)
 
     def __eq__(self, other):
         return self.tasks == other.tasks and self.graph == other.graph
 
+    def to_json(self):
+        "Return a JSON-able object representing the task graph, as documented"
+        named_links_dict = self.graph.named_links_dict()
+        # this dictionary may be keyed by label or by taskid, so let's just call it 'key'
+        tasks = {}
+        for key in self.graph.visit_postorder():
+            tasks[key] = self.tasks[key].to_json()
+            # overwrite dependencies with the information in the taskgraph's edges.
+            tasks[key]['dependencies'] = named_links_dict.get(key, {})
+        return tasks
+
     @classmethod
     def from_json(cls, tasks_dict):
         """
         This code is used to generate the a TaskGraph using a dictionary
         which is representative of the TaskGraph.
         """
         tasks = {}
         edges = set()
         for key, value in tasks_dict.iteritems():
-            # We get the implementation from JSON
-            implementation = value['kind_implementation']
-            # Loading the module and creating a Task from a dictionary
-            task_kind = find_object(implementation)
-            tasks[key] = task_kind.from_json(value)
+            tasks[key] = Task.from_json(value)
             if 'task_id' in value:
                 tasks[key].task_id = value['task_id']
             for depname, dep in value['dependencies'].iteritems():
                 edges.add((key, dep, depname))
         task_graph = cls(tasks, Graph(set(tasks), edges))
         return tasks, task_graph
--- a/taskcluster/taskgraph/test/test_create.py
+++ b/taskcluster/taskgraph/test/test_create.py
@@ -5,17 +5,17 @@
 from __future__ import absolute_import, print_function, unicode_literals
 
 import unittest
 import os
 
 from .. import create
 from ..graph import Graph
 from ..taskgraph import TaskGraph
-from .util import TestTask
+from ..task.base import Task
 
 from mozunit import main
 
 
 class TestCreate(unittest.TestCase):
 
     def setUp(self):
         self.old_task_id = os.environ.get('TASK_ID')
@@ -32,18 +32,18 @@ class TestCreate(unittest.TestCase):
         elif 'TASK_ID' in os.environ:
             del os.environ['TASK_ID']
 
     def fake_create_task(self, session, task_id, label, task_def):
         self.created_tasks[task_id] = task_def
 
     def test_create_tasks(self):
         tasks = {
-            'tid-a': TestTask(label='a', task={'payload': 'hello world'}),
-            'tid-b': TestTask(label='b', task={'payload': 'hello world'}),
+            'tid-a': Task(kind='test', label='a', attributes={}, task={'payload': 'hello world'}),
+            'tid-b': Task(kind='test', label='b', attributes={}, task={'payload': 'hello world'}),
         }
         label_to_taskid = {'a': 'tid-a', 'b': 'tid-b'}
         graph = Graph(nodes={'tid-a', 'tid-b'}, edges={('tid-a', 'tid-b', 'edge')})
         taskgraph = TaskGraph(tasks, graph)
 
         create.create_tasks(taskgraph, label_to_taskid, {'level': '4'})
 
         for tid, task in self.created_tasks.iteritems():
@@ -55,17 +55,17 @@ class TestCreate(unittest.TestCase):
                     # Don't look for decisiontask here
                     continue
                 self.assertIn(depid, self.created_tasks)
 
     def test_create_task_without_dependencies(self):
         "a task with no dependencies depends on the decision task"
         os.environ['TASK_ID'] = 'decisiontask'
         tasks = {
-            'tid-a': TestTask(label='a', task={'payload': 'hello world'}),
+            'tid-a': Task(kind='test', label='a', attributes={}, task={'payload': 'hello world'}),
         }
         label_to_taskid = {'a': 'tid-a'}
         graph = Graph(nodes={'tid-a'}, edges=set())
         taskgraph = TaskGraph(tasks, graph)
 
         create.create_tasks(taskgraph, label_to_taskid, {'level': '4'})
 
         for tid, task in self.created_tasks.iteritems():
--- a/taskcluster/taskgraph/test/test_decision.py
+++ b/taskcluster/taskgraph/test/test_decision.py
@@ -7,51 +7,21 @@ from __future__ import absolute_import, 
 import os
 import json
 import yaml
 import shutil
 import unittest
 import tempfile
 
 from .. import decision
-from ..graph import Graph
-from ..taskgraph import TaskGraph
-from .util import TestTask
 from mozunit import main
 
 
 class TestDecision(unittest.TestCase):
 
-    def test_taskgraph_to_json(self):
-        tasks = {
-            'a': TestTask(label='a', attributes={'attr': 'a-task'}),
-            'b': TestTask(label='b', task={'task': 'def'}),
-        }
-        graph = Graph(nodes=set('ab'), edges={('a', 'b', 'edgelabel')})
-        taskgraph = TaskGraph(tasks, graph)
-
-        res = taskgraph.to_json()
-
-        self.assertEqual(res, {
-            'a': {
-                'label': 'a',
-                'attributes': {'attr': 'a-task', 'kind': 'test'},
-                'task': {},
-                'dependencies': {'edgelabel': 'b'},
-                'kind_implementation': 'taskgraph.test.util:TestTask',
-            },
-            'b': {
-                'label': 'b',
-                'attributes': {'kind': 'test'},
-                'task': {'task': 'def'},
-                'dependencies': {},
-                'kind_implementation': 'taskgraph.test.util:TestTask',
-            }
-        })
-
     def test_write_artifact_json(self):
         data = [{'some': 'data'}]
         tmpdir = tempfile.mkdtemp()
         try:
             decision.ARTIFACTS_DIR = os.path.join(tmpdir, "artifacts")
             decision.write_artifact("artifact.json", data)
             with open(os.path.join(decision.ARTIFACTS_DIR, "artifact.json")) as f:
                 self.assertEqual(json.load(f), data)
--- a/taskcluster/taskgraph/test/test_optimize.py
+++ b/taskcluster/taskgraph/test/test_optimize.py
@@ -5,17 +5,17 @@
 from __future__ import absolute_import, print_function, unicode_literals
 
 import unittest
 
 from ..optimize import optimize_task_graph, resolve_task_references, optimization
 from ..optimize import annotate_task_graph, get_subgraph
 from ..taskgraph import TaskGraph
 from .. import graph
-from .util import TestTask
+from ..task.base import Task
 
 
 class TestResolveTaskReferences(unittest.TestCase):
 
     def do(self, input, output):
         taskid_for_edge_name = {'edge%d' % n: 'tid%d' % n for n in range(1, 4)}
         self.assertEqual(resolve_task_references('subject', input, taskid_for_edge_name), output)
 
@@ -62,28 +62,28 @@ class TestOptimize(unittest.TestCase):
         # set up some simple optimization functions
         optimization('no-optimize')(lambda self, params: (False, None))
         optimization('optimize-away')(lambda self, params: (True, None))
         optimization('optimize-to-task')(lambda self, params, task: (True, task))
         optimization('false-with-taskid')(lambda self, params: (False, 'some-taskid'))
 
     def make_task(self, label, optimization=None, task_def=None, optimized=None, task_id=None):
         task_def = task_def or {'sample': 'task-def'}
-        task = TestTask(label=label, task=task_def)
+        task = Task(kind='test', label=label, attributes={}, task=task_def)
         task.optimized = optimized
         if optimization:
             task.optimizations = [optimization]
         else:
             task.optimizations = []
         task.task_id = task_id
         return task
 
     def make_graph(self, *tasks_and_edges):
-        tasks = {t.label: t for t in tasks_and_edges if isinstance(t, TestTask)}
-        edges = {e for e in tasks_and_edges if not isinstance(e, TestTask)}
+        tasks = {t.label: t for t in tasks_and_edges if isinstance(t, Task)}
+        edges = {e for e in tasks_and_edges if not isinstance(e, Task)}
         return TaskGraph(tasks, graph.Graph(set(tasks), edges))
 
     def assert_annotations(self, graph, **annotations):
         def repl(task_id):
             return 'SLUGID' if task_id and len(task_id) == 22 else task_id
         got_annotations = {
             t.label: (t.optimized, repl(t.task_id)) for t in graph.tasks.itervalues()
         }
--- a/taskcluster/taskgraph/test/test_target_tasks.py
+++ b/taskcluster/taskgraph/test/test_target_tasks.py
@@ -5,17 +5,17 @@
 from __future__ import absolute_import, print_function, unicode_literals
 
 import unittest
 
 from .. import target_tasks
 from .. import try_option_syntax
 from ..graph import Graph
 from ..taskgraph import TaskGraph
-from .util import TestTask
+from ..task.base import Task
 from mozunit import main
 
 
 class FakeTryOptionSyntax(object):
 
     def __init__(self, message, task_graph):
         self.trigger_tests = 0
         self.talos_trigger_tests = 0
@@ -29,18 +29,19 @@ class FakeTryOptionSyntax(object):
         return 'at-at' in attributes
 
 
 class TestTargetTasks(unittest.TestCase):
 
     def default_matches(self, run_on_projects, project):
         method = target_tasks.get_method('default')
         graph = TaskGraph(tasks={
-            'a': TestTask(kind='build', label='a',
-                          attributes={'run_on_projects': run_on_projects}),
+            'a': Task(kind='build', label='a',
+                      attributes={'run_on_projects': run_on_projects},
+                      task={}),
         }, graph=Graph(nodes={'a'}, edges=set()))
         parameters = {'project': project}
         return 'a' in method(graph, parameters)
 
     def test_default_all(self):
         """run_on_projects=[all] includes release, integration, and other projects"""
         self.assertTrue(self.default_matches(['all'], 'mozilla-central'))
         self.assertTrue(self.default_matches(['all'], 'mozilla-inbound'))
@@ -62,18 +63,18 @@ class TestTargetTasks(unittest.TestCase)
     def test_default_nothing(self):
         """run_on_projects=[] includes nothing"""
         self.assertFalse(self.default_matches([], 'mozilla-central'))
         self.assertFalse(self.default_matches([], 'mozilla-inbound'))
         self.assertFalse(self.default_matches([], 'baobab'))
 
     def test_try_option_syntax(self):
         tasks = {
-            'a': TestTask(kind=None, label='a'),
-            'b': TestTask(kind=None, label='b', attributes={'at-at': 'yep'}),
+            'a': Task(kind=None, label='a', attributes={}, task={}),
+            'b': Task(kind=None, label='b', attributes={'at-at': 'yep'}, task={}),
         }
         graph = Graph(nodes=set('ab'), edges=set())
         tg = TaskGraph(tasks, graph)
         params = {'message': 'try me'}
 
         orig_TryOptionSyntax = try_option_syntax.TryOptionSyntax
         try:
             try_option_syntax.TryOptionSyntax = FakeTryOptionSyntax
--- a/taskcluster/taskgraph/test/test_taskgraph.py
+++ b/taskcluster/taskgraph/test/test_taskgraph.py
@@ -2,51 +2,77 @@
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
 from __future__ import absolute_import, print_function, unicode_literals
 
 import unittest
 
 from ..graph import Graph
-from ..task.docker_image import DockerImageTask
-from ..task.transform import TransformTask
+from ..task.base import Task
 from ..taskgraph import TaskGraph
 from mozunit import main
 
 
 class TestTaskGraph(unittest.TestCase):
 
-    def test_from_json(self):
-        task = {
-            "routes": [],
-            "extra": {
-                "imageMeta": {
-                    "contextHash": "<hash>",
-                    "imageName": "<image>",
-                    "level": "1"
-                }
+    maxDiff = None
+
+    def test_taskgraph_to_json(self):
+        tasks = {
+            'a': Task(kind='test', label='a',
+                      attributes={'attr': 'a-task'},
+                      task={'taskdef': True}),
+            'b': Task(kind='test', label='b',
+                      attributes={},
+                      task={'task': 'def'},
+                      optimizations=[['seta']],
+                      # note that this dep is ignored, superseded by that
+                      # from the taskgraph's edges
+                      dependencies={'first': 'a'}),
+        }
+        graph = Graph(nodes=set('ab'), edges={('a', 'b', 'edgelabel')})
+        taskgraph = TaskGraph(tasks, graph)
+
+        res = taskgraph.to_json()
+
+        self.assertEqual(res, {
+            'a': {
+                'kind': 'test',
+                'label': 'a',
+                'attributes': {'attr': 'a-task', 'kind': 'test'},
+                'task': {'taskdef': True},
+                'dependencies': {'edgelabel': 'b'},
+                'optimizations': [],
+            },
+            'b': {
+                'kind': 'test',
+                'label': 'b',
+                'attributes': {'kind': 'test'},
+                'task': {'task': 'def'},
+                'dependencies': {},
+                'optimizations': [['seta']],
             }
-        }
+        })
+
+    def test_round_trip(self):
         graph = TaskGraph(tasks={
-            'a': TransformTask(
+            'a': Task(
                 kind='fancy',
-                task={
-                    'label': 'a',
-                    'attributes': {},
-                    'dependencies': {},
-                    'when': {},
-                    'task': {'task': 'def'},
-                }),
-            'b': DockerImageTask(kind='docker-image',
-                                 task={
-                                     'label': 'b',
-                                     'attributes': {},
-                                     'task': task,
-                                 }),
-        }, graph=Graph(nodes={'a', 'b'}, edges=set()))
+                label='a',
+                attributes={},
+                dependencies={'prereq': 'b'},  # must match edges, below
+                optimizations=[['seta']],
+                task={'task': 'def'}),
+            'b': Task(
+                kind='pre',
+                label='b',
+                attributes={},
+                dependencies={},
+                optimizations=[['seta']],
+                task={'task': 'def2'}),
+        }, graph=Graph(nodes={'a', 'b'}, edges={('a', 'b', 'prereq')}))
 
         tasks, new_graph = TaskGraph.from_json(graph.to_json())
-        self.assertEqual(graph.tasks['a'], new_graph.tasks['a'])
         self.assertEqual(graph, new_graph)
 
 if __name__ == '__main__':
     main()
--- a/taskcluster/taskgraph/test/test_try_option_syntax.py
+++ b/taskcluster/taskgraph/test/test_try_option_syntax.py
@@ -5,35 +5,35 @@
 from __future__ import absolute_import, print_function, unicode_literals
 
 import unittest
 
 from ..try_option_syntax import TryOptionSyntax
 from ..try_option_syntax import RIDEALONG_BUILDS
 from ..graph import Graph
 from ..taskgraph import TaskGraph
-from .util import TestTask
+from ..task.base import Task
 from mozunit import main
 
 # an empty graph, for things that don't look at it
 empty_graph = TaskGraph({}, Graph(set(), set()))
 
 
 def unittest_task(n, tp):
-    return (n, TestTask('test', n, {
+    return (n, Task('test', n, {
         'unittest_try_name': n,
         'test_platform': tp,
-    }))
+    }, {}))
 
 
 def talos_task(n, tp):
-    return (n, TestTask('test', n, {
+    return (n, Task('test', n, {
         'talos_try_name': n,
         'test_platform': tp,
-    }))
+    }, {}))
 
 tasks = {k: v for k, v in [
     unittest_task('mochitest-browser-chrome', 'linux'),
     unittest_task('mochitest-browser-chrome-e10s', 'linux64'),
     unittest_task('mochitest-chrome', 'linux'),
     unittest_task('mochitest-webgl', 'linux'),
     unittest_task('crashtest-e10s', 'linux'),
     unittest_task('gtest', 'linux64'),
deleted file mode 100644
--- a/taskcluster/taskgraph/test/util.py
+++ /dev/null
@@ -1,24 +0,0 @@
-# This Source Code Form is subject to the terms of the Mozilla Public
-# License, v. 2.0. If a copy of the MPL was not distributed with this
-# file, You can obtain one at http://mozilla.org/MPL/2.0/.
-
-from __future__ import absolute_import, print_function, unicode_literals
-
-from ..task import base
-
-
-class TestTask(base.Task):
-
-    def __init__(self, kind=None, label=None, attributes=None, task=None):
-        super(TestTask, self).__init__(
-                kind or 'test',
-                label or 'test-label',
-                attributes or {},
-                task or {})
-
-    @classmethod
-    def load_tasks(cls, kind, path, config, parameters):
-        return []
-
-    def get_dependencies(self, taskgraph):
-        return []
--- a/taskcluster/taskgraph/transforms/docker_image.py
+++ b/taskcluster/taskgraph/transforms/docker_image.py
@@ -62,33 +62,24 @@ def fill_template(config, tasks):
         # for a given context hash, a worker within Taskcluster does not need to contain
         # the same image per branch.
         optimizations = [['index-search', '{}.level-{}.{}.hash.{}'.format(
             INDEX_PREFIX, level, image_name, context_hash)]
             for level in reversed(range(int(config.params['level']), 4))]
 
         # include some information that is useful in reconstructing this task
         # from JSON
-        extra = {
-            'imageMeta': {
-                'level': config.params['level'],
-                'contextHash': context_hash,
-                'imageName': image_name,
-            },
-        }
-
         taskdesc = {
             'label': 'build-docker-image-' + image_name,
             'description': description,
             'attributes': {'image_name': image_name},
             'expires-after': '1 year',
             'routes': routes,
             'optimizations': optimizations,
             'scopes': ['secrets:get:project/taskcluster/gecko/hgfingerprint'],
-            'extra': extra,
             'treeherder': {
                 'symbol': job_symbol,
                 'platform': 'taskcluster-images/opt',
                 'kind': 'other',
                 'tier': 1,
             },
             'run-on-projects': [],
             'worker-type': 'aws-provisioner-v1/gecko-images',