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