1 files changed, 106 insertions, 0 deletions
diff --git a/python/mozperftest/mozperftest/environment.py b/python/mozperftest/mozperftest/environment.py
new file mode 100644
index 0000000000..7be0f1f639
--- /dev/null
+++ b/python/mozperftest/mozperftest/environment.py
@@ -0,0 +1,106 @@
+# 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/.
+import contextlib
+import copy
+
+from mozperftest.argparser import FLAVORS
+from mozperftest.hooks import Hooks
+from mozperftest.layers import Layers, StopRunError
+from mozperftest.metrics import pick_metrics
+from mozperftest.system import pick_system
+from mozperftest.test import pick_test
+from mozperftest.utils import MachLogger
+
+SYSTEM, TEST, METRICS = 0, 1, 2
+
+
+class MachEnvironment(MachLogger):
+    def __init__(self, mach_cmd, flavor="desktop-browser", hooks=None, **kwargs):
+        MachLogger.__init__(self, mach_cmd)
+        self._mach_cmd = mach_cmd
+        self._mach_args = dict(
+            [(self._normalize(key), value) for key, value in kwargs.items()]
+        )
+        self.layers = []
+        if flavor not in FLAVORS:
+            raise NotImplementedError(flavor)
+        for layer in (pick_system, pick_test, pick_metrics):
+            self.add_layer(layer(self, flavor, mach_cmd))
+        if hooks is None:
+            # we just load an empty Hooks instance
+            hooks = Hooks(mach_cmd)
+        self.hooks = hooks
+
+    @contextlib.contextmanager
+    def frozen(self):
+        self.freeze()
+        try:
+            # used to trigger __enter__/__exit__
+            with self:
+                yield self
+        finally:
+            self.unfreeze()
+
+    def _normalize(self, name):
+        if name.startswith("--"):
+            name = name[2:]
+        return name.replace("-", "_")
+
+    def set_arg(self, name, value):
+        """Sets the argument"""
+        # see if we want to restrict to existing keys
+        self._mach_args[self._normalize(name)] = value
+
+    def get_arg(self, name, default=None, layer=None):
+        name = self._normalize(name)
+        marker = object()
+        res = self._mach_args.get(name, marker)
+        if res is marker:
+            # trying with the name prefixed with the layer name
+            if layer is not None and not name.startswith(layer.name):
+                name = "%s_%s" % (layer.name, name)
+                return self._mach_args.get(name, default)
+            return default
+        return res
+
+    def get_layer(self, name):
+        for layer in self.layers:
+            if isinstance(layer, Layers):
+                found = layer.get_layer(name)
+                if found is not None:
+                    return found
+            elif layer.name == name:
+                return layer
+        return None
+
+    def add_layer(self, layer):
+        self.layers.append(layer)
+
+    def freeze(self):
+        # freeze args  (XXX do we need to freeze more?)
+        self._saved_mach_args = copy.deepcopy(self._mach_args)
+
+    def unfreeze(self):
+        self._mach_args = self._saved_mach_args
+        self._saved_mach_args = None
+
+    def run(self, metadata):
+        # run the system and test layers
+        try:
+            with self.layers[SYSTEM] as syslayer, self.layers[TEST] as testlayer:
+                metadata = testlayer(syslayer(metadata))
+
+            # then run the metrics layers
+            with self.layers[METRICS] as metrics:
+                metadata = metrics(metadata)
+        except StopRunError:
+            # ends the cycle but without bubbling up the error
+            pass
+        return metadata
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type, value, traceback):
+        return