Skip to content

Commit

Permalink
feat!: update display mode management
Browse files Browse the repository at this point in the history
  • Loading branch information
tysmith committed Dec 30, 2024
1 parent c785338 commit c2423e2
Show file tree
Hide file tree
Showing 6 changed files with 176 additions and 149 deletions.
57 changes: 17 additions & 40 deletions src/ffpuppet/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,13 +26,11 @@

CREATE_SUSPENDED = 0x00000004

with suppress(ImportError):
from xvfbwrapper import Xvfb

from typing import TYPE_CHECKING

from .bootstrapper import Bootstrapper
from .checks import CheckLogContents, CheckLogSize, CheckMemoryUsage
from .display import DISPLAYS, DisplayMode
from .exceptions import BrowserExecutionError, InvalidPrefs, LaunchError
from .helpers import prepare_environment, wait_on_files
from .minidump_parser import MDSW_URL, MinidumpParser
Expand Down Expand Up @@ -87,13 +85,6 @@ class Reason(IntEnum):
class FFPuppet:
"""FFPuppet manages launching and monitoring the browser process(es).
This includes setting up the environment, collecting logs and some debugger support.
Attributes:
debugger: Debugger to use.
headless: Headless mode to use.
use_profile: Path to existing user profile.
use_xvfb: Use Xvfb (DEPRECATED).
working_path: Path to use as base directory for temporary files.
"""

LAUNCH_TIMEOUT_MIN = 10 # minimum amount of time to wait for the browser to launch
Expand All @@ -104,13 +95,12 @@ class FFPuppet:
"_bin_path",
"_checks",
"_dbg",
"_headless",
"_display",
"_launches",
"_logs",
"_proc_tree",
"_profile_template",
"_working_path",
"_xvfb",
"marionette",
"profile",
"reason",
Expand All @@ -119,42 +109,34 @@ class FFPuppet:
def __init__(
self,
debugger: Debugger = Debugger.NONE,
headless: str | None = None,
display_mode: DisplayMode = DisplayMode.DEFAULT,
use_profile: Path | None = None,
use_xvfb: bool = False,
working_path: str | None = None,
) -> None:
"""
Args:
debugger: Debugger to use.
display_mode: Display mode to use.
use_profile: Path to existing profile to use.
working_path: Path to use as base directory for temporary files.
"""
# tokens used to notify log scanner to kill the browser process
self._abort_tokens: set[Pattern[str]] = set()
self._bin_path: Path | None = None
self._checks: list[CheckLogContents | CheckLogSize | CheckMemoryUsage] = []
self._dbg = debugger
self._dbg_sanity_check(self._dbg)
self._headless = headless
self._launches = 0 # number of successful browser launches
self._logs = PuppetLogger(base_path=working_path)
self._display = DISPLAYS[display_mode]()
# number of successful browser launches
self._launches = 0
self._proc_tree: ProcessTree | None = None
self._profile_template = use_profile
self._xvfb: Xvfb | None = None
self._working_path = working_path
self.marionette: int | None = None
self.profile: Profile | None = None
self.reason: Reason | None = Reason.CLOSED

if use_xvfb:
self._headless = "xvfb"

if self._headless == "xvfb":
try:
self._xvfb = Xvfb(width=1280, height=1024)
except NameError:
self._logs.clean_up(ignore_errors=True)
raise OSError(
"Please install xvfbwrapper (Only supported on Linux)"
) from None
self._xvfb.start()
else:
assert self._headless in (None, "default")
self._logs = PuppetLogger(base_path=working_path)

def __enter__(self) -> FFPuppet:
return self
Expand Down Expand Up @@ -322,8 +304,7 @@ def build_launch_cmd(
if bin_path.lower().endswith(".py"):
cmd.append(executable)
cmd += [bin_path, "-new-instance"]
if self._headless == "default":
cmd.append("-headless")
cmd.extend(self._display.args)
if self.profile is not None:
cmd += ["-profile", str(self.profile)]

Expand Down Expand Up @@ -437,16 +418,13 @@ def clean_up(self) -> None:
Returns:
None
"""
self._display.close()
if self._launches < 0:
LOG.debug("clean_up() call ignored")
return
LOG.debug("clean_up() called")
self.close(force_close=True)
self._logs.clean_up(ignore_errors=True)
# close Xvfb
if self._xvfb is not None:
self._xvfb.stop()
self._xvfb = None
# at this point everything should be cleaned up
assert self.reason is not None
assert self._logs.closed
Expand Down Expand Up @@ -738,8 +716,7 @@ def launch(
# https://developer.gimp.org/api/2.0/glib/glib-running.html#G_DEBUG
env_mod["G_DEBUG"] = "gc-friendly"
env_mod["MOZ_CRASHREPORTER_DISABLE"] = "1"
if self._headless == "xvfb":
env_mod["MOZ_ENABLE_WAYLAND"] = "0"
env_mod.update(self._display.env)

# create a profile
self.profile = Profile(
Expand Down
101 changes: 101 additions & 0 deletions src/ffpuppet/display.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,101 @@
# 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/.
"""ffpuppet display module"""

from __future__ import annotations

from contextlib import suppress
from enum import Enum, auto, unique
from logging import getLogger
from platform import system
from types import MappingProxyType

with suppress(ImportError):
from xvfbwrapper import Xvfb

from typing import TYPE_CHECKING

if TYPE_CHECKING:
from collections.abc import Iterable, Mapping


LOG = getLogger(__name__)


@unique
class DisplayMode(Enum):
"""Supported display modes."""

DEFAULT = auto()
HEADLESS = auto()
if system() == "Linux":
XVFB = auto()


class Display:
"""Default display mode.
Attributes:
args: Extra Firefox command line arguments to use.
env: Extra environment variables to use.
mode: DisplayMode enum name.
"""

__slots__ = ("args", "env", "mode")

def __init__(self) -> None:
self.args: Iterable[str] = ()
self.env: Mapping[str, str] = MappingProxyType({})
self.mode: str = DisplayMode.DEFAULT.name

def close(self) -> None:
"""Perform any required operations to shutdown and cleanup.
Args:
None
Returns:
None
"""


class HeadlessDisplay(Display):
"""Headless display mode."""

def __init__(self) -> None:
super().__init__()
self.args = ("-headless",)
self.mode = DisplayMode.HEADLESS.name


class XvfbDisplay(Display):
"""Xvfb display mode."""

__slots__ = ("_xvfb",)

def __init__(self) -> None:
super().__init__()
self.env = MappingProxyType({"MOZ_ENABLE_WAYLAND": "0"})
self.mode = DisplayMode.XVFB.name
try:
self._xvfb: Xvfb | None = Xvfb(width=1280, height=1024)
except NameError:
LOG.error("Missing xvfbwrapper")
raise
self._xvfb.start()

def close(self) -> None:
if self._xvfb is not None:
self._xvfb.stop()
self._xvfb = None


_displays: dict[DisplayMode, type[Display]] = {
DisplayMode.DEFAULT: Display,
DisplayMode.HEADLESS: HeadlessDisplay,
}
if system() == "Linux":
_displays[DisplayMode.XVFB] = XvfbDisplay

DISPLAYS = MappingProxyType(_displays)
37 changes: 9 additions & 28 deletions src/ffpuppet/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@

from .bootstrapper import Bootstrapper
from .core import Debugger, FFPuppet, Reason
from .display import DisplayMode
from .exceptions import BrowserExecutionError
from .helpers import certutil_available, certutil_find
from .profile import Profile
Expand Down Expand Up @@ -116,6 +117,12 @@ def parse_args(argv: list[str] | None = None) -> Namespace:
type=Path,
help="Install trusted certificates.",
)
cfg_group.add_argument(
"--display",
choices=sorted(x.name.lower() for x in DisplayMode),
default=DisplayMode.DEFAULT.name,
help="Display mode.",
)
cfg_group.add_argument(
"-e",
"--extension",
Expand All @@ -124,17 +131,6 @@ def parse_args(argv: list[str] | None = None) -> Namespace:
help="Install extensions. Specify the path to the xpi or the directory "
"containing the unpacked extension.",
)
headless_choices = ["default"]
if system() == "Linux":
headless_choices.append("xvfb")
cfg_group.add_argument(
"--headless",
choices=headless_choices,
const="default",
default=None,
nargs="?",
help="Headless mode. 'default' uses browser's built-in headless mode.",
)
cfg_group.add_argument(
"--marionette",
const=0,
Expand All @@ -160,14 +156,6 @@ def parse_args(argv: list[str] | None = None) -> Namespace:
cfg_group.add_argument(
"-u", "--url", help="Server URL or path to local file to load."
)
if system() == "Linux":
cfg_group.add_argument(
"--xvfb",
action="store_true",
help="DEPRECATED! Please use '--headless xvfb'",
)
else:
cfg_group.set_defaults(xvfb=False)

report_group = parser.add_argument_group("Issue Detection & Reporting")
report_group.add_argument(
Expand Down Expand Up @@ -290,19 +278,12 @@ def parse_args(argv: list[str] | None = None) -> Namespace:
args.memory *= 1_048_576
if args.prefs is not None and not args.prefs.is_file():
parser.error(f"Invalid prefs.js file '{args.prefs}'")
if args.xvfb:
LOG.warning("'--xvfb' is DEPRECATED. Please use '--headless xvfb'")
args.headless = "xvfb"

return args


def main(argv: list[str] | None = None) -> None:
"""
FFPuppet main entry point
Run with --help for usage
"""
"""FFPuppet main entry point."""
args = parse_args(argv)
# set output verbosity
if args.log_level == DEBUG:
Expand All @@ -315,7 +296,7 @@ def main(argv: list[str] | None = None) -> None:

ffp = FFPuppet(
debugger=args.debugger,
headless=args.headless,
display_mode=DisplayMode[args.display.upper()],
use_profile=args.profile,
)
for a_token in args.abort_token:
Expand Down
28 changes: 28 additions & 0 deletions src/ffpuppet/test_display.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
# 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/.
"""display.py tests"""
from platform import system

from pytest import mark, raises

from .display import DISPLAYS, DisplayMode, XvfbDisplay


@mark.parametrize("mode", tuple(x for x in DisplayMode))
def test_displays(mocker, mode):
"""test Displays()"""
if system() == "Linux":
mocker.patch("ffpuppet.display.Xvfb", autospec=True)
display = DISPLAYS[mode]()
assert display
assert display.mode == mode.name
display.close()


@mark.skipif(system() != "Linux", reason="Only supported on Linux")
def test_xvfb_missing_deps(mocker):
"""test XvfbDisplay() missing deps"""
mocker.patch("ffpuppet.display.Xvfb", side_effect=NameError("test"))
with raises(NameError):
XvfbDisplay()
Loading

0 comments on commit c2423e2

Please sign in to comment.