From d6f9e35e0a243578b4ce989c75635abd2bbfe443 Mon Sep 17 00:00:00 2001 From: Tyson Smith Date: Wed, 17 Apr 2024 11:00:24 -0700 Subject: [PATCH] Add type hints to target --- grizzly/common/runner.py | 10 +- grizzly/common/test_runner.py | 8 +- grizzly/target/assets.py | 57 +++--- grizzly/target/puppet_target.py | 136 +++++++------ grizzly/target/target.py | 277 ++++++++++++++++++++------ grizzly/target/target_monitor.py | 80 ++++++-- grizzly/target/test_assets.py | 12 +- grizzly/target/test_puppet_target.py | 46 +++-- grizzly/target/test_target.py | 20 +- grizzly/target/test_target_monitor.py | 5 +- 10 files changed, 438 insertions(+), 213 deletions(-) diff --git a/grizzly/common/runner.py b/grizzly/common/runner.py index 572b23d3..a449bce9 100644 --- a/grizzly/common/runner.py +++ b/grizzly/common/runner.py @@ -95,7 +95,9 @@ def __init__( if idle_threshold > 0: assert idle_delay > 0 LOG.debug("using idle check, th %d, delay %ds", idle_threshold, idle_delay) - self._idle = _IdleChecker(target.is_idle, idle_threshold, idle_delay) + self._idle = _IdleChecker( + target.monitor.is_idle, idle_threshold, idle_delay + ) else: self._idle = None assert close_delay > 0 @@ -332,13 +334,13 @@ def run( LOG.debug("relaunch/shutdown limit hit") # ideally all browser tabs should be closed at this point # and the browser should exit on its own - # NOTE: this will take the full duration if target.is_idle() - # is not implemented + # NOTE: this will take the full duration if target.monitor.is_idle() + # is unable to detect if the target is idle for close_delay in range(max(int(self._close_delay / 0.5), 1)): if not self._target.monitor.is_healthy(): break # wait 3 seconds (6 passes) before attempting idle exit - if close_delay > 5 and self._target.is_idle(10): + if close_delay > 5 and self._target.monitor.is_idle(10): # NOTE: this will always trigger on systems where the # browser does not exit when the last window is closed LOG.debug("target idle") diff --git a/grizzly/common/test_runner.py b/grizzly/common/test_runner.py index f2465e72..63faf8e1 100644 --- a/grizzly/common/test_runner.py +++ b/grizzly/common/test_runner.py @@ -75,7 +75,7 @@ def test_runner_02(mocker): required=serv_files, ) # single run/iteration relaunch (not idle exit) - target.is_idle.return_value = False + target.monitor.is_idle.return_value = False runner = Runner(server, target, relaunch=1) assert runner._relaunch == 1 smap = ServerMap() @@ -83,7 +83,7 @@ def test_runner_02(mocker): assert runner.initial assert result.attempted assert target.close.call_count == 1 - assert target.is_idle.call_count > 0 + assert target.monitor.is_idle.call_count > 0 assert target.monitor.is_healthy.call_count > 0 assert result.status == Result.NONE assert result.served == serv_files @@ -93,7 +93,7 @@ def test_runner_02(mocker): target.reset_mock() testcase.reset_mock() # single run/iteration relaunch (idle exit) - target.is_idle.return_value = True + target.monitor.is_idle.return_value = True runner = Runner(server, target, relaunch=1) assert runner._relaunch == 1 result = runner.run([], ServerMap(), testcase) @@ -122,7 +122,7 @@ def test_runner_02(mocker): assert not runner.initial assert result.attempted assert target.close.call_count == 1 - assert target.is_idle.call_count == 0 + assert target.monitor.is_idle.call_count == 0 assert target.monitor.is_healthy.call_count == 1 assert result.status == Result.NONE assert result.served == serv_files diff --git a/grizzly/target/assets.py b/grizzly/target/assets.py index 341a74b0..026de2d6 100644 --- a/grizzly/target/assets.py +++ b/grizzly/target/assets.py @@ -5,6 +5,7 @@ from pathlib import Path from shutil import copyfile, copytree, move, rmtree from tempfile import mkdtemp +from typing import Any, Dict, List, Optional __all__ = ("AssetError", "AssetManager") __author__ = "Tyson Smith" @@ -20,30 +21,28 @@ class AssetError(Exception): class AssetManager: __slots__ = ("assets", "path") - def __init__(self, base_path=None): - self.assets = {} + def __init__(self, base_path: Optional[Path] = None) -> None: + self.assets: Dict[str, str] = {} self.path = Path(mkdtemp(prefix="assets_", dir=base_path)) - def __enter__(self): + def __enter__(self) -> "AssetManager": return self - def __exit__(self, *exc): + def __exit__(self, *exc: Any) -> None: self.cleanup() - def add(self, asset, path, copy=True): + def add(self, asset: str, path: Path, copy: bool = True) -> Path: """Add asset to the AssetManager. Args: - asset (str): Name of asset. - path (Path): Location on disk. - copy (bool): Copy or move the content. + asset: Name of asset. + path: Location on disk. + copy: Copy or move the content. Returns: - str: Path to the asset on the filesystem. + Path to the asset on the filesystem. """ - assert isinstance(asset, str) - assert isinstance(path, Path) - assert self.path, "cleanup() was called" + assert asset if not path.exists(): raise OSError(f"'{path}' does not exist") dst = self.path / path.name @@ -66,11 +65,11 @@ def add(self, asset, path, copy=True): LOG.debug("%s asset %r to '%s'", "copied" if copy else "moved", asset, dst) return dst - def add_batch(self, assets): + def add_batch(self, assets: List[List[str]]) -> None: """Add collection of assets to the AssetManager. Args: - assets (list(list(str, str))): List of list that contain asset, path pairs. + assets: List of list that contain asset, path pairs. Returns: None @@ -78,7 +77,7 @@ def add_batch(self, assets): for asset, path in assets: self.add(asset, Path(path)) - def cleanup(self): + def cleanup(self) -> None: """Remove asset files from filesystem. Args: @@ -90,54 +89,54 @@ def cleanup(self): if self.path: rmtree(self.path, ignore_errors=True) self.assets.clear() - self.path = None - def get(self, asset): + def get(self, asset: str) -> Optional[Path]: """Get path to content on filesystem for given asset. Args: - asset (str): Asset to lookup. + asset: Asset to lookup. Returns: - Path: Path to asset content or None if asset does not exist. + Path to asset content or None if asset does not exist. """ item = self.assets.get(asset, None) return self.path / item if item else None - def is_empty(self): + def is_empty(self) -> bool: """Check if AssetManager contains entries. Args: None Returns: - bool: True if AssetManager contains entries else False. + True if AssetManager contains entries else False. """ return not self.assets @classmethod - def load(cls, assets, src_path, base_path=None): + def load( + cls, assets: Dict[str, str], src_path: Path, base_path: Optional[Path] = None + ) -> "AssetManager": """Load assets from filesystem. Args: - asset (dict): Asset paths on filesystem relative to src_path, keyed on - asset name. - src_path (Path): Path to scan for assets. - base_path (str): Base path to use to create local storage. + asset: Asset paths on filesystem relative to src_path, keyed on asset name. + src_path: Path to scan for assets. + base_path: Base path to use to create local storage. Returns: - AssetManager: Populated with contents provided by assets argument. + AssetManager populated with contents provided by assets argument. """ obj = cls(base_path=base_path) for asset, src_name in assets.items(): obj.add(asset, src_path / src_name) return obj - def remove(self, asset): + def remove(self, asset: str) -> None: """Remove asset from AssetManager if asset exists. Args: - asset (str): Asset to remove. + asset: Asset to remove. Returns: None diff --git a/grizzly/target/puppet_target.py b/grizzly/target/puppet_target.py index dacbc5cf..fecbb1b7 100644 --- a/grizzly/target/puppet_target.py +++ b/grizzly/target/puppet_target.py @@ -6,14 +6,17 @@ from os import kill from pathlib import Path from platform import system -from signal import SIGABRT +from signal import SIGABRT, Signals +from tempfile import TemporaryDirectory, mkdtemp +from time import sleep, time +from typing import Dict, Optional, Set, cast try: - from signal import SIGUSR1 + from signal import SIGUSR1 # pylint: disable=ungrouped-imports + + COVERAGE_SIG: Optional[Signals] = SIGUSR1 except ImportError: - SIGUSR1 = None -from tempfile import TemporaryDirectory, mkdtemp -from time import sleep, time + COVERAGE_SIG = None from ffpuppet import BrowserTimeoutError, Debugger, FFPuppet, LaunchError, Reason from ffpuppet.helpers import certutil_available, certutil_find @@ -21,7 +24,7 @@ from prefpicker import PrefPicker from psutil import AccessDenied, NoSuchProcess, Process, process_iter, wait_procs -from ..common.reporter import Report +from ..common.report import Report from ..common.utils import grz_tmp from .target import Result, Target, TargetLaunchError, TargetLaunchTimeout from .target_monitor import TargetMonitor @@ -34,24 +37,31 @@ class PuppetMonitor(TargetMonitor): - def __init__(self, puppet): + def __init__(self, puppet: FFPuppet) -> None: self._puppet = puppet - def clone_log(self, log_id, offset=0): + def clone_log(self, log_id: str, offset: int = 0) -> Optional[Path]: return self._puppet.clone_log(log_id, offset=offset) - def is_running(self): - return self._puppet.is_running() - - def is_healthy(self): + def is_healthy(self) -> bool: return self._puppet.is_healthy() + def is_idle(self, threshold: int) -> bool: + # assert 0 <= threshold <= 100 + for _, cpu in self._puppet.cpu_usage(): + if cpu >= threshold: + return False + return True + + def is_running(self) -> bool: + return self._puppet.is_running() + @property - def launches(self): + def launches(self) -> int: return self._puppet.launches - def log_length(self, log_id): - return self._puppet.log_length(log_id) + def log_length(self, log_id: str) -> int: + return self._puppet.log_length(log_id) or 0 class PuppetTarget(Target): @@ -84,7 +94,14 @@ class PuppetTarget(Target): __slots__ = ("use_valgrind", "_debugger", "_extension", "_prefs", "_puppet") - def __init__(self, binary, launch_timeout, log_limit, memory_limit, **kwds): + def __init__( + self, + binary: Path, + launch_timeout: int, + log_limit: int, + memory_limit: int, + **kwds, + ) -> None: certs = kwds.pop("certs", None) # only pass certs to FFPuppet if certutil is available # otherwise certs can't be used @@ -110,8 +127,8 @@ def __init__(self, binary, launch_timeout, log_limit, memory_limit, **kwds): if kwds.pop("valgrind", False): self.use_valgrind = True self._debugger = Debugger.VALGRIND - self._extension = None - self._prefs = None + self._extension: Optional[Path] = None + self._prefs: Optional[Path] = None # create Puppet object self._puppet = FFPuppet( @@ -122,26 +139,26 @@ def __init__(self, binary, launch_timeout, log_limit, memory_limit, **kwds): if kwds: LOG.debug("PuppetTarget ignoring unsupported kwargs: %s", ", ".join(kwds)) - def _cleanup(self): + def _cleanup(self) -> None: # prevent parallel calls to FFPuppet.close() and/or FFPuppet.clean_up() with self._lock: self._puppet.clean_up() - def close(self, force_close=False): + def close(self, force_close: bool = False) -> None: # prevent parallel calls to FFPuppet.close() and/or FFPuppet.clean_up() with self._lock: self._puppet.close(force_close=force_close) @property - def closed(self): + def closed(self) -> bool: return self._puppet.reason is not None - def create_report(self, is_hang=False): + def create_report(self, is_hang: bool = False) -> Report: logs = Path(mkdtemp(prefix="logs_", dir=grz_tmp("logs"))) self.save_logs(logs) return Report(logs, self.binary, is_hang=is_hang) - def filtered_environ(self): + def filtered_environ(self) -> Dict[str, str]: # remove context specific entries from environment filtered = dict(self.environ) opts = SanitizerOptions() @@ -157,19 +174,13 @@ def filtered_environ(self): # remove empty entries return {k: v for k, v in filtered.items() if v} - def is_idle(self, threshold): - for _, cpu in self._puppet.cpu_usage(): - if cpu >= threshold: - return False - return True - @property - def monitor(self): + def monitor(self) -> PuppetMonitor: if self._monitor is None: self._monitor = PuppetMonitor(self._puppet) - return self._monitor + return cast(PuppetMonitor, self._monitor) - def check_result(self, ignored): + def check_result(self, ignored: Set[str]) -> Result: result = Result.NONE # check if there has been a crash, hangs will appear as SIGABRT if not self._puppet.is_healthy(): @@ -194,12 +205,15 @@ def check_result(self, ignored): result = Result.IGNORED LOG.debug("log size limit exceeded") else: + assert self._puppet.reason # crash or hang (forced SIGABRT) has been detected LOG.debug("result detected (%s)", self._puppet.reason.name) result = Result.FOUND return result - def handle_hang(self, ignore_idle=True, ignore_timeout=False): + def handle_hang( + self, ignore_idle: bool = True, ignore_timeout: bool = False + ) -> bool: # only send SIGABRT in certain case send_abort = ( not ignore_timeout @@ -229,15 +243,15 @@ def handle_hang(self, ignore_idle=True, ignore_timeout=False): self.close() return was_idle - def https(self): + def https(self) -> bool: return self._https - def dump_coverage(self, timeout=5): + def dump_coverage(self, timeout: int = 5) -> None: if system() != "Linux": LOG.debug("dump_coverage() only supported on Linux") return - assert SIGUSR1 is not None + assert COVERAGE_SIG is not None pid = self._puppet.get_pid() if pid is None or not self._puppet.is_healthy(): LOG.debug("Skipping coverage dump (target is not in a good state)") @@ -245,30 +259,31 @@ def dump_coverage(self, timeout=5): # If at this point, the browser is in a good state, i.e. no crashes # or hangs, so signal the browser to dump coverage. running_procs = 0 - signaled_pids = [] + signaled_pids: Set[int] = set() try: - # send SIGUSR1 to browser processes + # send COVERAGE_SIG (SIGUSR1) to browser processes + # TODO: this should use FFPuppet.processes() parent_proc = Process(pid) for proc in chain([parent_proc], parent_proc.children(recursive=True)): - # avoid sending SIGUSR1 to non-browser processes + # avoid sending signal to non-browser processes if Path(proc.exe()).name.startswith("firefox"): LOG.debug( - "Sending SIGUSR1 to %d (%s)", + "Sending signal to %d (%s)", proc.pid, "parent" if proc.pid == pid else "child", ) try: - kill(proc.pid, SIGUSR1) - signaled_pids.append(proc.pid) + kill(proc.pid, COVERAGE_SIG) + signaled_pids.add(proc.pid) except OSError: - LOG.warning("Failed to send SIGUSR1 to pid %d", proc.pid) + LOG.warning("Failed to send signal to pid %d", proc.pid) if proc.is_running(): running_procs += 1 except (AccessDenied, NoSuchProcess): # pragma: no cover pass if not signaled_pids: LOG.warning( - "SIGUSR1 not sent, no browser processes found (%d process(es) running)", + "Signal not sent, no browser processes found (%d process(es) running)", running_procs, ) return @@ -324,11 +339,12 @@ def dump_coverage(self, timeout=5): break sleep(delay) - def launch(self, location): + def launch(self, location: str) -> None: # setup environment env_mod = dict(self.environ) # do not allow network connections to non local endpoints env_mod["MOZ_DISABLE_NONLOCAL_CONNECTIONS"] = "1" + # we always want the browser to exit when a crash is detected env_mod["MOZ_CRASHREPORTER_SHUTDOWN"] = "1" try: self._puppet.launch( @@ -338,8 +354,8 @@ def launch(self, location): log_limit=self.log_limit, memory_limit=self.memory_limit, prefs_js=self._prefs, - extension=self._extension, - env_mod=env_mod, + extension=[self._extension] if self._extension else None, + env_mod=cast(Dict[str, Optional[str]], env_mod), cert_files=[self.certs.root] if self.certs else None, ) except LaunchError as exc: @@ -348,10 +364,15 @@ def launch(self, location): raise TargetLaunchTimeout(str(exc)) from None raise TargetLaunchError(str(exc), self.create_report()) from None - def log_size(self): - return self._puppet.log_length("stderr") + self._puppet.log_length("stdout") + def log_size(self) -> int: + total = 0 + for log in ("stderr", "stdout"): + length = self._puppet.log_length(log) + if length: + total += length + return total - def merge_environment(self, extra): + def merge_environment(self, extra: Dict[str, str]) -> None: output = dict(extra) if self.environ: # prioritize existing environment variables @@ -360,7 +381,7 @@ def merge_environment(self, extra): org = SanitizerOptions() out = SanitizerOptions() for san in ("ASAN", "LSAN", "TSAN", "UBSAN"): - opts = "_".join((san, "OPTIONS")) + opts = f"{san}_OPTIONS" org.load_options(self.environ.get(opts, "")) if not org: # nothing to add from original @@ -371,7 +392,7 @@ def merge_environment(self, extra): output[opts] = str(out) self.environ = output - def process_assets(self): + def process_assets(self) -> None: self._extension = self.asset_mgr.get("extension") self._prefs = self.asset_mgr.get("prefs") # generate temporary prefs.js with prefpicker @@ -380,11 +401,12 @@ def process_assets(self): with TemporaryDirectory(dir=grz_tmp("target")) as tmp_path: prefs = Path(tmp_path) / "prefs.js" template = PrefPicker.lookup_template("browser-fuzzing.yml") + assert template PrefPicker.load_template(template).create_prefsjs(prefs) self._prefs = self.asset_mgr.add("prefs", prefs, copy=False) abort_tokens = self.asset_mgr.get("abort-tokens") if abort_tokens: - LOG.debug("loading 'abort tokens' from %r", abort_tokens) + LOG.debug("loading 'abort tokens' from '%s'", abort_tokens) with (self.asset_mgr.path / abort_tokens).open() as in_fp: for line in in_fp: line = line.strip() @@ -404,7 +426,9 @@ def process_assets(self): "suppressions", f"'{self.asset_mgr.get(asset)}'", overwrite=True ) elif opts.get("suppressions"): - path = Path(opts.pop("suppressions").strip("\"'")) + suppressions = opts.pop("suppressions") + assert suppressions + path = Path(suppressions.strip("\"'")) if path.is_file(): # use environment specified suppression file LOG.debug("using %r from environment", asset) @@ -422,5 +446,5 @@ def process_assets(self): LOG.debug("updating suppressions in %r", var_name) self.environ[var_name] = str(opts) - def save_logs(self, *args, **kwargs): - self._puppet.save_logs(*args, **kwargs) + def save_logs(self, dst: Path) -> None: + self._puppet.save_logs(dst) diff --git a/grizzly/target/target.py b/grizzly/target/target.py index f449e972..2847c018 100644 --- a/grizzly/target/target.py +++ b/grizzly/target/target.py @@ -2,13 +2,17 @@ # 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/. from abc import ABCMeta, abstractmethod -from enum import Enum, unique +from enum import IntEnum, unique from logging import getLogger from os import environ +from pathlib import Path from threading import Lock +from typing import Any, Dict, Optional, Set, Tuple +from ..common.report import Report from ..common.utils import CertificateBundle, grz_tmp from .assets import AssetManager +from .target_monitor import TargetMonitor __all__ = ("Result", "Target", "TargetError", "TargetLaunchError") __author__ = "Tyson Smith" @@ -18,8 +22,8 @@ @unique -class Result(Enum): - """Target results codes""" +class Result(IntEnum): + """Target result codes""" NONE = 0 FOUND = 1 @@ -33,7 +37,7 @@ class TargetError(Exception): class TargetLaunchError(TargetError): """Raised if a failure during launch occurs""" - def __init__(self, message, report): + def __init__(self, message: str, report: Report) -> None: super().__init__(message) self.report = report @@ -43,8 +47,8 @@ class TargetLaunchTimeout(TargetError): class Target(metaclass=ABCMeta): - SUPPORTED_ASSETS = None - TRACKED_ENVVARS = () + SUPPORTED_ASSETS: Tuple[str, ...] = () + TRACKED_ENVVARS: Tuple[str, ...] = () __slots__ = ( "_asset_mgr", @@ -61,12 +65,12 @@ class Target(metaclass=ABCMeta): def __init__( self, - binary, - launch_timeout, - log_limit, - memory_limit, - certs=None, - ): + binary: Path, + launch_timeout: int, + log_limit: int, + memory_limit: int, + certs: Optional[CertificateBundle] = None, + ) -> None: assert launch_timeout > 0 assert log_limit >= 0 assert memory_limit >= 0 @@ -75,7 +79,7 @@ def __init__( self._asset_mgr = AssetManager(base_path=grz_tmp("target")) self._https = False self._lock = Lock() - self._monitor = None + self._monitor: Optional[TargetMonitor] = None self.binary = binary self.certs = certs self.environ = self.scan_environment(dict(environ), self.TRACKED_ENVVARS) @@ -83,103 +87,246 @@ def __init__( self.log_limit = log_limit self.memory_limit = memory_limit - def __enter__(self): + def __enter__(self) -> "Target": return self - def __exit__(self, *exc): + def __exit__(self, *exc: Any) -> None: self.cleanup() @property - def asset_mgr(self): + def asset_mgr(self) -> AssetManager: + """Get current AssetManager. + + Args: + None + + Returns: + AssetManager. + """ return self._asset_mgr @asset_mgr.setter - def asset_mgr(self, asset_mgr): + def asset_mgr(self, asset_mgr: AssetManager) -> None: + """Set AssetManager and cleanup previous AssetManager. + + Args: + None + + Returns: + AssetManager. + """ self._asset_mgr.cleanup() assert isinstance(asset_mgr, AssetManager) self._asset_mgr = asset_mgr @abstractmethod - def _cleanup(self): - pass + def _cleanup(self) -> None: + """Cleanup method to be implemented by subclass. + + Args: + None + + Returns: + None. + """ @abstractmethod - def check_result(self, ignored): - pass + def check_result(self, ignored: Set[str]) -> Result: + """Check for results. - def cleanup(self): - # call target specific _cleanup first + Args: + ignored: Result types that are currently ignored. + + Returns: + Result code. + """ + + def cleanup(self) -> None: + """Perform necessary cleanup. DO NOT OVERRIDE. + + Args: + ignored: Types of results to ignore. + + Returns: + Result code. + """ + # call target specific _cleanup method first self._cleanup() self._asset_mgr.cleanup() @abstractmethod - def close(self, force_close=False): - pass + def close(self, force_close: bool = False) -> None: + """Close target. + + Args: + force_close: Close as quickly as possible. Logs will not be collected. + + Returns: + None. + """ @property @abstractmethod - def closed(self): - pass + def closed(self) -> bool: + """Check if the target is closed. + + Args: + None + + Returns: + True if closed otherwise False. + """ + + @abstractmethod + def create_report(self, is_hang: bool = False) -> Report: + """Process logs and create a Report. + + Args: + is_hang: Indicate whether the results is due to a hang/timeout. + + Returns: + Report object. + """ @abstractmethod - def create_report(self, is_hang=False): - pass + def dump_coverage(self, timeout: int = 0) -> None: + """Trigger target coverage data dump. + + Args: + timeout: Amount of time to wait for data to be written. + + Returns: + None. + """ + + def filtered_environ(self) -> Dict[str, str]: + """Used to collect the environment to add to a testcase. - def dump_coverage(self, _timeout=0): - LOG.warning("dump_coverage() is not supported!") + Args: + None - def filtered_environ(self): - """Used to collect the environment to add to a testcase""" + Returns: + Environment variables. + """ return dict(self.environ) @abstractmethod - def handle_hang(self, ignore_idle=True, ignore_timeout=False): - pass + def handle_hang( + self, + ignore_idle: bool = True, + ignore_timeout: bool = False, + ) -> bool: + """Handle a target hang. + + Args: + ignore_idle: Do not treat as a hang if target is idle. + ignore_timeout: Indicates if a timeout will be ignored. + + Returns: + True if the target was idle otherwise False. + """ @abstractmethod - def https(self): - pass + def https(self) -> bool: + """Target configured for HTTPS. - # TODO: move to monitor? - def is_idle(self, _threshold): - LOG.debug("Target.is_idle() not implemented! returning False") - return False + Args: + None + + Returns: + True if HTTPS can be used otherwise False. + """ @abstractmethod - def launch(self, location): - pass + def launch(self, location: str) -> None: + """Launch the target. + + Args: + location: URL to load. - def log_size(self): - LOG.debug("log_size() not implemented! returning 0") - return 0 + Returns: + None. + """ @abstractmethod - def merge_environment(self, extra): - pass + def log_size(self) -> int: + """Calculate the amount of data contained in target log files. + + Args: + None + + Returns: + Total data size of log files in bytes. + """ + + @abstractmethod + def merge_environment(self, extra: Dict[str, str]) -> None: + """Add to existing environment. + + Args: + extra: Environment variables to add. + + Returns: + None. + """ @property @abstractmethod - def monitor(self): - pass + def monitor(self) -> TargetMonitor: + """TargetMonitor. + + Args: + extra: Environment variables to add. + + Returns: + TargetMonitor + """ @abstractmethod - def process_assets(self): - pass + def process_assets(self) -> None: + """Prepare assets for use by the target. + + Args: + None + + Returns: + None. + """ + + def reverse(self, remote: int, local: int) -> None: + """Configure port mappings. Remote -> device, local -> desktop (current system). - def reverse(self, remote, local): - # remote->device, local->desktop - pass + Args: + remote: Port on remote device. + local: Port on local machine. + + Returns: + None. + """ @staticmethod - def scan_environment(to_scan, tracked): - # scan environment for tracked environment variables - env = {} - if tracked: - for var in tracked: - if var in to_scan: - env[var] = to_scan[var] - return env + def scan_environment( + env: Dict[str, str], + include: Optional[Tuple[str, ...]], + ) -> Dict[str, str]: + """Scan environment for tracked environment variables. + + Args: + env: Environment to scan. + include: Variables to include in output. + + Returns: + Tracked variables found in scanned environment. + """ + return {var: env[var] for var in include if var in env} if include else {} @abstractmethod - def save_logs(self, *args, **kwargs): - pass + def save_logs(self, dst: Path) -> None: + """Save logs to specified location. + + Args: + dst: Location to save logs. + + Returns: + None. + """ diff --git a/grizzly/target/target_monitor.py b/grizzly/target/target_monitor.py index 314abe0b..7fd4e1a7 100644 --- a/grizzly/target/target_monitor.py +++ b/grizzly/target/target_monitor.py @@ -2,7 +2,8 @@ # 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/. from abc import ABCMeta, abstractmethod -from os import remove +from pathlib import Path +from typing import Optional __all__ = ("TargetMonitor",) __author__ = "Tyson Smith" @@ -11,32 +12,69 @@ class TargetMonitor(metaclass=ABCMeta): @abstractmethod - def clone_log(self, log_id, offset=0): - pass + def clone_log(self, log_id: str, offset: int = 0) -> Optional[Path]: + """Create a copy of a log. + + Args: + log_id: Log identifier. + offset: Number of bytes to seek into log before copying data. + + Returns: + Copy of specified log. + """ @abstractmethod - def is_healthy(self): - pass + def is_healthy(self) -> bool: + """Check for failures such as assertions, crashes, etc. + + Args: + None + + Returns: + True if no target failures are found otherwise False. + """ + + @abstractmethod + def is_idle(self, threshold: int) -> bool: + """Check if target is idle. + + Args: + threshold: Maximum allowed CPU usage as percentage (per process). + + Returns: + True if CPU usage for all processes is below the threshold otherwise False. + """ @abstractmethod - def is_running(self): - pass + def is_running(self) -> bool: + """Check if target is running. + + Args: + None + + Returns: + True if target is running otherwise False. + """ @property @abstractmethod - def launches(self): - pass - - def log_data(self, log_id, offset=0): - data = None - log_file = self.clone_log(log_id, offset=offset) - if log_file: - try: - data = log_file.read_bytes() - finally: - remove(log_file) - return data + def launches(self) -> int: + """Number of successful target launches. + + Args: + None + + Returns: + Number of successful launches. + """ @abstractmethod - def log_length(self, log_id): - pass + def log_length(self, log_id: str) -> int: + """Calculate the length of a specific log file. + + Args: + log_id: Log identifier. + + Returns: + Log file size in bytes. + """ diff --git a/grizzly/target/test_assets.py b/grizzly/target/test_assets.py index f14af421..cbf5f3b8 100644 --- a/grizzly/target/test_assets.py +++ b/grizzly/target/test_assets.py @@ -45,10 +45,12 @@ def test_asset_manager_01(tmp_path): assert (asset_mgr.path / "example_path/b/2.txt").is_file() assert len(asset_mgr.assets) == 2 # get - assert (asset_mgr.path / "example_2.txt").samefile( - asset_mgr.get("example_file") - ) - assert (asset_mgr.path / "example_path").samefile(asset_mgr.get("example_path")) + example = asset_mgr.get("example_file") + assert example + assert (asset_mgr.path / "example_2.txt").samefile(example) + example = asset_mgr.get("example_path") + assert example + assert (asset_mgr.path / "example_path").samefile(example) # remove directory asset_mgr.remove("example_path") assert len(asset_mgr.assets) == 1 @@ -61,7 +63,7 @@ def test_asset_manager_01(tmp_path): # cleanup asset_mgr.cleanup() assert not asset_mgr.assets - assert asset_mgr.path is None + assert not asset_mgr.path.is_dir() def test_asset_manager_02(tmp_path): diff --git a/grizzly/target/test_puppet_target.py b/grizzly/target/test_puppet_target.py index bc523a93..702ed375 100644 --- a/grizzly/target/test_puppet_target.py +++ b/grizzly/target/test_puppet_target.py @@ -9,8 +9,9 @@ from pytest import mark, raises from ..common.utils import CertificateBundle +from .assets import AssetManager from .puppet_target import PuppetTarget -from .target import AssetManager, Result, TargetLaunchError, TargetLaunchTimeout +from .target import Result, TargetLaunchError, TargetLaunchTimeout def test_puppet_target_01(mocker, tmp_path): @@ -26,7 +27,7 @@ def test_puppet_target_01(mocker, tmp_path): assert target.launch_timeout == 300 assert target.log_limit == 25 assert target.memory_limit == 5000 - assert target.check_result([]) == Result.NONE + assert target.check_result(set()) == Result.NONE assert not target.https() assert target.log_size() == 1124 fake_ffp.return_value.log_length.assert_any_call("stderr") @@ -260,18 +261,6 @@ def test_puppet_target_05(mocker, tmp_path): def test_puppet_target_06(mocker, tmp_path): - """test PuppetTarget.is_idle()""" - fake_ffp = mocker.patch("grizzly.target.puppet_target.FFPuppet", autospec=True) - fake_ffp.return_value.cpu_usage.return_value = [(999, 30), (998, 20), (997, 10)] - fake_file = tmp_path / "fake" - fake_file.touch() - with PuppetTarget(fake_file, 300, 25, 5000) as target: - assert not target.is_idle(0) - assert not target.is_idle(25) - assert target.is_idle(50) - - -def test_puppet_target_07(mocker, tmp_path): """test PuppetTarget.monitor""" fake_ffp = mocker.patch("grizzly.target.puppet_target.FFPuppet", autospec=True) fake_file = tmp_path / "fake" @@ -294,6 +283,18 @@ def test_puppet_target_07(mocker, tmp_path): assert fake_ffp.return_value.clone_log.call_count == 1 +def test_puppet_target_07(mocker, tmp_path): + """test PuppetTarget.monitor.is_idle()""" + fake_ffp = mocker.patch("grizzly.target.puppet_target.FFPuppet", autospec=True) + fake_ffp.return_value.cpu_usage.return_value = [(999, 30), (998, 20), (997, 10)] + fake_file = tmp_path / "fake" + fake_file.touch() + with PuppetTarget(fake_file, 300, 25, 5000) as target: + assert not target.monitor.is_idle(0) + assert not target.monitor.is_idle(25) + assert target.monitor.is_idle(50) + + def test_puppet_target_08(mocker, tmp_path): """test PuppetTarget.process_assets()""" mocker.patch("grizzly.target.puppet_target.FFPuppet", autospec=True) @@ -303,16 +304,19 @@ def test_puppet_target_08(mocker, tmp_path): with PuppetTarget(fake_file, 300, 25, 5000) as target: assert target.asset_mgr.get("prefs") is None target.process_assets() - assert target.asset_mgr.get("prefs").is_file() - assert target.asset_mgr.get("prefs").name == "prefs.js" + asset = target.asset_mgr.get("prefs") + assert asset + assert asset.is_file() + assert asset.name == "prefs.js" # prefs file provided with AssetManager(base_path=tmp_path) as asset_mgr: asset_mgr.add("prefs", fake_file) with PuppetTarget(fake_file, 300, 25, 5000) as target: target.asset_mgr = asset_mgr target.process_assets() - assert target.asset_mgr.get("prefs").is_file() - assert target.asset_mgr.get("prefs").name == "fake" + asset = target.asset_mgr.get("prefs") + assert asset + assert asset.name == "fake" # abort tokens file provided with AssetManager(base_path=tmp_path) as asset_mgr: asset_mgr.add("abort-tokens", fake_file) @@ -323,8 +327,10 @@ def test_puppet_target_08(mocker, tmp_path): assert target._puppet.add_abort_token.call_count == 0 target.asset_mgr = asset_mgr target.process_assets() - assert target.asset_mgr.get("abort-tokens").is_file() - assert target.asset_mgr.get("abort-tokens").name == "fake" + asset = target.asset_mgr.get("abort-tokens") + assert asset + assert asset.is_file() + assert asset.name == "fake" assert target._puppet.add_abort_token.call_count == 2 diff --git a/grizzly/target/test_target.py b/grizzly/target/test_target.py index 1b50efb4..d8849ad2 100644 --- a/grizzly/target/test_target.py +++ b/grizzly/target/test_target.py @@ -19,20 +19,26 @@ def close(self, force_close=False): @property def closed(self): - pass + return True def create_report(self, is_hang=False): pass - def handle_hang(self, ignore_idle=True, ignore_timeout=False): + def dump_coverage(self, timeout=0): pass + def handle_hang(self, ignore_idle=True, ignore_timeout=False): + return False + def https(self): return self._https def launch(self, location): pass + def log_size(self): + return 0 + @property def monitor(self): return self._monitor @@ -43,7 +49,7 @@ def merge_environment(self, extra): def process_assets(self): pass - def save_logs(self, *_args, **_kwargs): + def save_logs(self, dst): pass @@ -59,14 +65,12 @@ def test_target_01(tmp_path): assert target.asset_mgr.path != org_path assert not target.environ assert not target.filtered_environ() - assert not target.is_idle(0) assert target.launch_timeout == 10 assert target.log_size() == 0 assert target.log_limit == 2 assert target.memory_limit == 3 assert target.monitor is None # test stubs - target.dump_coverage() target.reverse(1, 2) @@ -87,6 +91,6 @@ def test_target_02(mocker, tmp_path): def test_target_03(): """test Target.scan_environment()""" - assert not Target.scan_environment({"a": "1"}, []) - assert not Target.scan_environment({}, ["a"]) - assert Target.scan_environment({"a": "1", "b": "2"}, ["a"]) == {"a": "1"} + assert not Target.scan_environment({"a": "1"}, ()) + assert not Target.scan_environment({}, ("a",)) + assert Target.scan_environment({"a": "1", "b": "2"}, ("a",)) == {"a": "1"} diff --git a/grizzly/target/test_target_monitor.py b/grizzly/target/test_target_monitor.py index bd78be76..4dd13c16 100644 --- a/grizzly/target/test_target_monitor.py +++ b/grizzly/target/test_target_monitor.py @@ -16,6 +16,9 @@ def clone_log(self, log_id, offset=0): def is_healthy(self): return True + def is_idle(self, threshold): + return False + def is_running(self): return True @@ -30,7 +33,7 @@ def log_length(self, log_id): test_log = mon.clone_log("test_log", offset=0) assert test_log.is_file() assert mon.is_healthy() + assert not mon.is_idle(0) assert mon.is_running() assert mon.launches == 1 - assert mon.log_data("test_log") == b"test" assert mon.log_length("test_log") == 100