TensorRT-LLMs/tests/unittest/disaggregated/test_agent_multi_backends.py
Shi Xiaowei 944c304bbb
[TRTLLM-9527][feat] Python transceiver components (step 2) (#10494)
Signed-off-by: Shixiaowei02 <39303645+Shixiaowei02@users.noreply.github.com>
2026-01-22 10:14:50 -08:00

33 lines
848 B
Python

import os
import subprocess
import pytest
@pytest.mark.parametrize("use_py_nixl", ["0", "1"])
def test_run_with_different_env(use_py_nixl):
os.environ["TRTLLM_USE_PY_NIXL_KVCACHE"] = use_py_nixl
print(f"Running tests with TRTLLM_USE_PY_NIXL_KVCACHE={use_py_nixl}")
test_file_path = os.path.join(os.path.dirname(__file__), "test_agent.py")
print(f"Running tests in: {test_file_path}")
result = subprocess.run(
["pytest", "--capture=no", test_file_path],
env=os.environ.copy(),
capture_output=True,
text=True,
)
print(result.stdout)
if result.returncode != 0:
print("Test failed. stderr output:")
print(result.stderr)
assert result.returncode == 0, f"Tests failed with TRTLLM_USE_PY_NIXL_KVCACHE={use_py_nixl}"
if __name__ == "__main__":
pytest.main()