166 lines
4.6 KiB
Python
166 lines
4.6 KiB
Python
"""Test JSON serialization for BaseModel objects."""
|
|
|
|
import json
|
|
import pytest
|
|
from typing import Optional, List
|
|
|
|
from zai.core._base_models import BaseModel
|
|
from zai.core._json_encoder import ZAIJSONEncoder, json_dumps
|
|
from zai.core._http_client import HttpClient
|
|
|
|
|
|
class ModelForTest(BaseModel):
|
|
"""Test model for JSON serialization."""
|
|
name: str
|
|
age: Optional[int] = None
|
|
tags: Optional[List[str]] = None
|
|
|
|
|
|
class MessageForTest(BaseModel):
|
|
"""Test message model similar to CompletionMessage."""
|
|
content: Optional[str] = None
|
|
role: str
|
|
tool_calls: Optional[List[dict]] = None
|
|
|
|
|
|
def test_base_model_dict_conversion():
|
|
"""Test that BaseModel can be converted to dict."""
|
|
model = ModelForTest(name="test", age=25, tags=["tag1", "tag2"])
|
|
|
|
# Test model_dump() method
|
|
model_dict = model.model_dump()
|
|
assert model_dict == {"name": "test", "age": 25, "tags": ["tag1", "tag2"]}
|
|
|
|
# Test keys/values/items methods
|
|
assert list(model.keys()) == ["name", "age", "tags"]
|
|
assert list(model.values()) == ["test", 25, ["tag1", "tag2"]]
|
|
assert list(model.items()) == [("name", "test"), ("age", 25), ("tags", ["tag1", "tag2"])]
|
|
|
|
|
|
def test_zai_json_encoder():
|
|
"""Test custom JSON encoder for BaseModel objects."""
|
|
model = ModelForTest(name="test", age=25)
|
|
|
|
# Test with custom encoder
|
|
json_str = json.dumps(model, cls=ZAIJSONEncoder)
|
|
result = json.loads(json_str)
|
|
assert result == {"name": "test", "age": 25}
|
|
|
|
# Test with nested BaseModel
|
|
message = MessageForTest(content="hello", role="user")
|
|
nested_data = {"message": message, "other": "data"}
|
|
|
|
json_str = json.dumps(nested_data, cls=ZAIJSONEncoder)
|
|
result = json.loads(json_str)
|
|
assert result == {
|
|
"message": {"content": "hello", "role": "user"},
|
|
"other": "data"
|
|
}
|
|
|
|
|
|
def test_json_dumps_utility():
|
|
"""Test the json_dumps utility function."""
|
|
model = ModelForTest(name="test", age=25)
|
|
|
|
json_str = json_dumps(model)
|
|
result = json.loads(json_str)
|
|
assert result == {"name": "test", "age": 25}
|
|
|
|
|
|
def test_http_client_json_preparation():
|
|
"""Test that HttpClient can prepare JSON data with BaseModel objects."""
|
|
# We'll create a minimal HttpClient to test the _prepare_json_data method
|
|
import httpx
|
|
from zai.core._http_client import HttpClient
|
|
|
|
# Create a minimal client for testing
|
|
client = HttpClient(
|
|
base_url="https://test.com",
|
|
version="v1",
|
|
_strict_response_validation=True,
|
|
timeout=10.0
|
|
)
|
|
|
|
# Test with BaseModel
|
|
model = ModelForTest(name="test", age=25)
|
|
result = client._prepare_json_data(model)
|
|
assert result == {"name": "test", "age": 25}
|
|
|
|
# Test with list containing BaseModel
|
|
models = [ModelForTest(name="test1", age=25), ModelForTest(name="test2", age=30)]
|
|
result = client._prepare_json_data(models)
|
|
expected = [
|
|
{"name": "test1", "age": 25},
|
|
{"name": "test2", "age": 30}
|
|
]
|
|
assert result == expected
|
|
|
|
# Test with dict containing BaseModel
|
|
data = {
|
|
"message": MessageForTest(content="hello", role="user"),
|
|
"count": 1
|
|
}
|
|
result = client._prepare_json_data(data)
|
|
expected = {
|
|
"message": {"content": "hello", "role": "user"},
|
|
"count": 1
|
|
}
|
|
assert result == expected
|
|
|
|
# Test with None
|
|
assert client._prepare_json_data(None) is None
|
|
|
|
# Test with regular data
|
|
regular_data = {"key": "value", "number": 42}
|
|
assert client._prepare_json_data(regular_data) == regular_data
|
|
|
|
|
|
def test_completion_message_like_serialization():
|
|
"""Test serialization for CompletionMessage-like objects."""
|
|
# Simulate a CompletionMessage-like object
|
|
message = MessageForTest(
|
|
content="Hello, how can I help you?",
|
|
role="assistant",
|
|
tool_calls=[{"id": "call_123", "function": {"name": "test_func"}}]
|
|
)
|
|
|
|
# Test that it can be serialized
|
|
json_str = json_dumps(message)
|
|
result = json.loads(json_str)
|
|
|
|
expected = {
|
|
"content": "Hello, how can I help you?",
|
|
"role": "assistant",
|
|
"tool_calls": [{"id": "call_123", "function": {"name": "test_func"}}]
|
|
}
|
|
assert result == expected
|
|
|
|
# Test in a messages list (similar to the original error case)
|
|
messages = [
|
|
{"role": "user", "content": "Help me"},
|
|
message # This is a BaseModel object
|
|
]
|
|
|
|
# This should work now with our custom encoder
|
|
json_str = json_dumps(messages)
|
|
result = json.loads(json_str)
|
|
|
|
expected = [
|
|
{"role": "user", "content": "Help me"},
|
|
{
|
|
"content": "Hello, how can I help you?",
|
|
"role": "assistant",
|
|
"tool_calls": [{"id": "call_123", "function": {"name": "test_func"}}]
|
|
}
|
|
]
|
|
assert result == expected
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# Run basic tests
|
|
test_base_model_dict_conversion()
|
|
test_zai_json_encoder()
|
|
test_json_dumps_utility()
|
|
test_http_client_json_preparation()
|
|
test_completion_message_like_serialization()
|
|
print("All tests passed!")
|