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