Spaces:
Running
Running
File size: 3,930 Bytes
5f923cd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 | # Copyright 2026 Google LLC.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Pytest fixtures and hooks for litert_lm engine tests.
This file contains pytest configurations, fixtures, and hooks used by the
tests in this directory. It handles command-line argument parsing,
dynamic test parametrization, and provides helper functions for running
the litert_lm C++ engine.
"""
import os
import subprocess
from typing import Callable
import pytest
def pytest_addoption(parser: pytest.Parser) -> None:
"""Adds custom command line arguments to pytest."""
parser.addoption(
"--model-path",
action="store",
required=True,
help="Absolute or relative path to the .litertlm model file",
)
parser.addoption(
"--build-system",
action="store",
choices=["bazel", "cmake"],
required=True,
help="The build system to use for the engine binary (bazel or cmake).",
)
@pytest.fixture(scope="session")
def model_path(request: pytest.FixtureRequest) -> str:
"""Extracts the model path from the CLI and ensures the file exists."""
path = request.config.getoption("--model-path")
if not os.path.exists(path):
pytest.fail(f"CRITICAL: Model file not found at {path}")
return path
def pytest_generate_tests(metafunc: pytest.Metafunc) -> None:
"""Generates test parameters for the 'engine_binary' fixture.
This function is a pytest hook that dynamically parametrizes tests
that depend on the 'engine_binary' fixture, based on the build system
specified via command-line options.
Args:
metafunc: The metafunc object provided by pytest.
"""
if "engine_binary" not in metafunc.fixturenames:
return
build_system = metafunc.config.getoption("--build-system")
repo_root = os.path.abspath(
os.path.join(os.path.dirname(__file__), "..", "..")
)
exe_ext = ".exe" if os.name == "nt" else ""
if build_system == "bazel":
binary_path = os.path.join(
repo_root,
"bazel-bin",
"runtime",
"engine",
f"litert_lm_main{exe_ext}",
)
else:
binary_path = os.path.join(
repo_root, "cmake", "build", f"litert_lm_main{exe_ext}"
)
if not os.path.exists(binary_path):
pytest.fail(
f"CRITICAL: Engine binary not found at {binary_path}. "
f"Did you run the {build_system} build?"
)
metafunc.parametrize("engine_binary", [binary_path])
@pytest.fixture
def run_engine(
engine_binary: str, model_path: str
) -> Callable[..., str]:
"""Provides a helper fixture that tests can call to execute the C++ engine.
It automatically handles timeouts and fatal crashes.
Args:
engine_binary: Path to the engine executable.
model_path: Path to the model file.
"""
def _run(prompt: str, timeout: int = 120) -> str:
cmd = [
engine_binary,
"--backend=cpu",
f"--model_path={model_path}",
f"--input_prompt={prompt}",
]
result = subprocess.run(
cmd, capture_output=True, text=True, timeout=timeout, check=False
)
# Instantly fail the test if the C++ engine segfaults or OOMs
if result.returncode != 0:
error_tail = (result.stderr if result.stderr else result.stdout).strip()[
-800:
]
pytest.fail(
f"Engine crashed with code {result.returncode}!\nCrash"
f" Tail:\n{error_tail}"
)
return result.stdout + result.stderr
return _run
|