from __future__ import annotations import json import os import socket import statistics import subprocess import sys import time import unittest from datetime import datetime, timezone from urllib.request import Request, urlopen from tests.real_server import RealUvicornServer IMPORT_ERROR = None try: from backend import app except ModuleNotFoundError as exc: app = None IMPORT_ERROR = exc RESULTS_PATH = os.path.join(os.path.dirname(__file__), "benchmark-results.json") def _git_sha(): try: repo = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) return subprocess.run( ["git", "-c", f"safe.directory={repo}", "-C", repo, "log", "-1", "--format=%H"], capture_output=True, text=True, timeout=5, ).stdout.strip() except Exception: return "" def _save_results(project, queries): existing = {"results": []} if os.path.isfile(RESULTS_PATH): try: with open(RESULTS_PATH) as f: existing = json.load(f) except (json.JSONDecodeError, OSError): pass existing.setdefault("results", []).append({ "timestamp": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"), "project": project, "git_sha": _git_sha(), "queries": queries, }) existing["results"] = existing["results"][-50:] with open(RESULTS_PATH, "w") as f: json.dump(existing, f, indent=2, ensure_ascii=False) def _print_and_collect(results_list, label, iterations, times): median = statistics.median(times) mn = min(times) mx = max(times) print(f" {label:<46s} median={median:>8.3f}ms min={mn:>8.3f}ms max={mx:>8.3f}ms (n={iterations})") results_list.append({ "label": label.strip(), "iterations": iterations, "median_ms": round(median, 3), "min_ms": round(mn, 3), "max_ms": round(mx, 3), }) def free_port(): with socket.socket() as sock: sock.bind(("127.0.0.1", 0)) return sock.getsockname()[1] # ---- shared server for API & browser tests ---- _server_proc = None _server_port = None _server = None def start_server(): global _server_proc, _server_port, _server if _server_proc is not None: return _server = RealUvicornServer().start() _server_proc, _server_port = _server.process, _server.port def stop_server(): global _server_proc, _server_port, _server if _server: _server.stop() _server.stderr_file.close() _server = _server_proc = _server_port = None # ===== local benchmarks (pure function) ===== @unittest.skipIf(app is None, f"app dependencies unavailable: {IMPORT_ERROR}") class LocalSearchBenchmarkTests(unittest.TestCase): _results: list = [] @classmethod def setUpClass(cls): if app and not app.records: app.load_data() def _bench(self, label, fn, iterations=10): times = [] for _ in range(iterations): start = time.perf_counter() fn() elapsed = (time.perf_counter() - start) * 1000 times.append(elapsed) _print_and_collect(self._results, label, iterations, times) def test_benchmarks(self): self._bench('"手机" normal', lambda: app.search("手机")) self._bench('"手机" exact', lambda: app.search("手机", exact=True)) self._bench('"文化革命" exact', lambda: app.search("文化革命", exact=True)) self._bench('"文化 革命" normal', lambda: app.search("文化 革命")) self._bench('"手。机" exact', lambda: app.search("手。机", exact=True)) self._bench('"手*机" exact', lambda: app.search("手*机", exact=True)) self._bench('"手?机" exact', lambda: app.search("手?机", exact=True)) self._bench('"文" exact', lambda: app.search("文", exact=True)) self._bench('"ABC" normal', lambda: app.search("ABC")) self._bench('"手机" search_folders=False', lambda: app.search("手机", search_folders=False)) self._bench('sort by name', lambda: app.search(page=1, page_size=100, sort="name")) self._bench('sort by size', lambda: app.search(page=1, page_size=100, sort="size")) @classmethod def tearDownClass(cls): if cls._results: _save_results("huggingface-Search-local", cls._results) # ===== API benchmarks (HTTP to local server) ===== class ApiSearchBenchmarkTests(unittest.TestCase): _results: list = [] def _bench(self, label, path, payload=None, iterations=5): body = json.dumps(payload).encode() if payload else None req = Request(f"http://127.0.0.1:{_server_port}{path}", data=body, headers={"Content-Type": "application/json"} if body else {}) times = [] for i in range(iterations + 1): start = time.perf_counter() with urlopen(req, timeout=60) as r: r.read() elapsed = (time.perf_counter() - start) * 1000 if i > 0: times.append(elapsed) _print_and_collect(self._results, label, iterations, times) def test_search_api_benchmarks(self): self._bench('"手机" normal', "/api/search", {"q": "手机"}) self._bench('"手机" exact', "/api/search", {"q": "手机", "exact": True}) self._bench('"文化革命" exact', "/api/search", {"q": "文化革命", "exact": True}) self._bench('"文化 革命" normal', "/api/search", {"q": "文化 革命"}) self._bench('"手。机" exact', "/api/search", {"q": "手。机", "exact": True}) self._bench('"手*机" exact', "/api/search", {"q": "手*机", "exact": True}) self._bench('"手?机" exact', "/api/search", {"q": "手?机", "exact": True}) self._bench('"文" exact', "/api/search", {"q": "文", "exact": True}) self._bench('"ABC" normal', "/api/search", {"q": "ABC"}) self._bench('repositories', "/api/repos") @classmethod def tearDownClass(cls): if cls._results: _save_results("huggingface-Search-api", cls._results) # ===== browser benchmarks (Playwright) ===== try: from playwright.sync_api import sync_playwright except ImportError: sync_playwright = None @unittest.skipUnless(sync_playwright, "install requirements-test.txt, run: playwright install chromium") class BrowserSearchBenchmarkTests(unittest.TestCase): _results: list = [] @classmethod def setUpClass(cls): cls.playwright = sync_playwright().start() cls.browser = cls.playwright.chromium.launch( headless=True, args=["--no-sandbox", "--disable-gpu", "--disable-dev-shm-usage"]) @classmethod def tearDownClass(cls): if cls._results: _save_results("huggingface-Search-browser", cls._results) cls.browser.close() cls.playwright.stop() def _bench_search(self, label, query, iterations=3): ctx = self.browser.new_context(viewport={"width": 1440, "height": 900}) page = ctx.new_page() times = [] for i in range(iterations + 1): page.goto(f"http://127.0.0.1:{_server_port}/") page.wait_for_selector("#results-list .result-item", timeout=30000) inp = page.locator("#search-input") start = time.perf_counter() inp.fill(query) page.wait_for_function("""query => STATE.query === query && !STATE.isLoading && displayedSearchView?.key === getSearchViewKey() """, arg=query, timeout=30000) elapsed = (time.perf_counter() - start) * 1000 if i > 0: times.append(elapsed) ctx.close() _print_and_collect(self._results, label, iterations, times) def test_browser_search_benchmarks(self): self._bench_search('"手机" normal', "手机") self._bench_search('"文化革命" exact', "文化革命") self._bench_search('"手*机" exact', "手*机") def setUpModule(): if os.environ.get("HF_SEARCH_RUN_BENCHMARK") != "1": raise unittest.SkipTest("set HF_SEARCH_RUN_BENCHMARK=1 to run benchmarks") start_server() def tearDownModule(): stop_server()