Source code
Revision control
Copy as Markdown
Other Tools
# 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 argparse
import getpass
import importlib.util
import json
import os
import subprocess
import urllib.request
from pathlib import Path
from typing import Optional
import tomllib
from mach.decorators import Command, CommandArgument, SubCommand
from mozbuild.base import MachCommandBase
# Firefox Account OAuth configuration
FXA_SCOPES = ("profile",)
FXA_CLIENT_ID = "5882386c6d801776"
FXA_ACCOUNT_SERVER_URL = "https://api.accounts.firefox.com"
FXA_OAUTH_SERVER_URL = "https://oauth.accounts.firefox.com"
ON_TRY = "MOZ_AUTOMATION" in os.environ
def _dispatch_perftest(mach_context, test_path: str, extra_args: Optional[list] = None):
perftest_args = [test_path]
if extra_args:
# Strip leading dashes so mochitest args match perftest expectations.
extra_args = [arg.lstrip("-") for arg in extra_args]
perftest_args.extend(["--mochitest-extra-args", *extra_args])
# Forward directly to perftest with translated mochitest arguments.
return mach_context.commands.dispatch(
"perftest",
mach_context,
perftest_args,
)
class EvalCommand(MachCommandBase):
"""Forward eval runs to mozperftest and mochitests."""
@Command(
"eval",
category="testing",
description="Run evaluation tests, backed by perftests and mochitests.",
)
@CommandArgument(
"test_path",
help="The path to an individual .js test or a .toml manifest relative to the "
"repo root. The test must be an evaluation test located in a browser_eval folder.",
)
@CommandArgument(
"extra_args",
nargs=argparse.REMAINDER,
help="Additional mochitest arguments passed through to perftest.",
)
def run_eval(self, test_path: str, extra_args: Optional[list] = None):
if Path(test_path).is_dir():
test_path = str(Path(test_path) / "eval.toml")
if Path(test_path).suffix not in (".js", ".toml"):
print(
"Error: test_path must be a .js test file or a .toml manifest or a dir with an eval.toml."
)
return 1
if not Path(test_path).exists():
print(f"Error: {test_path} does not exist.")
return 1
if test_path.endswith(".toml"):
with open(test_path, "rb") as f:
manifest = tomllib.load(f)
toml_dir = Path(test_path).resolve().parent
test_files = [
str(toml_dir / key) for key in manifest if key.endswith(".js")
]
print(f"Discovered {len(test_files)} test(s) in manifest {test_path}:")
for test_file in test_files:
print(f" {test_file}")
if not test_files:
print(f"No JavaScript tests found in {test_path}")
return 1
result = 0
for test_file in test_files:
r = _dispatch_perftest(self._mach_context, test_file, extra_args)
if r:
return r
return result
return _dispatch_perftest(self._mach_context, test_path, extra_args)
def _get_fxa_token(mach_cmd):
mach_cmd.activate_virtualenv()
if importlib.util.find_spec("fxa") is None:
try:
mach_cmd.virtualenv_manager.install_pip_package("PyFxA==0.8.1")
except Exception as exception:
print(f"Failed to install 'fxa' package: {exception}")
raise exception
from fxa import core, oauth
from fxa.errors import ClientError
from fxa.tools.bearer import get_bearer_token
from fxa.tools.unblock import send_unblock_code
print("Login to your Firefox Account (accounts.firefox.com)")
email = input("Email: ").strip()
password = getpass.getpass("Password: ").strip()
try:
token = get_bearer_token(
email,
password,
scopes=FXA_SCOPES,
client_id=FXA_CLIENT_ID,
account_server_url=FXA_ACCOUNT_SERVER_URL,
oauth_server_url=FXA_OAUTH_SERVER_URL,
)
except ClientError as exception:
try:
if "Unconfirmed session" not in str(exception):
raise
try:
send_unblock_code(email, FXA_ACCOUNT_SERVER_URL)
except ClientError:
print("Login failed: unable to send unblock code.")
raise exception
print("\nAn authorization code was sent to your email, enter it here.")
unblock_code = input("Code: ").strip()
try:
# Attempt to login without 2 factor authentication.
session = core.Client(server_url=FXA_ACCOUNT_SERVER_URL).login(
email,
password,
unblock_code=unblock_code,
)
token = oauth.Client(
client_id=FXA_CLIENT_ID,
server_url=FXA_OAUTH_SERVER_URL,
).authorize_token(session, " ".join(FXA_SCOPES))
except ClientError:
# Two factor is required, try again.
session = core.Client(server_url=FXA_ACCOUNT_SERVER_URL).login(
email,
password,
unblock_code=unblock_code,
verification_method="totp-2fa",
)
print(
"\nTwo-factor authorization is enabled, open your app and enter the code:"
)
totp_code = input("Code: ").strip()
if not session.totp_verify(totp_code):
print("Login failed: invalid two-factor code.")
raise ClientError("Invalid two-factor code.")
token = oauth.Client(
client_id=FXA_CLIENT_ID,
server_url=FXA_OAUTH_SERVER_URL,
).authorize_token(session, " ".join(FXA_SCOPES))
except Exception as retry_exception:
print(f"Login failed: {retry_exception}")
raise retry_exception
except Exception as exception:
print(f"Login failed: {exception}")
raise exception
return token
def _get_mlpa_token():
project_id = "moz-fx-llm-proxy-nonprod-fd0e"
environment = "stage"
secret_name = f"{environment}-gke-app-secrets"
url = (
f"/projects/{project_id}/secrets/{secret_name}/versions/latest:access"
)
access_token = subprocess.check_output(
["gcloud", "auth", "print-access-token"], text=True
).strip()
req = urllib.request.Request(
url, headers={"Authorization": f"Bearer {access_token}"}
)
with urllib.request.urlopen(req) as resp:
data = json.loads(resp.read())
import base64
payload = json.loads(base64.b64decode(data["payload"]["data"]))
token = payload["mlpa-experimentation-authorization-token"]
return token
class EvalToolsCommand(MachCommandBase):
"""Helper utilities for eval workflows."""
@Command(
"eval-tools",
category="testing",
description="Helper utilities for ML evals.",
)
def eval_tools(self):
print("Eval helper utilities.\n\nRun `./mach eval-tools --help` for details.")
return 0
@SubCommand(
"eval-tools",
"login",
description="Login helper to fetch a bearer token (interactive).",
)
def eval_tools_login(self):
if ON_TRY:
print("This command is not intended to be run in try/CI environments.")
return 1
if os.environ.get("MOZ_FXA_BEARER_TOKEN"):
print("MOZ_FXA_BEARER_TOKEN already set; skipping login.")
print("Unset with:\nunset MOZ_FXA_BEARER_TOKEN")
fxa_token = os.environ.get("MOZ_FXA_BEARER_TOKEN")
else:
try:
fxa_token = _get_fxa_token(self)
except Exception as exception:
print(f"Error during FXA login: {exception}")
return 1
if os.environ.get("MOZ_MLPA_AUTHORIZATION_TOKEN"):
print(
"MOZ_MLPA_AUTHORIZATION_TOKEN already set; skipping MLPA token fetch."
)
print("Unset with: unset MOZ_MLPA_AUTHORIZATION_TOKEN")
mlpa_token = os.environ.get("MOZ_MLPA_AUTHORIZATION_TOKEN")
else:
try:
mlpa_token = _get_mlpa_token()
except Exception as exception:
print(f"Error fetching MLPA token: {exception}")
return 1
# don't print tokens in CI/try to avoid leaking secrets in logs
if not ON_TRY:
print(
"\nCopy and paste the following in your terminal to persist your login:\n"
)
print(f"export MOZ_FXA_BEARER_TOKEN='{fxa_token}'")
print(f"export MOZ_MLPA_AUTHORIZATION_TOKEN='{mlpa_token}'")
return 0