Files
2026-07-13 12:24:33 +08:00

369 lines
13 KiB
Python

# SPDX-License-Identifier: Apache-2.0
"""Tests for InteractiveState."""
# Standard
import argparse
import json
# First Party
from lmcache.cli.commands.bench.engine_bench.interactive.state import (
InteractiveState,
)
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _make_full_state(**overrides: object) -> InteractiveState:
"""Build a state with all required fields set."""
values: dict[str, object] = {
"engine_url": "http://localhost:8000",
"workload": "long-doc-qa",
"tokens_per_gb_kvcache": 6553,
}
values.update(overrides)
state = InteractiveState()
for k, v in values.items():
state.set(k, v)
return state
def _make_cli_args(**overrides: object) -> argparse.Namespace:
"""Build a Namespace mimicking bench engine CLI output."""
defaults: dict[str, object] = {
"engine_url": None,
"workload": None,
"model": None,
"lmcache_url": None,
"tokens_per_gb_kvcache": None,
"kv_cache_volume": 100.0,
"seed": 42,
"output_dir": ".",
"no_csv": False,
"json": False,
"quiet": False,
"ldqa_document_length": 10000,
"ldqa_query_per_document": 2,
"ldqa_shuffle_policy": "random",
"ldqa_num_inflight_requests": 3,
"mrc_shared_prompt_length": 2000,
"mrc_chat_history_length": 10000,
"mrc_user_input_length": 50,
"mrc_output_length": 200,
"mrc_qps": 1.0,
"mrc_duration": 60.0,
"rp_request_length": 10000,
"rp_num_requests": 50,
"format": None,
"output": None,
"config": None,
"bench_target": "engine",
}
defaults.update(overrides)
return argparse.Namespace(**defaults)
# ---------------------------------------------------------------------------
# Basic accessors
# ---------------------------------------------------------------------------
class TestBasicAccessors:
def test_empty_state(self) -> None:
state = InteractiveState()
assert not state.is_set("engine_url")
assert state.get("engine_url") is None
assert state.get("engine_url", "fallback") == "fallback"
def test_set_and_get(self) -> None:
state = InteractiveState()
state.set("engine_url", "http://localhost:8000")
assert state.is_set("engine_url")
assert state.get("engine_url") == "http://localhost:8000"
def test_values_copy(self) -> None:
state = _make_full_state()
vals = state.values
vals["engine_url"] = "mutated"
assert state.get("engine_url") == "http://localhost:8000"
# ---------------------------------------------------------------------------
# Readiness
# ---------------------------------------------------------------------------
class TestReadiness:
def test_empty_not_ready(self) -> None:
state = InteractiveState()
assert not state.is_ready()
def test_full_state_ready(self) -> None:
state = _make_full_state()
assert state.is_ready()
def test_missing_engine_url(self) -> None:
state = _make_full_state()
state._values.pop("engine_url")
assert not state.is_ready()
def test_missing_workload(self) -> None:
state = _make_full_state()
state._values.pop("workload")
assert not state.is_ready()
def test_tokens_per_gb_not_required_when_lmcache_set(self) -> None:
state = InteractiveState()
state.set("engine_url", "http://localhost:8000")
state.set("workload", "long-doc-qa")
state.set("lmcache_url", "http://localhost:8080")
# tokens_per_gb_kvcache not set, but lmcache_url is → condition
# for tokens_per_gb is not met → not required
assert state.is_ready()
# ---------------------------------------------------------------------------
# Missing required
# ---------------------------------------------------------------------------
class TestMissingRequired:
def test_all_missing(self) -> None:
state = InteractiveState()
missing = state.get_missing_required()
keys = [item.key for item in missing]
assert "engine_url" in keys
assert "workload" in keys
def test_none_missing(self) -> None:
state = _make_full_state()
assert state.get_missing_required() == []
def test_tokens_per_gb_skipped_with_lmcache(self) -> None:
state = InteractiveState()
state.set("engine_url", "http://localhost:8000")
state.set("workload", "long-doc-qa")
state.set("lmcache_url", "http://localhost:8080")
missing = state.get_missing_required()
keys = [item.key for item in missing]
assert "tokens_per_gb_kvcache" not in keys
# ---------------------------------------------------------------------------
# Workload items
# ---------------------------------------------------------------------------
class TestWorkloadItems:
def test_long_doc_qa(self) -> None:
state = _make_full_state(workload="long-doc-qa")
items = state.get_workload_items()
keys = [item.key for item in items]
assert "ldqa_document_length" in keys
assert "mrc_qps" not in keys
assert "rp_num_requests" not in keys
def test_multi_round_chat(self) -> None:
state = _make_full_state(workload="multi-round-chat")
items = state.get_workload_items()
keys = [item.key for item in items]
assert "mrc_shared_prompt_length" in keys
assert "ldqa_document_length" not in keys
def test_random_prefill(self) -> None:
state = _make_full_state(workload="random-prefill")
items = state.get_workload_items()
keys = [item.key for item in items]
assert "rp_request_length" in keys
assert "ldqa_document_length" not in keys
def test_workload_items_all_default_initially(self) -> None:
state = _make_full_state(workload="long-doc-qa")
assert state.workload_items_all_default()
def test_workload_items_not_default_after_set(self) -> None:
state = _make_full_state(workload="long-doc-qa")
state.set("ldqa_document_length", 5000)
assert not state.workload_items_all_default()
# ---------------------------------------------------------------------------
# Fill defaults
# ---------------------------------------------------------------------------
class TestFillDefaults:
def test_fills_kv_cache_volume(self) -> None:
state = _make_full_state(workload="long-doc-qa")
state.fill_defaults()
assert state.get("kv_cache_volume") == 100.0
def test_fills_workload_defaults(self) -> None:
state = _make_full_state(workload="long-doc-qa")
state.fill_defaults()
assert state.get("ldqa_document_length") == 10000
assert state.get("ldqa_query_per_document") == 2
def test_does_not_overwrite_set_values(self) -> None:
state = _make_full_state(workload="long-doc-qa")
state.set("ldqa_document_length", 5000)
state.fill_defaults()
assert state.get("ldqa_document_length") == 5000
# ---------------------------------------------------------------------------
# from_cli_args
# ---------------------------------------------------------------------------
class TestFromCliArgs:
def test_picks_up_explicit_values(self) -> None:
args = _make_cli_args(
engine_url="http://host:8000",
workload="long-doc-qa",
tokens_per_gb_kvcache=6553,
)
state = InteractiveState.from_cli_args(args)
assert state.get("engine_url") == "http://host:8000"
assert state.get("workload") == "long-doc-qa"
assert state.get("tokens_per_gb_kvcache") == 6553
def test_none_values_not_set(self) -> None:
args = _make_cli_args() # all defaults
state = InteractiveState.from_cli_args(args)
assert not state.is_set("engine_url")
assert not state.is_set("workload")
assert not state.is_set("model")
def test_default_values_not_set(self) -> None:
# kv_cache_volume defaults to 100.0 — should not be marked as set
args = _make_cli_args(kv_cache_volume=100.0)
state = InteractiveState.from_cli_args(args)
assert not state.is_set("kv_cache_volume")
def test_non_default_values_set(self) -> None:
args = _make_cli_args(kv_cache_volume=50.0)
state = InteractiveState.from_cli_args(args)
assert state.is_set("kv_cache_volume")
assert state.get("kv_cache_volume") == 50.0
# ---------------------------------------------------------------------------
# to_namespace
# ---------------------------------------------------------------------------
class TestToNamespace:
def test_round_trip(self) -> None:
state = _make_full_state(workload="long-doc-qa")
ns = state.to_namespace()
assert ns.engine_url == "http://localhost:8000"
assert ns.workload == "long-doc-qa"
assert ns.tokens_per_gb_kvcache == 6553
# Defaults filled
assert ns.kv_cache_volume == 100.0
assert ns.ldqa_document_length == 10000
assert ns.seed == 42
def test_has_bench_target(self) -> None:
state = _make_full_state()
ns = state.to_namespace()
assert ns.bench_target == "engine"
def test_has_output_attrs(self) -> None:
state = _make_full_state()
ns = state.to_namespace()
assert hasattr(ns, "output_dir")
assert hasattr(ns, "no_csv")
assert hasattr(ns, "quiet")
assert hasattr(ns, "format")
assert hasattr(ns, "output")
# ---------------------------------------------------------------------------
# JSON round-trip
# ---------------------------------------------------------------------------
class TestJsonRoundTrip:
def test_to_json_and_back(self) -> None:
state = _make_full_state(workload="long-doc-qa")
state.set("ldqa_document_length", 5000)
data = state.to_json()
# engine_url excluded from export (environment-specific)
assert "engine_url" not in data
assert data["ldqa_document_length"] == 5000
restored = InteractiveState.from_json(data)
assert restored.get("ldqa_document_length") == 5000
def test_save_and_load(self, tmp_path: object) -> None:
# Standard
import pathlib
path = str(pathlib.Path(str(tmp_path)) / "test_config.json")
state = _make_full_state(workload="long-doc-qa")
state.save_json(path)
loaded = InteractiveState.load_json(path)
assert "engine_url" not in loaded.values
assert loaded.get("workload") == "long-doc-qa"
def test_json_file_is_valid_json(self, tmp_path: object) -> None:
# Standard
import pathlib
path = str(pathlib.Path(str(tmp_path)) / "test.json")
state = _make_full_state()
state.save_json(path)
with open(path) as f:
data = json.load(f)
assert isinstance(data, dict)
assert "engine_url" not in data
assert data["workload"] == "long-doc-qa"
def test_merge_cli_args(self) -> None:
state = _make_full_state(workload="long-doc-qa")
args = _make_cli_args(
engine_url="http://override:9000",
)
state.merge_cli_args(args)
assert state.get("engine_url") == "http://override:9000"
# Original values preserved
assert state.get("workload") == "long-doc-qa"
# ---------------------------------------------------------------------------
# Summary
# ---------------------------------------------------------------------------
class TestSummary:
def test_summary_lines_non_empty(self) -> None:
state = _make_full_state(workload="long-doc-qa")
state.fill_defaults()
lines = state.summary_lines()
assert len(lines) > 0
def test_summary_includes_engine_url(self) -> None:
state = _make_full_state(workload="long-doc-qa")
state.fill_defaults()
labels = [label for label, _ in state.summary_lines()]
assert "Engine URL" in labels
def test_summary_shows_auto_detect_for_empty_model(self) -> None:
state = _make_full_state(workload="long-doc-qa")
state.fill_defaults()
for label, value in state.summary_lines():
if label == "Model name":
assert value == "(auto-detect)"
break
def test_summary_excludes_wrong_workload_items(self) -> None:
state = _make_full_state(workload="long-doc-qa")
state.fill_defaults()
labels = [label for label, _ in state.summary_lines()]
assert "Queries per second" not in labels # mrc item
assert "Number of requests" not in labels # rp item