1# This Source Code Form is subject to the terms of the Mozilla Public
2# License, v. 2.0. If a copy of the MPL was not distributed with this
3# file, You can obtain one at http://mozilla.org/MPL/2.0/.
4import copy
5import contextlib
6
7from mozperftest.test import pick_test
8from mozperftest.system import pick_system
9from mozperftest.metrics import pick_metrics
10from mozperftest.layers import Layers, StopRunError
11from mozperftest.utils import MachLogger
12from mozperftest.hooks import Hooks
13from mozperftest.argparser import FLAVORS
14
15
16SYSTEM, TEST, METRICS = 0, 1, 2
17
18
19class MachEnvironment(MachLogger):
20    def __init__(self, mach_cmd, flavor="desktop-browser", hooks=None, **kwargs):
21        MachLogger.__init__(self, mach_cmd)
22        self._mach_cmd = mach_cmd
23        self._mach_args = dict(
24            [(self._normalize(key), value) for key, value in kwargs.items()]
25        )
26        self.layers = []
27        if flavor not in FLAVORS:
28            raise NotImplementedError(flavor)
29        for layer in (pick_system, pick_test, pick_metrics):
30            self.add_layer(layer(self, flavor, mach_cmd))
31        if hooks is None:
32            # we just load an empty Hooks instance
33            hooks = Hooks(mach_cmd)
34        self.hooks = hooks
35
36    @contextlib.contextmanager
37    def frozen(self):
38        self.freeze()
39        try:
40            # used to trigger __enter__/__exit__
41            with self:
42                yield self
43        finally:
44            self.unfreeze()
45
46    def _normalize(self, name):
47        if name.startswith("--"):
48            name = name[2:]
49        return name.replace("-", "_")
50
51    def set_arg(self, name, value):
52        """Sets the argument"""
53        # see if we want to restrict to existing keys
54        self._mach_args[self._normalize(name)] = value
55
56    def get_arg(self, name, default=None, layer=None):
57        name = self._normalize(name)
58        marker = object()
59        res = self._mach_args.get(name, marker)
60        if res is marker:
61            # trying with the name prefixed with the layer name
62            if layer is not None and not name.startswith(layer.name):
63                name = "%s_%s" % (layer.name, name)
64                return self._mach_args.get(name, default)
65            return default
66        return res
67
68    def get_layer(self, name):
69        for layer in self.layers:
70            if isinstance(layer, Layers):
71                found = layer.get_layer(name)
72                if found is not None:
73                    return found
74            elif layer.name == name:
75                return layer
76        return None
77
78    def add_layer(self, layer):
79        self.layers.append(layer)
80
81    def freeze(self):
82        # freeze args  (XXX do we need to freeze more?)
83        self._saved_mach_args = copy.deepcopy(self._mach_args)
84
85    def unfreeze(self):
86        self._mach_args = self._saved_mach_args
87        self._saved_mach_args = None
88
89    def run(self, metadata):
90        # run the system and test layers
91        try:
92            with self.layers[SYSTEM] as syslayer, self.layers[TEST] as testlayer:
93                metadata = testlayer(syslayer(metadata))
94
95            # then run the metrics layers
96            with self.layers[METRICS] as metrics:
97                metadata = metrics(metadata)
98        except StopRunError:
99            # ends the cycle but without bubbling up the error
100            pass
101        return metadata
102
103    def __enter__(self):
104        return self
105
106    def __exit__(self, type, value, traceback):
107        return
108