codekingpro/portable-devtools

1"""LangSmith Pytest hooks."""2 3import importlib.util4import json5import logging6import os7import time8from collections import defaultdict9from threading import Lock10from typing import Any11 12import pytest13 14from langsmith import utils as ls_utils15from langsmith.testing._internal import test as ls_test16 17logger = logging.getLogger(__name__)18 19 20def pytest_addoption(parser):21    """Set a boolean flag for LangSmith output.22 23    Skip if --langsmith-output is already defined.24    """25    try:26        # Try to add the option, will raise if it already exists27        group = parser.getgroup("langsmith", "LangSmith")28        group.addoption(29            "--langsmith-output",30            action="store_true",31            default=False,32            help="Use LangSmith output (requires 'rich').",33        )34    except ValueError:35        # Option already exists36        logger.warning(37            "LangSmith output flag cannot be added because it's already defined."38        )39 40 41def _handle_output_args(args):42    """Handle output arguments."""43    if any(opt in args for opt in ["--langsmith-output"]):44        # Only add --quiet if it's not already there45        if not any(a in args for a in ["-qq"]):46            args.insert(0, "-qq")47        # Disable built-in output capturing48        if not any(a in args for a in ["-s", "--capture=no"]):49            args.insert(0, "-s")50 51 52if pytest.__version__.startswith("7."):53 54    def pytest_cmdline_preparse(config, args):55        """Call immediately after command line options are parsed (pytest v7)."""56        _handle_output_args(args)57 58else:59 60    def pytest_load_initial_conftests(args):61        """Handle args in pytest v8+."""62        _handle_output_args(args)63 64 65@pytest.hookimpl(hookwrapper=True)66def pytest_runtest_call(item):67    """Apply LangSmith tracking to tests marked with @pytest.mark.langsmith."""68    marker = item.get_closest_marker("langsmith")69    if marker:70        # Get marker kwargs if any (e.g.,71        # @pytest.mark.langsmith(output_keys=["expected"]))72        kwargs = marker.kwargs if marker else {}73        # Inject experiment_metadata from session-scoped fixture if available74        # and not already explicitly set in the marker kwargs75        request_obj = getattr(item, "_request", None)76        if "experiment_metadata" not in kwargs and request_obj is not None:77            try:78                experiment_metadata = request_obj.getfixturevalue(79                    "langsmith_experiment_metadata"80                )81                if experiment_metadata is not None:82                    kwargs = {**kwargs, "experiment_metadata": experiment_metadata}83            except pytest.FixtureLookupError:84                pass85        # Wrap the test function with our test decorator86        original_func = item.obj87        item.obj = ls_test(**kwargs)(original_func)88        if request_obj is not None and "request" not in item.funcargs:89            item.funcargs["request"] = request_obj90        if request_obj is not None and "request" not in item._fixtureinfo.argnames:91            # Create a new FuncFixtureInfo instance with updated argnames92            item._fixtureinfo = type(item._fixtureinfo)(93                argnames=item._fixtureinfo.argnames + ("request",),94                initialnames=item._fixtureinfo.initialnames,95                names_closure=item._fixtureinfo.names_closure,96                name2fixturedefs=item._fixtureinfo.name2fixturedefs,97            )98    yield99 100 101@pytest.hookimpl102def pytest_report_teststatus(report, config):103    """Remove the short test-status character outputs ("./F")."""104    # The hook normally returns a 3-tuple: (short_letter, verbose_word, color)105    # By returning empty strings, the progress characters won't show.106    if config.getoption("--langsmith-output"):107        return "", "", ""108 109 110class LangSmithPlugin:111    """Plugin for rendering LangSmith results."""112 113    def __init__(self):114        """Initialize."""115        from rich.console import Console  # type: ignore[import-not-found]116        from rich.live import Live  # type: ignore[import-not-found]117 118        self.test_suites = defaultdict(list)119        self.test_suite_urls = {}120 121        self.process_status = {}  # Track process status122        self.status_lock = Lock()  # Thread-safe updates123        self.console = Console()124 125        self.live = Live(126            self.generate_tables(), console=self.console, refresh_per_second=10127        )128        self.live.start()129        self.live.console.print("Collecting tests...")130 131    def pytest_collection_finish(self, session):132        """Call after collection phase is completed and session.items is populated."""133        self.collected_nodeids = set()134        for item in session.items:135            self.collected_nodeids.add(item.nodeid)136 137    def add_process_to_test_suite(self, test_suite, process_id):138        """Group a test case with its test suite."""139        self.test_suites[test_suite].append(process_id)140 141    def update_process_status(self, process_id, status):142        """Update test results."""143        # First update144        if not self.process_status:145            self.live.console.print("Running tests...")146 147        with self.status_lock:148            current_status = self.process_status.get(process_id, {})149            self.process_status[process_id] = _merge_statuses(150                status,151                current_status,152                unpack=["feedback", "inputs", "reference_outputs", "outputs"],153            )154        self.live.update(self.generate_tables())155 156    def pytest_runtest_logstart(self, nodeid):157        """Initialize live display when first test starts."""158        self.update_process_status(nodeid, {"status": "running"})159 160    def generate_tables(self):161        """Generate a collection of tables—one per suite.162 163        Returns a 'Group' object so it can be rendered simultaneously by Rich Live.164        """165        from rich.console import Group166 167        tables = []168        for suite_name in self.test_suites:169            table = self._generate_table(suite_name)170            tables.append(table)171        group = Group(*tables)172        return group173 174    def _generate_table(self, suite_name: str):175        """Generate results table."""176        from rich.table import Table  # type: ignore[import-not-found]177 178        process_ids = self.test_suites[suite_name]179 180        title = f"""Test Suite: [bold]{suite_name}[/bold]181LangSmith URL: [bright_cyan]{self.test_suite_urls[suite_name]}[/bright_cyan]"""  # noqa: E501182        table = Table(title=title, title_justify="left")183        table.add_column("Test")184        table.add_column("Inputs")185        table.add_column("Ref outputs")186        table.add_column("Outputs")187        table.add_column("Status")188        table.add_column("Feedback")189        table.add_column("Duration")190 191        # Test, inputs, ref outputs, outputs col width192        max_status = len("status")193        max_duration = len("duration")194        now = time.time()195        durations = []196        numeric_feedbacks = defaultdict(list)197        # Gather data only for this suite198        suite_statuses = {pid: self.process_status[pid] for pid in process_ids}199        for pid, status in suite_statuses.items():200            duration = status.get("end_time", now) - status.get("start_time", now)201            durations.append(duration)202            for k, v in status.get("feedback", {}).items():203                if isinstance(v, (float, int, bool)):204                    numeric_feedbacks[k].append(v)205            max_duration = max(len(f"{duration:.2f}s"), max_duration)206            max_status = max(len(status.get("status", "queued")), max_status)207 208        passed_count = sum(s.get("status") == "passed" for s in suite_statuses.values())209        failed_count = sum(s.get("status") == "failed" for s in suite_statuses.values())210 211        # You could arrange a row to show the aggregated data—here, in the last column:212        if passed_count + failed_count:213            rate = passed_count / (passed_count + failed_count)214            color = "green" if rate == 1 else "red"215            aggregate_status = f"[{color}]{rate:.0%}[/{color}]"216        else:217            aggregate_status = "Passed: --"218        if durations:219            aggregate_duration = f"{sum(durations) / len(durations):.2f}s"220        else:221            aggregate_duration = "--s"222        if numeric_feedbacks:223            aggregate_feedback = "\n".join(224                f"{k}: {sum(v) / len(v)}" for k, v in numeric_feedbacks.items()225            )226        else:227            aggregate_feedback = "--"228 229        max_duration = max(max_duration, len(aggregate_duration))230        max_dynamic_col_width = (self.console.width - (max_status + max_duration)) // 5231        max_dynamic_col_width = max(max_dynamic_col_width, 8)232 233        for pid, status in suite_statuses.items():234            status_color = {235                "running": "yellow",236                "passed": "green",237                "failed": "red",238                "skipped": "cyan",239            }.get(status.get("status", "queued"), "white")240 241            duration = status.get("end_time", now) - status.get("start_time", now)242            feedback = "\n".join(243                f"{_abbreviate(k, max_len=max_dynamic_col_width)}: {int(v) if isinstance(v, bool) else v}"  # noqa: E501244                for k, v in status.get("feedback", {}).items()245            )246            inputs = _dumps_with_fallback(status.get("inputs", {}))247            reference_outputs = _dumps_with_fallback(248                status.get("reference_outputs", {})249            )250            outputs = _dumps_with_fallback(status.get("outputs", {}))251            table.add_row(252                _abbreviate_test_name(str(pid), max_len=max_dynamic_col_width),253                _abbreviate(inputs, max_len=max_dynamic_col_width),254                _abbreviate(reference_outputs, max_len=max_dynamic_col_width),255                _abbreviate(outputs, max_len=max_dynamic_col_width)[256                    -max_dynamic_col_width:257                ],258                f"[{status_color}]{status.get('status', 'queued')}[/{status_color}]",259                feedback,260                f"{duration:.2f}s",261            )262 263        # Add a blank row or a section separator if you like:264        table.add_row("", "", "", "", "", "", "")265        # Finally, our “footer” row:266        table.add_row(267            "[bold]Averages[/bold]",268            "",269            "",270            "",271            aggregate_status,272            aggregate_feedback,273            aggregate_duration,274        )275 276        return table277 278    def pytest_configure(self, config):279        """Disable warning reporting and show no warnings in output."""280        # Disable general warning reporting281        config.option.showwarnings = False282 283        # Disable warning summary284        reporter = config.pluginmanager.get_plugin("warnings-plugin")285        if reporter:286            reporter.warning_summary = lambda *args, **kwargs: None287 288    def pytest_sessionfinish(self, session):289        """Stop Rich Live rendering at the end of the session."""290        self.live.stop()291        self.live.console.print("\nFinishing up...")292 293 294def pytest_configure(config):295    """Register the 'langsmith' marker."""296    config.addinivalue_line(297        "markers", "langsmith: mark test to be tracked in LangSmith"298    )299    if config.getoption("--langsmith-output"):300        if not importlib.util.find_spec("rich"):301            msg = (302                "Must have 'rich' installed to use --langsmith-output. "303                "Please install with: `pip install -U 'langsmith[pytest]'`"304            )305            raise ValueError(msg)306        if os.environ.get("PYTEST_XDIST_TESTRUNUID"):307            msg = (308                "--langsmith-output not supported with pytest-xdist. "309                "Please remove the '--langsmith-output' option or '-n' option."310            )311            raise ValueError(msg)312        if ls_utils.test_tracking_is_disabled():313            msg = (314                "--langsmith-output not supported when env var"315                "LANGSMITH_TEST_TRACKING='false'. Please remove the"316                "'--langsmith-output' option "317                "or enable test tracking."318            )319            raise ValueError(msg)320        config.pluginmanager.register(LangSmithPlugin(), "langsmith_output_plugin")321        # Suppress warnings summary322        config.option.showwarnings = False323 324 325def _abbreviate(x: str, max_len: int) -> str:326    if len(x) > max_len:327        return x[: max_len - 3] + "..."328    else:329        return x330 331 332def _abbreviate_test_name(test_name: str, max_len: int) -> str:333    if len(test_name) > max_len:334        file, test = test_name.split("::")335        if len(".py::" + test) > max_len:336            return "..." + test[-(max_len - 3) :]337        file_len = max_len - len("...::" + test)338        return "..." + file[-file_len:] + "::" + test339    else:340        return test_name341 342 343def _merge_statuses(update: dict, current: dict, *, unpack: list[str]) -> dict:344    for path in unpack:345        if path_update := update.pop(path, None):346            path_current = current.get(path, {})347            if isinstance(path_update, dict) and isinstance(path_current, dict):348                current[path] = {**path_current, **path_update}349            else:350                current[path] = path_update351    return {**current, **update}352 353 354def _dumps_with_fallback(obj: Any) -> str:355    try:356        return json.dumps(obj)357    except Exception:358        return "unserializable"359