TensorRT-LLMs/tests/unittest/llmapi/test_utils.py
Yan Chunwei 5eae3184fa
[None][chore] add missing tests to test list (#6590)
Signed-off-by: Superjomn <328693+Superjomn@users.noreply.github.com>
2025-08-06 22:12:27 +08:00

27 lines
790 B
Python

from tensorrt_llm.llmapi.utils import ApiStatusRegistry
def test_api_status_registry():
@ApiStatusRegistry.set_api_status("beta")
def _my_method(self, *args, **kwargs):
pass
assert ApiStatusRegistry.get_api_status(_my_method) == "beta"
@ApiStatusRegistry.set_api_status("prototype")
def _my_method(self, *args, **kwargs):
pass
# will always keep the first status, and the behaviour will be unknown if
# one method is registered with a different status in different files.
assert ApiStatusRegistry.get_api_status(_my_method) == "beta"
class App:
@ApiStatusRegistry.set_api_status("beta")
def _my_method(self, *args, **kwargs):
pass
assert ApiStatusRegistry.get_api_status(App._my_method) == "beta"