1  import pytest
 2  
 3  from mlflow.types.agent import ChatAgentChunk, ChatAgentMessage, ChatAgentResponse
 4  
 5  
 6  def test_chat_agent_message_throws_on_invalid_data():
 7      # Missing 'content' or 'tool_calls'
 8      data = {"role": "user", "name": "test_user"}
 9      with pytest.raises(ValueError, match="Either 'content' or 'tool_calls'"):
10          ChatAgentMessage(**data)
11  
12      # Missing 'name' for tool message
13      data = {"role": "tool", "content": "test_content", "tool_call_id": "test_tool_call_id"}
14      with pytest.raises(ValueError, match="Both 'name' and 'tool_call_id'."):
15          ChatAgentMessage(**data)
16  
17      # Missing 'tool_call_id' for tool message
18      data = {"role": "tool", "name": "test_user", "content": "test_content"}
19      with pytest.raises(ValueError, match="Both 'name' and 'tool_call_id'."):
20          ChatAgentMessage(**data)
21  
22      # Missing both 'name' and 'tool_call_id' for tool message
23      data = {"role": "tool", "content": "test_content"}
24      with pytest.raises(ValueError, match="Both 'name' and 'tool_call_id'."):
25          ChatAgentMessage(**data)
26  
27  
28  def test_chat_agent_message_allows_empty_content():
29      # Empty string content should be allowed (not None)
30      data = {"role": "assistant", "content": ""}
31      message = ChatAgentMessage(**data)
32      assert message.content == ""
33      assert message.tool_calls is None
34  
35  
36  def test_chat_agent_response_throws_on_missing_id():
37      data = {"messages": [{"role": "user", "content": "a"}]}
38      with pytest.raises(ValueError, match="All ChatAgentMessage objects in field `messages`"):
39          ChatAgentResponse(**data)
40  
41  
42  def test_chat_agent_response_throws_on_duplicate_ids():
43      data = {
44          "messages": [
45              {"role": "user", "content": "a", "id": "1"},
46              {"role": "user", "content": "b", "id": "1"},
47          ]
48      }
49      with pytest.raises(ValueError, match="unique IDs"):
50          ChatAgentResponse(**data)
51  
52  
53  def test_chat_agent_response_throws_on_updated_id():
54      data = {
55          "messages": [
56              {"role": "user", "content": "a", "id": "1"},
57          ]
58      }
59      # shouldn't raise an error
60      response = ChatAgentResponse(**data)
61      with pytest.raises(ValueError, match="unique IDs"):
62          response.messages = [
63              ChatAgentMessage(**{"role": "user", "content": "b", "id": "1"}),
64              ChatAgentMessage(**{"role": "user", "content": "a", "id": "1"}),
65          ]
66  
67  
68  def test_chat_agent_chunk_throws_on_missing_id():
69      data = {"delta": {"role": "user", "content": "a"}}
70      with pytest.raises(ValueError, match="The field `delta` of ChatAgentChunk"):
71          ChatAgentChunk(**data)
72  
73  
74  def test_chat_agent_chunk_throws_on_updated_id():
75      data = {"delta": {"role": "user", "content": "a", "id": "1"}}
76      chunk = ChatAgentChunk(**data)
77      with pytest.raises(ValueError, match="The field `delta` of ChatAgentChunk"):
78          chunk.delta = ChatAgentMessage(**{"role": "user", "content": "b"})