wehub-resource-sync
138 lines
4.3 KiB
Python
138 lines
4.3 KiB
Python
import sys
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
from claude_agent_sdk.types import AssistantMessage, ResultMessage, TextBlock, UserMessage
|
|
|
|
import mlflow.anthropic
|
|
from mlflow.anthropic.autolog import patched_claude_sdk_init
|
|
|
|
|
|
def test_anthropic_autolog_without_claude_sdk():
|
|
sys.modules.pop("claude_agent_sdk", None)
|
|
|
|
with (
|
|
patch.dict(
|
|
"sys.modules",
|
|
{
|
|
"anthropic": MagicMock(__version__="0.35.0"),
|
|
"anthropic.resources": MagicMock(Messages=MagicMock, AsyncMessages=MagicMock),
|
|
},
|
|
),
|
|
patch("mlflow.anthropic.safe_patch"),
|
|
):
|
|
mlflow.anthropic.autolog()
|
|
|
|
|
|
def _patch_sdk_init(mock_self, response_messages):
|
|
original_init = MagicMock()
|
|
|
|
async def fake_receive_response():
|
|
for msg in response_messages:
|
|
yield msg
|
|
|
|
mock_self.receive_response = fake_receive_response
|
|
patched_claude_sdk_init(original_init, mock_self)
|
|
return original_init
|
|
|
|
|
|
def test_patched_claude_sdk_init_wraps_receive_response():
|
|
mock_self = MagicMock()
|
|
|
|
async def fake_receive_response():
|
|
yield "msg1"
|
|
|
|
mock_self.receive_response = fake_receive_response
|
|
original_init = MagicMock()
|
|
patched_claude_sdk_init(original_init, mock_self)
|
|
|
|
original_init.assert_called_once_with(mock_self, None)
|
|
assert mock_self.receive_response is not fake_receive_response
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_receive_response_builds_trace():
|
|
mock_self = MagicMock()
|
|
messages = [
|
|
UserMessage(content="Hello"),
|
|
AssistantMessage(content=[TextBlock(text="Hi!")], model="claude-sonnet-4-20250514"),
|
|
ResultMessage(
|
|
subtype="success",
|
|
duration_ms=5000,
|
|
duration_api_ms=4000,
|
|
is_error=False,
|
|
num_turns=1,
|
|
session_id="test-session",
|
|
usage={"input_tokens": 100, "output_tokens": 20},
|
|
),
|
|
]
|
|
_patch_sdk_init(mock_self, messages)
|
|
|
|
with (
|
|
patch("mlflow.utils.autologging_utils.autologging_is_disabled", return_value=False),
|
|
patch("mlflow.claude_code.tracing.process_sdk_messages") as mock_process,
|
|
):
|
|
[msg async for msg in mock_self.receive_response()]
|
|
|
|
mock_process.assert_called_once()
|
|
called_messages = mock_process.call_args[0][0]
|
|
assert len(called_messages) == 3
|
|
result_messages = [m for m in called_messages if isinstance(m, ResultMessage)]
|
|
assert len(result_messages) == 1
|
|
assert result_messages[0].usage == {"input_tokens": 100, "output_tokens": 20}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_query_captures_async_generator_prompt():
|
|
mock_self = MagicMock()
|
|
|
|
async def fake_query(prompt, *args, **kwargs):
|
|
# Consume the generator like the real SDK would
|
|
async for _ in prompt:
|
|
pass
|
|
|
|
mock_self.query = fake_query
|
|
|
|
response_messages = [
|
|
AssistantMessage(content=[TextBlock(text="Hi!")], model="claude-sonnet-4-20250514"),
|
|
ResultMessage(
|
|
subtype="success",
|
|
duration_ms=1000,
|
|
duration_api_ms=800,
|
|
is_error=False,
|
|
num_turns=1,
|
|
session_id="s",
|
|
),
|
|
]
|
|
_patch_sdk_init(mock_self, response_messages)
|
|
|
|
async def prompt_generator():
|
|
yield {"type": "user", "message": {"role": "user", "content": "Hello from generator"}}
|
|
|
|
with (
|
|
patch("mlflow.utils.autologging_utils.autologging_is_disabled", return_value=False),
|
|
patch("mlflow.claude_code.tracing.process_sdk_messages") as mock_process,
|
|
):
|
|
await mock_self.query(prompt_generator())
|
|
[msg async for msg in mock_self.receive_response()]
|
|
|
|
mock_process.assert_called_once()
|
|
called_messages = mock_process.call_args[0][0]
|
|
user_messages = [m for m in called_messages if isinstance(m, UserMessage)]
|
|
assert len(user_messages) == 1
|
|
assert user_messages[0].content == "Hello from generator"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_receive_response_skips_when_autologging_disabled():
|
|
mock_self = MagicMock()
|
|
_patch_sdk_init(mock_self, ["msg1", "msg2"])
|
|
|
|
with (
|
|
patch("mlflow.utils.autologging_utils.autologging_is_disabled", return_value=True),
|
|
patch("mlflow.claude_code.tracing.process_sdk_messages") as mock_process,
|
|
):
|
|
[msg async for msg in mock_self.receive_response()]
|
|
|
|
mock_process.assert_not_called()
|