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/.
4
5# Integrates Talos mozharness with mach
6
7from __future__ import absolute_import, print_function, unicode_literals
8
9import logging
10import os
11import six
12import sys
13import json
14import socket
15
16from mozbuild.base import (
17    MozbuildObject,
18    BinaryNotFoundException,
19)
20from mach.decorators import Command
21
22HERE = os.path.dirname(os.path.realpath(__file__))
23
24
25class TalosRunner(MozbuildObject):
26    def run_test(self, talos_args):
27        """
28        We want to do couple of things before running Talos
29        1. Clone mozharness
30        2. Make config for Talos Mozharness
31        3. Run mozharness
32        """
33
34        try:
35            self.init_variables(talos_args)
36        except BinaryNotFoundException as e:
37            self.log(logging.ERROR, "talos", {"error": str(e)}, "ERROR: {error}")
38            self.log(logging.INFO, "raptor", {"help": e.help()}, "{help}")
39            return 1
40
41        self.make_config()
42        self.write_config()
43        self.make_args()
44        return self.run_mozharness()
45
46    def init_variables(self, talos_args):
47        self.talos_dir = os.path.join(self.topsrcdir, "testing", "talos")
48        self.mozharness_dir = os.path.join(self.topsrcdir, "testing", "mozharness")
49        self.talos_json = os.path.join(self.talos_dir, "talos.json")
50        self.config_file_path = os.path.join(
51            self._topobjdir, "testing", "talos-in_tree_conf.json"
52        )
53        self.binary_path = self.get_binary_path()
54        self.virtualenv_script = os.path.join(
55            self.topsrcdir, "third_party", "python", "virtualenv", "virtualenv.py"
56        )
57        self.virtualenv_path = os.path.join(self._topobjdir, "testing", "talos-venv")
58        self.python_interp = sys.executable
59        self.talos_args = talos_args
60
61    def make_config(self):
62        default_actions = ["populate-webroot"]
63        default_actions.extend(
64            [
65                "create-virtualenv",
66                "run-tests",
67            ]
68        )
69        self.config = {
70            "run_local": True,
71            "talos_json": self.talos_json,
72            "binary_path": self.binary_path,
73            "repo_path": self.topsrcdir,
74            "obj_path": self.topobjdir,
75            "log_name": "talos",
76            "virtualenv_path": self.virtualenv_path,
77            "pypi_url": "http://pypi.python.org/simple",
78            "base_work_dir": self.mozharness_dir,
79            "exes": {
80                "python": self.python_interp,
81                "virtualenv": [self.python_interp, self.virtualenv_script],
82            },
83            "title": socket.gethostname(),
84            "default_actions": default_actions,
85            "talos_extra_options": ["--develop"] + self.talos_args,
86            "python3_manifest": {
87                "win32": "python3.manifest",
88                "win64": "python3_x64.manifest",
89            },
90        }
91
92    def make_args(self):
93        self.args = {
94            "config": {},
95            "initial_config_file": self.config_file_path,
96        }
97
98    def write_config(self):
99        try:
100            config_file = open(self.config_file_path, "wb")
101            config_file.write(six.ensure_binary(json.dumps(self.config)))
102        except IOError as e:
103            err_str = "Error writing to Talos Mozharness config file {0}:{1}"
104            print(err_str.format(self.config_file_path, str(e)))
105            raise e
106
107    def run_mozharness(self):
108        sys.path.insert(0, self.mozharness_dir)
109        from mozharness.mozilla.testing.talos import Talos
110
111        talos_mh = Talos(
112            config=self.args["config"],
113            initial_config_file=self.args["initial_config_file"],
114        )
115        return talos_mh.run()
116
117
118def create_parser():
119    sys.path.insert(0, HERE)  # allow to import the talos package
120    from talos.cmdline import create_parser
121
122    return create_parser(mach_interface=True)
123
124
125@Command(
126    "talos-test",
127    category="testing",
128    description="Run talos tests (performance testing).",
129    parser=create_parser,
130)
131def run_talos_test(command_context, **kwargs):
132    talos = command_context._spawn(TalosRunner)
133
134    try:
135        return talos.run_test(sys.argv[2:])
136    except Exception as e:
137        print(str(e))
138        return 1
139