Bug 808336 - Part 1: Refactor mach command handler management; r=jhammel
authorGregory Szorc <gps@mozilla.com>
Tue, 06 Nov 2012 16:57:41 -0800
changeset 120336 e0e1220d2c5a1086a4911be8d69e2937b842d116
parent 120335 bef9c0bcb62fdab7ca7895307117195f593570b0
child 120337 4f96fdb5e4f9b9e101e2e3c5f535bff8e0fe8f6c
push id1997
push userakeybl@mozilla.com
push dateMon, 07 Jan 2013 21:25:26 +0000
treeherdermozilla-beta@4baf45cdcf21 [default view] [failures only]
perfherder[talos] [build metrics] [platform microbench] (compared to previous push)
reviewersjhammel
bugs808336
milestone19.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 808336 - Part 1: Refactor mach command handler management; r=jhammel Previously we were tighly coupled with MozbuildObject. This was not in the spirit of mach being a generic tool. Now, instead of passing multiple arguments to __init__ of the class providing the mach command we either pass 0 or 1. The number of arguments is detected when processing the @CommandProvider decorator. The optional argument is a named tuple containing mach run-time state. Capturing of mach command provider information is now captured in a class (as opposed to an anoymous tuple). We also capture these in a rich data structure which is passed as part of the run-time data to the command provider class. This allows mach commands to inspect the mach environment. Mach decorators have been moved to mach.decorators. mach.base is reserved for generic mach data/container classes. Existing mach command classes derived from MozbuildObject have been replaced with either object or mozbuild.base.MachCommandBase. This latter class translates the mach context instance passed to __init__ into the constructor arguments for MozbuildObject.__init__. Support for registering function handlers with mach has been removed. All handlers must be inside classes.
layout/tools/reftest/mach_commands.py
python/mach/README.rst
python/mach/mach/base.py
python/mach/mach/commands/settings.py
python/mach/mach/decorators.py
python/mach/mach/main.py
python/mach/mach/registrar.py
python/mach/mach/test/common.py
python/mozboot/mozboot/mach_commands.py
python/mozbuild/mach/commands/build.py
python/mozbuild/mach/commands/warnings.py
python/mozbuild/mozbuild/base.py
testing/mochitest/mach_commands.py
testing/xpcshell/mach_commands.py
--- a/layout/tools/reftest/mach_commands.py
+++ b/layout/tools/reftest/mach_commands.py
@@ -1,21 +1,24 @@
 # 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 unicode_literals
 
 import os
 
-from mozbuild.base import MozbuildObject
+from mozbuild.base import (
+    MachCommandBase,
+    MozbuildObject,
+)
 
 from moztesting.util import parse_test_path
 
-from mach.base import (
+from mach.decorators import (
     CommandArgument,
     CommandProvider,
     Command,
 )
 
 
 generic_help = 'Test to run. Can be specified as a single file, a ' +\
 'directory, or omitted. If omitted, the entire test suite is executed.'
@@ -70,17 +73,17 @@ class ReftestRunner(MozbuildObject):
         else:
             env = {}
 
         # TODO hook up harness via native Python
         self._run_make(directory='.', target=suite, append_env=env)
 
 
 @CommandProvider
-class MachCommands(MozbuildObject):
+class MachCommands(MachCommandBase):
     @Command('reftest', help='Run a reftest.')
     @CommandArgument('test_file', default=None, nargs='?', metavar='TEST',
         help=generic_help)
     def run_reftest(self, test_file):
         self._run_reftest(test_file, 'reftest')
 
     @Command('crashtest', help='Run a crashtest.')
     @CommandArgument('test_file', default=None, nargs='?', metavar='TEST',
--- a/python/mach/README.rst
+++ b/python/mach/README.rst
@@ -8,37 +8,40 @@ instantiating the *Mach* class from the 
 
 Implementing mach Commands
 --------------------------
 
 The *mach* driver follows the convention of popular tools like Git,
 Subversion, and Mercurial and provides a common driver for multiple
 subcommands.
 
-Subcommands are implemented by decorating a class inheritting from
-mozbuild.base.MozbuildObject and by decorating methods that act as
-subcommand handlers.
+Subcommands are implemented by decorating a class and by decorating
+methods that act as subcommand handlers.
 
-Relevant decorators are defined in the *mach.base* module. There are
+Relevant decorators are defined in the *mach.decorators* module. There are
 the *Command* and *CommandArgument* decorators, which should be used
 on methods to denote that a specific method represents a handler for
 a mach subcommand. There is also the *CommandProvider* decorator,
 which is applied to a class to denote that it contains mach subcommands.
 
+Classes with the *@CommandProvider* decorator *must* have an *__init__*
+method that accepts 1 or 2 arguments. If it accepts 2 arguments, the
+2nd argument will be a *MachCommandContext* instance. This is just a named
+tuple containing references to objects provided by the mach driver.
+
 Here is a complete example:
 
-    from mozbuild.base import MozbuildObject
-
-    from mach.base import CommandArgument
-    from mach.base import CommandProvider
-    from mach.base import Command
+    from mach.decorators import (
+        CommandArgument,
+        CommandProvider,
+        Command,
+    )
 
     @CommandProvider
-    class MyClass(MozbuildObject):
-
+    class MyClass(object):
         @Command('doit', help='Do ALL OF THE THINGS.')
         @CommandArgument('--force', '-f', action='store_true',
             help='Force doing it.')
         def doit(self, force=False):
             # Do stuff here.
 
 
 When the module is loaded, the decorators tell mach about all handlers.
--- a/python/mach/mach/base.py
+++ b/python/mach/mach/base.py
@@ -1,94 +1,52 @@
 # 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 unicode_literals
 
-import types
-
-from mach.registrar import register_method_handler
-
-
-def CommandProvider(cls):
-    """Class decorator to denote that it provides subcommands for Mach.
-
-    When this decorator is present, mach looks for commands being defined by
-    methods inside the class.
-    """
-
-    # The implementation of this decorator relies on the parse-time behavior of
-    # decorators. When the module is imported, the method decorators (like
-    # @Command and @CommandArgument) are called *before* this class decorator.
-    # The side-effect of the method decorators is to store specifically-named
-    # attributes on the function types. We just scan over all functions in the
-    # class looking for the side-effects of the method decorators.
+from collections import namedtuple
 
-    # We scan __dict__ because we only care about the classes own attributes,
-    # not inherited ones. If we did inherited attributes, we could potentially
-    # define commands multiple times. We also sort keys so commands defined in
-    # the same class are grouped in a sane order.
-    for attr in sorted(cls.__dict__.keys()):
-        value = cls.__dict__[attr]
-
-        if not isinstance(value, types.FunctionType):
-            continue
-
-        parser_args = getattr(value, '_mach_command', None)
-        if parser_args is None:
-            continue
-
-        arguments = getattr(value, '_mach_command_args', None)
-
-        register_method_handler(cls, attr, (parser_args[0], parser_args[1]),
-            arguments or [])
-
-    return cls
+# Holds mach run-time state so it can easily be passed to command providers.
+CommandContext = namedtuple('CommandContext', ['topdir', 'cwd',
+    'settings', 'log_manager', 'commands'])
 
 
-class Command(object):
-    """Decorator for functions or methods that provide a mach subcommand.
-
-    The decorator accepts arguments that would be passed to add_parser() of an
-    ArgumentParser instance created via add_subparsers(). Essentially, it
-    accepts the arguments one would pass to add_argument().
-
-    For example:
+class MethodHandler(object):
+    """Describes a Python method that implements a mach command.
 
-        @Command('foo', help='Run the foo action')
-        def foo(self):
-            pass
+    Instances of these are produced by mach when it processes classes
+    defining mach commands.
     """
-    def __init__(self, *args, **kwargs):
-        self._command_args = (args, kwargs)
+    __slots__ = (
+        # The Python class providing the command. This is the class type not
+        # an instance of the class. Mach will instantiate a new instance of
+        # the class if the command is executed.
+        'cls',
 
-    def __call__(self, func):
-        func._mach_command = self._command_args
-
-        return func
-
+        # Whether the __init__ method of the class should receive a mach
+        # context instance. This should only affect the mach driver and how
+        # it instantiates classes.
+        'pass_context',
 
-class CommandArgument(object):
-    """Decorator for additional arguments to mach subcommands.
+        # The name of the method providing the command. In other words, this
+        # is the str name of the attribute on the class type corresponding to
+        # the name of the function.
+        'method',
 
-    This decorator should be used to add arguments to mach commands. Arguments
-    to the decorator are proxied to ArgumentParser.add_argument().
-
-    For example:
+        # Arguments passed to add_parser() on the main mach subparser. This is
+        # a 2-tuple of positional and named arguments, respectively.
+        'parser_args',
 
-        @Command('foo', help='Run the foo action')
-        @CommandArgument('-b', '--bar', action='store_true', default=False,
-            help='Enable bar mode.')
-        def foo(self):
-            pass
-    """
-    def __init__(self, *args, **kwargs):
-        self._command_args = (args, kwargs)
+        # Arguments added to this command's parser. This is a 2-tuple of
+        # positional and named arguments, respectively.
+        'arguments',
+    )
 
-    def __call__(self, func):
-        command_args = getattr(func, '_mach_command_args', [])
+    def __init__(self, cls, method, parser_args, arguments=None,
+        pass_context=False):
 
-        command_args.append(self._command_args)
-
-        func._mach_command_args = command_args
-
-        return func
+        self.cls = cls
+        self.method = method
+        self.parser_args = parser_args
+        self.arguments = arguments or []
+        self.pass_context = pass_context
--- a/python/mach/mach/commands/settings.py
+++ b/python/mach/mach/commands/settings.py
@@ -1,28 +1,33 @@
 # 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 print_function, unicode_literals
 
 from textwrap import TextWrapper
 
-from mozbuild.base import MozbuildObject
-from mach.base import CommandProvider
-from mach.base import Command
+from mach.decorators import (
+    CommandProvider,
+    Command,
+)
+
 
 #@CommandProvider
-class Settings(MozbuildObject):
+class Settings(object):
     """Interact with settings for mach.
 
     Currently, we only provide functionality to view what settings are
     available. In the future, this module will be used to modify settings, help
     people create configs via a wizard, etc.
     """
+    def __init__(self, context):
+        self.settings = context.settings
+
     @Command('settings-list', help='Show available config settings.')
     def list_settings(self):
         """List available settings in a concise list."""
         for section in sorted(self.settings):
             for option in sorted(self.settings[section]):
                 short, full = self.settings.option_help(section, option)
                 print('%s.%s -- %s' % (section, option, short))
 
copy from python/mach/mach/base.py
copy to python/mach/mach/decorators.py
--- a/python/mach/mach/base.py
+++ b/python/mach/mach/decorators.py
@@ -1,51 +1,70 @@
 # 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 unicode_literals
 
+import inspect
 import types
 
-from mach.registrar import register_method_handler
+from .base import MethodHandler
+from .registrar import Registrar
 
 
 def CommandProvider(cls):
     """Class decorator to denote that it provides subcommands for Mach.
 
     When this decorator is present, mach looks for commands being defined by
     methods inside the class.
     """
 
     # The implementation of this decorator relies on the parse-time behavior of
     # decorators. When the module is imported, the method decorators (like
     # @Command and @CommandArgument) are called *before* this class decorator.
     # The side-effect of the method decorators is to store specifically-named
     # attributes on the function types. We just scan over all functions in the
     # class looking for the side-effects of the method decorators.
 
+    # Tell mach driver whether to pass context argument to __init__.
+    pass_context = False
+
+    if inspect.ismethod(cls.__init__):
+        spec = inspect.getargspec(cls.__init__)
+
+        if len(spec.args) > 2:
+            msg = 'Mach @CommandProvider class %s implemented incorrectly. ' + \
+                  '__init__() must take 1 or 2 arguments. From %s'
+            msg = msg % (cls.__name__, inspect.getsourcefile(cls))
+            raise Exception(msg)
+
+        if len(spec.args) == 2:
+            pass_context = True
+
     # We scan __dict__ because we only care about the classes own attributes,
     # not inherited ones. If we did inherited attributes, we could potentially
     # define commands multiple times. We also sort keys so commands defined in
     # the same class are grouped in a sane order.
     for attr in sorted(cls.__dict__.keys()):
         value = cls.__dict__[attr]
 
         if not isinstance(value, types.FunctionType):
             continue
 
         parser_args = getattr(value, '_mach_command', None)
         if parser_args is None:
             continue
 
         arguments = getattr(value, '_mach_command_args', None)
 
-        register_method_handler(cls, attr, (parser_args[0], parser_args[1]),
-            arguments or [])
+        handler = MethodHandler(cls, attr, (parser_args[0], parser_args[1]),
+            arguments=arguments, pass_context=pass_context)
+
+        Registrar.register_command_handler(handler)
 
     return cls
 
 
 class Command(object):
     """Decorator for functions or methods that provide a mach subcommand.
 
     The decorator accepts arguments that would be passed to add_parser() of an
--- a/python/mach/mach/main.py
+++ b/python/mach/mach/main.py
@@ -14,45 +14,47 @@ import logging
 import os
 import sys
 import traceback
 import uuid
 import sys
 
 from mozbuild.base import BuildConfig
 
-from .base import (
+from .base import CommandContext
+
+from .decorators import (
     CommandArgument,
     CommandProvider,
     Command,
 )
 
 from .config import ConfigSettings
 from .logging import LoggingManager
 
-from .registrar import populate_argument_parser
+from .registrar import Registrar
 
 
 # Classes inheriting from ConfigProvider that provide settings.
 # TODO this should come from auto-discovery somehow.
 SETTINGS_PROVIDERS = [
     BuildConfig,
 ]
 
 # Settings for argument parser that don't get proxied to sub-module. i.e. these
 # are things consumed by the driver itself.
 CONSUMED_ARGUMENTS = [
     'settings_file',
     'verbose',
     'logfile',
     'log_interval',
     'command',
-    'cls',
-    'method',
-    'func',
+    'mach_class',
+    'mach_method',
+    'mach_pass_context',
 ]
 
 MACH_ERROR = r'''
 The error occurred in mach itself. This is likely a bug in mach itself or a
 fundamental problem with a loaded module.
 
 Please consider filing a bug against mach by going to the URL:
 
@@ -299,28 +301,31 @@ To see more help for a specific command,
         self.log_manager.add_terminal_logging(level=log_level,
             write_interval=args.log_interval)
 
         self.load_settings(args)
 
         stripped = {k: getattr(args, k) for k in vars(args) if k not in
             CONSUMED_ARGUMENTS}
 
-        # If the command is associated with a class, instantiate and run it.
-        # All classes must be Base-derived and take the expected argument list.
-        if hasattr(args, 'cls'):
-            cls = getattr(args, 'cls')
-            instance = cls(self.cwd, self.settings, self.log_manager)
-            fn = getattr(instance, getattr(args, 'method'))
+        context = CommandContext(topdir=self.cwd, cwd=self.cwd,
+            settings=self.settings, log_manager=self.log_manager,
+            commands=Registrar)
+
+        if not hasattr(args, 'mach_class'):
+            raise Exception('ArgumentParser result missing mach_class.')
 
-        # If the command is associated with a function, call it.
-        elif hasattr(args, 'func'):
-            fn = getattr(args, 'func')
+        cls = getattr(args, 'mach_class')
+
+        if getattr(args, 'mach_pass_context'):
+            instance = cls(context)
         else:
-            raise Exception('Dispatch configuration error in module.')
+            instance = cls()
+
+        fn = getattr(instance, getattr(args, 'mach_method'))
 
         try:
             result = fn(**stripped)
 
             if not result:
                 result = 0
 
             assert isinstance(result, int)
@@ -448,17 +453,17 @@ To see more help for a specific command,
             metavar='FILENAME', type=argparse.FileType('ab'),
             help='Filename to write log data to.')
         global_group.add_argument('--log-interval', dest='log_interval',
             action='store_true', default=False,
             help='Prefix log line with interval from last message rather '
                 'than relative time. Note that this is NOT execution time '
                 'if there are parallel operations.')
 
-        populate_argument_parser(subparser)
+        Registrar.populate_argument_parser(subparser)
 
         return parser
 
     @Command('help', help='Show mach usage info or help for a command.')
     @CommandArgument('subcommand', default=None, nargs='?',
         help='Command to show help info for.')
     def _help(self, subcommand=None):
         # The built-in handler doesn't pass the original ArgumentParser into
--- a/python/mach/mach/registrar.py
+++ b/python/mach/mach/registrar.py
@@ -1,21 +1,35 @@
 # 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 unicode_literals
 
-
-class_handlers = []
-
-def register_method_handler(cls, method, parser_args, arguments):
-    class_handlers.append((cls, method, parser_args, arguments))
+import collections
 
 
-def populate_argument_parser(parser):
-    for cls, method, parser_args, arguments in class_handlers:
-        p = parser.add_parser(*parser_args[0], **parser_args[1])
+class MachRegistrar(object):
+    """Container for mach command and config providers."""
+
+    def __init__(self):
+        self.command_handlers = {}
+
+    def register_command_handler(self, handler):
+        name = handler.parser_args[0][0]
+
+        self.command_handlers[name] = handler
 
-        for arg in arguments:
-            p.add_argument(*arg[0], **arg[1])
+    def populate_argument_parser(self, parser):
+        for command in sorted(self.command_handlers.keys()):
+            handler = self.command_handlers[command]
+            p = parser.add_parser(*handler.parser_args[0],
+                **handler.parser_args[1])
 
-        p.set_defaults(cls=cls, method=method)
+            for arg in handler.arguments:
+                p.add_argument(*arg[0], **arg[1])
+
+            p.set_defaults(mach_class=handler.cls, mach_method=handler.method,
+                mach_pass_context=handler.pass_context)
+
+
+Registrar = MachRegistrar()
+
--- a/python/mach/mach/test/common.py
+++ b/python/mach/mach/test/common.py
@@ -1,27 +1,27 @@
 # 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 unicode_literals
 
 import time
 
-from mozbuild.base import MozbuildObject
-
-from mach.base import CommandArgument
-from mach.base import CommandProvider
-from mach.base import Command
+from mach.base import (
+    CommandArgument,
+    CommandProvider,
+    Command,
+)
 
 import mach.test.common2 as common2
 
 
 @CommandProvider
-class TestCommandProvider(MozbuildObject):
+class TestCommandProvider(object):
     @Command('throw')
     @CommandArgument('--message', '-m', default='General Error')
     def throw(self, message):
         raise Exception(message)
 
     @Command('throw_deep')
     @CommandArgument('--message', '-m', default='General Error')
     def throw_deep(self, message):
--- a/python/mozboot/mozboot/mach_commands.py
+++ b/python/mozboot/mozboot/mach_commands.py
@@ -1,22 +1,23 @@
 # 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 unicode_literals
 
-from mozbuild.base import MozbuildObject
+from mach.decorators import (
+    CommandArgument,
+    CommandProvider,
+    Command,
+)
 
-from mach.base import CommandArgument
-from mach.base import CommandProvider
-from mach.base import Command
 
 @CommandProvider
-class Bootstrap(MozbuildObject):
+class Bootstrap(object):
     """Bootstrap system and mach for optimal development experience."""
 
     @Command('bootstrap',
         help='Install required system packages for building.')
     def bootstrap(self):
         from mozboot.bootstrap import Bootstrapper
 
         bootstrapper = Bootstrapper()
--- a/python/mozbuild/mach/commands/build.py
+++ b/python/mozbuild/mach/commands/build.py
@@ -2,23 +2,26 @@
 # 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 unicode_literals
 
 import logging
 import os
 
-from mach.base import CommandProvider
-from mach.base import Command
-from mozbuild.base import MozbuildObject
+from mach.decorators import (
+    CommandProvider,
+    Command,
+)
+
+from mozbuild.base import MachCommandBase
 
 
 @CommandProvider
-class Build(MozbuildObject):
+class Build(MachCommandBase):
     """Interface to build the tree."""
 
     @Command('build', help='Build the tree.')
     def build(self):
         # This code is only meant to be temporary until the more robust tree
         # building code in bug 780329 lands.
         from mozbuild.compilation.warnings import WarningsCollector
         from mozbuild.compilation.warnings import WarningsDatabase
--- a/python/mozbuild/mach/commands/warnings.py
+++ b/python/mozbuild/mach/commands/warnings.py
@@ -2,24 +2,27 @@
 # 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 print_function, unicode_literals
 
 import operator
 import os
 
-from mach.base import CommandArgument
-from mach.base import CommandProvider
-from mach.base import Command
-from mozbuild.base import MozbuildObject
+from mach.decorators import (
+    CommandArgument,
+    CommandProvider,
+    Command,
+)
+
+from mozbuild.base import MachCommandBase
 
 
 @CommandProvider
-class Warnings(MozbuildObject):
+class Warnings(MachCommandBase):
     """Provide commands for inspecting warnings."""
 
     @property
     def database_path(self):
         return self._get_state_filename('warnings.json')
 
     @property
     def database(self):
--- a/python/mozbuild/mozbuild/base.py
+++ b/python/mozbuild/mozbuild/base.py
@@ -291,8 +291,20 @@ class BuildConfig(ConfigProvider):
     @classmethod
     def _register_settings(cls):
         def register(section, option, type_cls, **kwargs):
             cls.register_setting(section, option, type_cls, domain='mozbuild',
                 **kwargs)
 
         register('build', 'threads', PositiveIntegerType,
             default=multiprocessing.cpu_count())
+
+
+class MachCommandBase(MozbuildObject):
+    """Base class for mach command providers that wish to be MozbuildObjects.
+
+    This provides a level of indirection so MozbuildObject can be refactored
+    without having to change everything that inherits from it.
+    """
+
+    def __init__(self, context):
+        MozbuildObject.__init__(self, context.topdir, context.settings,
+            context.log_manager)
--- a/testing/mochitest/mach_commands.py
+++ b/testing/mochitest/mach_commands.py
@@ -1,21 +1,24 @@
 # 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 unicode_literals
 
 import os
 
-from mozbuild.base import MozbuildObject
+from mozbuild.base import (
+    MachCommandBase,
+    MozbuildObject,
+)
 
 from moztesting.util import parse_test_path
 
-from mach.base import (
+from mach.decorators import (
     CommandArgument,
     CommandProvider,
     Command,
 )
 
 
 generic_help = 'Test to run. Can be specified as a single file, a ' +\
 'directory, or omitted. If omitted, the entire test suite is executed.'
@@ -78,17 +81,17 @@ class MochitestRunner(MozbuildObject):
             env = {'TEST_PATH': path}
         else:
             env = {}
 
         self._run_make(directory='.', target=target, append_env=env)
 
 
 @CommandProvider
-class MachCommands(MozbuildObject):
+class MachCommands(MachCommandBase):
     @Command('mochitest-plain', help='Run a plain mochitest.')
     @CommandArgument('test_file', default=None, nargs='?', metavar='TEST',
         help=generic_help)
     def run_mochitest_plain(self, test_file):
         self.run_mochitest(test_file, 'plain')
 
     @Command('mochitest-chrome', help='Run a chrome mochitest.')
     @CommandArgument('test_file', default=None, nargs='?', metavar='TEST',
--- a/testing/xpcshell/mach_commands.py
+++ b/testing/xpcshell/mach_commands.py
@@ -5,19 +5,22 @@
 # Integrates the xpcshell test runner with mach.
 
 from __future__ import unicode_literals
 
 import os
 
 from StringIO import StringIO
 
-from mozbuild.base import MozbuildObject
+from mozbuild.base import (
+    MachCommandBase,
+    MozbuildObject,
+)
 
-from mach.base import (
+from mach.decorators import (
     CommandArgument,
     CommandProvider,
     Command,
 )
 
 
 class XPCShellRunner(MozbuildObject):
     """Run xpcshell tests."""
@@ -106,17 +109,17 @@ class XPCShellRunner(MozbuildObject):
 
         # TODO do something with result.
         xpcshell.runTests(**args)
 
         self.log_manager.disable_unstructured()
 
 
 @CommandProvider
-class MachCommands(MozbuildObject):
+class MachCommands(MachCommandBase):
     @Command('xpcshell-test', help='Run an xpcshell test.')
     @CommandArgument('test_file', default='all', nargs='?', metavar='TEST',
         help='Test to run. Can be specified as a single JS file, a directory, '
              'or omitted. If omitted, the entire test suite is executed.')
     @CommandArgument('--debug', '-d', action='store_true',
         help='Run test in a debugger.')
     @CommandArgument('--interactive', '-i', action='store_true',
         help='Open an xpcshell prompt before running tests.')