mirror of https://github.com/microsoft/autogen.git
42 lines
1.5 KiB
Python
42 lines
1.5 KiB
Python
from typing import Dict, List, Optional, Tuple, Union
|
|
|
|
from pydantic import BaseModel, Field
|
|
from typing_extensions import Annotated
|
|
|
|
from autogen._pydantic import model_dump, model_dump_json, type2schema
|
|
|
|
|
|
def test_type2schema() -> None:
|
|
assert type2schema(str) == {"type": "string"}
|
|
assert type2schema(int) == {"type": "integer"}
|
|
assert type2schema(float) == {"type": "number"}
|
|
assert type2schema(bool) == {"type": "boolean"}
|
|
assert type2schema(None) == {"type": "null"}
|
|
assert type2schema(Optional[int]) == {"anyOf": [{"type": "integer"}, {"type": "null"}]}
|
|
assert type2schema(List[int]) == {"items": {"type": "integer"}, "type": "array"}
|
|
assert type2schema(Tuple[int, float, str]) == {
|
|
"maxItems": 3,
|
|
"minItems": 3,
|
|
"prefixItems": [{"type": "integer"}, {"type": "number"}, {"type": "string"}],
|
|
"type": "array",
|
|
}
|
|
assert type2schema(Dict[str, int]) == {"additionalProperties": {"type": "integer"}, "type": "object"}
|
|
assert type2schema(Annotated[str, "some text"]) == {"type": "string"}
|
|
assert type2schema(Union[int, float]) == {"anyOf": [{"type": "integer"}, {"type": "number"}]}
|
|
|
|
|
|
def test_model_dump() -> None:
|
|
class A(BaseModel):
|
|
a: str
|
|
b: int = 2
|
|
|
|
assert model_dump(A(a="aaa")) == {"a": "aaa", "b": 2}
|
|
|
|
|
|
def test_model_dump_json() -> None:
|
|
class A(BaseModel):
|
|
a: str
|
|
b: int = 2
|
|
|
|
assert model_dump_json(A(a="aaa")).replace(" ", "") == '{"a":"aaa","b":2}'
|