33 lines
1.4 KiB
Python
33 lines
1.4 KiB
Python
import logging
|
|
from uuid import uuid4
|
|
|
|
from deerflow.models.timing_callback import ModelTimingCallback
|
|
|
|
|
|
def test_model_timing_callback_logs_connection_and_completion(caplog):
|
|
callback = ModelTimingCallback(model_name="qwen", base_url="http://model/v1")
|
|
run_id = uuid4()
|
|
|
|
with caplog.at_level(logging.INFO, logger="deerflow.models.timing_callback"):
|
|
callback.on_chat_model_start({}, [[object()]], run_id=run_id)
|
|
callback.on_llm_new_token("", run_id=run_id, chunk=object())
|
|
callback.on_llm_end(object(), run_id=run_id)
|
|
|
|
messages = [record.getMessage() for record in caplog.records]
|
|
assert any("LLM_CONNECT_START" in message for message in messages)
|
|
assert any("LLM_CONNECT_OK" in message and "first_chunk_ms=" in message for message in messages)
|
|
assert any("LLM_REQUEST_DONE" in message and "total_ms=" in message for message in messages)
|
|
|
|
|
|
def test_model_timing_callback_logs_failure(caplog):
|
|
callback = ModelTimingCallback(model_name="qwen", base_url="http://model/v1")
|
|
run_id = uuid4()
|
|
|
|
with caplog.at_level(logging.INFO, logger="deerflow.models.timing_callback"):
|
|
callback.on_chat_model_start({}, [[object()]], run_id=run_id)
|
|
callback.on_llm_error(TimeoutError("upstream timeout"), run_id=run_id)
|
|
|
|
message = caplog.records[-1].getMessage()
|
|
assert "LLM_REQUEST_FAILED" in message
|
|
assert "TimeoutError" in message
|